feat(megatron): Add support for Gated Delta Net (GDN) & Kimi Delta Attention (KDA)#676
Open
clairesonglee wants to merge 36 commits into
Open
feat(megatron): Add support for Gated Delta Net (GDN) & Kimi Delta Attention (KDA)#676clairesonglee wants to merge 36 commits into
clairesonglee wants to merge 36 commits into
Conversation
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
- 300M pure GDN pretrain config aligned with FLA (micro_batch=128, global=1024) - 300M model architecture config (hidden=1024, layers=24, ffn=4096) - use_short_conv in language_model.yaml - Updated KDA pretrain config Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # examples/megatron/configs/MI300X/mamba_370M-pretrain.yaml # examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml # examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml # examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml # primus/backends/megatron/core/models/hybrid/__init__.py # primus/backends/megatron/core/models/hybrid/hybrid_block.py # primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py # primus/configs/models/megatron/language_model.yaml # primus/configs/models/megatron/mamba_370M.yaml # primus/configs/models/megatron/mamba_base.yaml # primus/configs/models/megatron/zebra_llama_1B.yaml # primus/configs/models/megatron/zebra_llama_3B.yaml # primus/configs/models/megatron/zebra_llama_8B.yaml # primus/core/utils/import_utils.py # primus/modules/trainer/lightmegatron/pre_trainer.py # primus/modules/trainer/megatron/pre_trainer.py # primus/modules/trainer/megatron/trainer.py
- Introduced `bash-docker.sh` for running the Primus environment with necessary device and network configurations. - Added `GDN_FLA_PARITY.md` to document changes for achieving GDN training parity with Flash Linear Attention (FLA). - Created `README_GDN.md` for a comprehensive guide on running the 300M pure GDN pretraining recipe in Primus, including setup and performance metrics. - Updated `README.md` to include details about Zebra-Llama models and their configurations. - Enhanced Megatron patch script to apply necessary modifications for GDN training. - Adjusted training configurations to ensure alignment with FLA performance metrics.
- Introduced `KDA_FLA_PARITY.md` to document changes for achieving Kimi Delta Attention (KDA) training parity with Flash Linear Attention (FLA). - Created `README_KDA.md` for a comprehensive guide on running the 300M pure KDA pretraining recipe in Primus, including setup and performance metrics. - Added new training configuration file `zebra_llama_300M_kda_pure-pretrain.yaml` aligned with FLA specifications. - Enhanced Megatron patch script to apply necessary modifications for KDA training. - Updated model code and runtime configurations to ensure alignment with FLA performance metrics, including fused operations and normalization techniques. - Implemented diagnostic features for capturing activations and iter-1 batch tokens for comparison with FLA outputs.
- Introduced `KDA_FLA_PARITY.md` to document changes for achieving Kimi Delta Attention (KDA) training parity with Flash Linear Attention (FLA). - Created `README_KDA.md` for a comprehensive guide on running the 300M pure KDA pretraining recipe in Primus, including setup and performance metrics. - Added new training configuration file `zebra_llama_300M_kda_pure-pretrain.yaml` aligned with FLA specifications. - Enhanced Megatron patch script to apply necessary modifications for KDA training. - Updated model code and runtime configurations to ensure alignment with FLA performance metrics, including optimizations for fused operations and memory management.
clairesonglee
force-pushed
the
clairlee/kda-optimized-training-patch
branch
from
June 3, 2026 12:04
774acb4 to
c301288
Compare
clairesonglee
force-pushed
the
clairlee/kda-optimized-training-patch
branch
from
June 8, 2026 08:51
ddc56fb to
a6cc909
Compare
clairesonglee
force-pushed
the
clairlee/kda-optimized-training-patch
branch
from
June 8, 2026 14:03
a6cc909 to
c791b85
Compare
- Expanded .gitignore to include additional scratch run artifacts and IDE/editor workspace files. - Removed the all_changes.md file as it is no longer needed for tracking changes across the codebase.
- Updated `GDN_FLA_PARITY.md` and `KDA_FLA_PARITY.md` to clarify the recommended toggle profiles, emphasizing the use of YAML knobs alongside environment variables for configuration. - Revised `README_GDN.md` and `README_KDA.md` to reflect the new toggle profiles and their effects on achieving bit-level parity with FLA. - Adjusted the `megatron_patch.sh` script to streamline the patching process and removed outdated diagnostic options. - Improved clarity in the documentation regarding the configuration of FLA-order dataset shims and their impact on training performance.
The upstream Megatron-LM MambaStackSubmodules dataclass does not have a moe_layer field. Only HybridStackSubmodules (Primus) supports it. This caused a RuntimeError at import time when the spec module was loaded for any hybrid model training. Co-authored-by: Cursor <cursoragent@cursor.com>
…paths Set mock_data: true, train_data_path: null, and load: null across all zebra_llama pretrain configs so they can run without depending on /home/vanbhati@amd.com data or checkpoint paths. Co-authored-by: Cursor <cursoragent@cursor.com>
Hybrid models have parameters that don't all participate in backward, causing Megatron's DDP buffer per_param_grad_ready_counts assertion to fail when overlap_grad_reduce is enabled. Set overlap_grad_reduce, overlap_param_gather, and gradient_accumulation_fusion to false. Co-authored-by: Cursor <cursoragent@cursor.com>
… overlap flags - Set train_iters to 50 across all 12 zebra_llama configs for quick perf benchmarking - 300M configs: reduce batch sizes (mbs=8, gbs=64) for 8-GPU runs - 1B GDN Pure / KDA Pure: restore original mbs=16, gbs=128 - 1B GDN Pure 100B: restore original mbs=64, gbs=512 - 300M configs: disable overlap_grad_reduce, overlap_param_gather, gradient_accumulation_fusion Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
clairesonglee
force-pushed
the
clairlee/kda-optimized-training-patch
branch
from
July 15, 2026 05:40
c791b85 to
8424d25
Compare
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| sys.exit(main()) |
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| sys.exit(main()) |
| with open(marker, "w") as fh: | ||
| fh.write(_msg + "\n") | ||
| fh.write(f"pid={os.getpid()} ts={time.time()}\n") | ||
| except Exception: |
| eps=self.config.layernorm_epsilon, | ||
| ) | ||
| self._use_fla_fused_gated_norm = True | ||
| except ImportError: |
| fh.write("sys.path[:6]:\n") | ||
| for p in sys.path[:6]: | ||
| fh.write(f" {p}\n") | ||
| except Exception: |
| # `return_dict` for any code paths that consult config defaults. | ||
| try: | ||
| hf_model.config.return_dict = True | ||
| except Exception: |
| for h in mg_hooks: | ||
| try: | ||
| h.remove() | ||
| except Exception: |
…d-training-patch Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # .gitignore # primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py # primus/backends/megatron/core/models/hybrid/hybrid_block.py # primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py # primus/modules/trainer/megatron/pre_trainer.py # primus/modules/trainer/megatron/trainer.py
FLA's chunk_gated_delta_rule kernel produces sporadic NaN when use_gate_in_kernel=True on ROCm/MI300X. Pre-compute the gate externally in fp32 and pass with use_gate_in_kernel=False instead. Also fix module_utils import paths after main refactor (primus.modules.module_utils -> primus.core.utils.module_utils). Co-authored-by: Cursor <cursoragent@cursor.com>
clairesonglee
marked this pull request as ready for review
July 20, 2026 21:33
clairesonglee
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
July 20, 2026 21:33
Comment on lines
+121
to
+122
| # try: | ||
| # user = input("User> ").strip() |
Comment on lines
+127
to
+134
| # if not user: | ||
| # continue | ||
| # if user == "/exit": | ||
| # return | ||
| # if user == "/reset": | ||
| # history.clear() | ||
| # print("[Info] History cleared.\n") | ||
| # continue |
Comment on lines
+120
to
+123
| # if 'layer_norm_weight' in new_key: | ||
| # new_key = new_key.replace('layer_norm_weight', 'weight') | ||
| # if 'layer_norm_bias' in new_key: | ||
| # new_key = new_key.replace('layer_norm_bias', 'bias') |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.