[Megatron-Bridge] Add support for Mamba SFT posttraining with Megatron-Bridge

examples/megatron/configs/MI300X/deepseek_v2-BF16-pretrain.yaml

@@ -95,5 +95,5 @@ modules:
       turbo_sync_free_moe_stage: 2
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/deepseek_v2-FP8-pretrain.yaml

@@ -94,8 +94,8 @@ modules:
       turbo_sync_free_moe_stage: 2
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml

@@ -90,5 +90,5 @@ modules:
       turbo_sync_free_moe_stage: 2
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/deepseek_v2_lite-FP8-pretrain.yaml

@@ -89,8 +89,8 @@ modules:
       turbo_sync_free_moe_stage: 2
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/deepseek_v3-BF16-pretrain.yaml

@@ -80,5 +80,5 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/deepseek_v3-FP8-pretrain.yaml

@@ -79,8 +79,8 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/llama2_70B-BF16-pretrain.yaml

@@ -77,5 +77,5 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/llama2_70B-FP8-pretrain.yaml

@@ -77,8 +77,8 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

@@ -80,5 +80,5 @@ modules:
       # sequence_parallel: 1
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/llama2_7B-FP8-pretrain.yaml

@@ -80,8 +80,8 @@ modules:
       # sequence_parallel: 1
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml

@@ -76,5 +76,5 @@ modules:
       use_turbo_attention: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml

@@ -76,8 +76,8 @@ modules:
       use_turbo_attention: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml

@@ -72,5 +72,5 @@ modules:
       use_turbo_grouped_mlp: false
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml

@@ -72,8 +72,8 @@ modules:
       use_turbo_grouped_mlp: false
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/llama3_70B-BF16-pretrain.yaml

@@ -77,5 +77,5 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/llama3_70B-FP8-pretrain.yaml

@@ -77,8 +77,8 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/llama3_8B-BF16-pretrain.yaml

@@ -73,5 +73,5 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/llama3_8B-FP8-pretrain.yaml

@@ -73,8 +73,8 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/mamba_370M-pretrain.yaml

@@ -0,0 +1,85 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:mamba_370M-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+
+    # model to run
+    model: mamba_370M.yaml
+    overrides:
+      # log
+      wandb_project: "Primus_Mamba_Pretrain"
+      # disable_wandb: false
+      # disable_tensorboard: false
+      stderr_sink_level: DEBUG
+
+      eval_iters: 0
+
+      log_avg_skip_iterations: 2
+      log_avg_reset_interval: 50
+
+      train_iters: 50
+      micro_batch_size: 4
+      global_batch_size: 256
+
+      seq_length: 2048
+      max_position_embeddings: 2048
+
+      lr: 3.0e-4
+      min_lr: 0.0
+      lr_warmup_iters: 50000
+      lr_decay_iters: 73192188
+      lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: true
+      init_method_std: 0.02
+      norm_epsilon: 1.0e-5
+
+      # Mamba-specific: must provide spec
+      spec: ['megatron.core.models.mamba.mamba_layer_specs', 'mamba_stack_spec']
+
+      # Tokenizer
+      tokenizer_type: HuggingFaceTokenizer
+      tokenizer_model: EleutherAI/gpt-neox-20b
+
+      # parallel
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      gradient_accumulation_fusion: false
+
+      # data
+      mock_data: true
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null
+
+      # ckpt
+      finetune: false
+      auto_continue_train: false
+      load: null
+      no_load_optim: null
+      no_load_rng: null
+      save: null
+      save_interval: 20000
+      no_save_optim: null
+      no_save_rng: null
+      disable_last_saving: true
+      ckpt_format: torch
+
+      # Turbo - may need to disable for Mamba if not supported
+      enable_primus_turbo: false
+      use_turbo_attention: false
+      use_turbo_grouped_mlp: false
+
+      # Cross entropy flags
+      # cross_entropy_fusion_impl: "native"
+      # cross_entropy_loss_fusion: false

examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-BF16-pretrain.yaml

@@ -73,5 +73,5 @@ modules:
       ckpt_format: torch
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-FP8-pretrain.yaml

@@ -72,8 +72,8 @@ modules:
       ckpt_format: torch
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml

@@ -68,5 +68,5 @@ modules:
       ckpt_format: torch
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-FP8-pretrain.yaml

@@ -67,8 +67,8 @@ modules:
       ckpt_format: torch
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/qwen2.5_72B-BF16-pretrain.yaml

@@ -81,5 +81,5 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/qwen2.5_72B-FP8-pretrain.yaml

@@ -81,8 +81,8 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid

examples/megatron/configs/MI300X/qwen2.5_7B-BF16-pretrain.yaml

@@ -74,5 +74,5 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true

examples/megatron/configs/MI300X/qwen2.5_7B-FP8-pretrain.yaml

@@ -74,8 +74,8 @@ modules:
       use_turbo_grouped_mlp: true
 
       # Cross entropy flags
-      # cross_entropy_fusion_impl: "te"
-      # cross_entropy_loss_fusion: true
+      cross_entropy_fusion_impl: "te"
+      cross_entropy_loss_fusion: true
 
       # enable fp8 training
       fp8: hybrid