[megatron] Implement gemma4

docs/02-user-guide/pretraining.md

@@ -358,3 +358,4 @@ The tables above in the Megatron, TorchTitan, and MaxText sections are curated M
 - [CLI reference](cli-reference.md): launcher usage
 - [Configuration system](configuration-system.md): YAML merge rules
 - Backend parameter references: [Megatron parameters](../03-configuration-reference/megatron-parameters.md), [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md), [MaxText parameters](../03-configuration-reference/maxtext-parameters.md)
+- [Gemma 4 pre-training on MI455X](../04-technical-guides/Gemma4_Bridge_MI455X_README.md): single-GPU gfx1250 Megatron-Bridge proxies for Gemma-4-26B MoE and 31B dense

docs/04-technical-guides/README.md

@@ -20,6 +20,7 @@ Deep technical topics for advanced users.
 - [Hybrid models](hybrid-models/README.md): Hylo hybrid hybrid recurrent-attention (Mamba/KDA/GDN + MLA) models, FLA-parity recipes, and checkpoint conversion
 - [Native SFT and LoRA](native-sft-lora.md): Megatron-native SFT/LoRA runbook (BF16 / FP8 / FP4), no Megatron-Bridge dependency
 - [Native SFT LoRA on MI455X](LoRA_Native_Trainer_MI455_README.md): 1-GPU gfx1250 recipes (Llama-3.2-1B + 70B/72B/235B proxies)
+- [Gemma 4 pre-training on MI455X](Gemma4_Bridge_MI455X_README.md): 1-GPU gfx1250 Megatron-Bridge proxies (26B MoE + 31B dense), hipBLASLt Tensile fix
 
 ---
 

docs/README.md

@@ -80,6 +80,7 @@ Deep technical topics for advanced users.
 - [Diffusion models](./04-technical-guides/diffusion-models/README.md): Flux diffusion architecture, data pipeline, and FP8 / MXFP4 training
 - [Native SFT and LoRA](./04-technical-guides/native-sft-lora.md): Megatron-native SFT/LoRA runbook (BF16 / FP8 / FP4), no Megatron-Bridge dependency
 - [Native SFT LoRA on MI455X](./04-technical-guides/LoRA_Native_Trainer_MI455_README.md): 1-GPU gfx1250 native LoRA recipes and launch notes
+- [Gemma 4 pre-training on MI455X](./04-technical-guides/Gemma4_Bridge_MI455X_README.md): 1-GPU gfx1250 Megatron-Bridge Gemma 4 proxies and launch notes
 
 ### [Operations](./05-operations/)
 
@@ -137,6 +138,7 @@ Long-form articles on the reasoning behind Primus features, most of them publish
 | Train a diffusion (Flux) model | [Diffusion models](./04-technical-guides/diffusion-models/README.md) |
 | Fine-tune with native SFT / LoRA | [Native SFT and LoRA](./04-technical-guides/native-sft-lora.md) |
 | Run native LoRA on MI455X (gfx1250) | [Native SFT LoRA on MI455X](./04-technical-guides/LoRA_Native_Trainer_MI455_README.md) |
+| Pre-train Gemma 4 on MI455X (gfx1250) | [Gemma 4 pre-training on MI455X](./04-technical-guides/Gemma4_Bridge_MI455X_README.md) |
 | Auto-tune my training configuration | [Tuning agent](./02-user-guide/tuning-agent.md) |
 | Profile a training run | [Profiling and observability](./04-technical-guides/profiling-and-observability.md) |
 | Track experiments (WandB/MLflow/TensorBoard) | [Logging and experiment tracking](./04-technical-guides/logging-and-experiment-tracking.md) |

docs/sphinx/_toc.yml.in

@@ -108,6 +108,8 @@ subtrees:
         title: Native SFT LoRA
       - file: 04-technical-guides/LoRA_Native_Trainer_MI455_README.md
         title: Native SFT LoRA on MI455X
+      - file: 04-technical-guides/Gemma4_Bridge_MI455X_README.md
+        title: Gemma 4 pre-training on MI455X
       - file: 04-technical-guides/hybrid-models/README.md
         title: Hybrid recurrent-attention models
         subtrees:

examples/megatron_bridge/configs/MI300X/gemma4_26b_pretrain.yaml

@@ -0,0 +1,63 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gemma4_26b_pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron_bridge
+    config: pretrain_trainer.yaml
+
+    # Model to run
+    model: gemma4_26b.yaml
+
+    overrides:
+      stderr_sink_level: DEBUG
+
+      # Recipe override: use pretrain config instead of finetune
+      flavor: gemma4_26b_pretrain_config
+      dataset: null
+
+      # Training configuration
+      train_iters: 50
+      global_batch_size: 512
+      micro_batch_size: 1
+      seq_length: ${PRIMUS_SEQ_LENGTH:8192}
+
+      # Nested overrides
+      log_interval: 1
+      eval_interval: 500
+      eval_iters: 0
+      skip_save: true
+
+      # Optimizer
+      lr: 1.0e-5
+      min_lr: 0.0
+      lr_warmup_iters: 2
+      lr_decay_iters: null
+
+      # Parallelism - MoE with expert parallelism.
+      # TP=1/EP=8 runs out of memory at this sequence length because the
+      # non-expert weights are replicated on every rank; TP=2/EP=4 fits.
+      # TP=1/EP=8 is still usable at shorter sequence lengths, and is the faster
+      # of the two: the expert all-to-all is latency-bound on small messages, so
+      # spreading the experts as wide as possible and pushing as many tokens as
+      # possible through each dispatch both help. Activation recompute is what
+      # buys the memory for the larger micro-batch:
+      #   PRIMUS_GEMMA4_SET="model.recompute_granularity=full;\
+      #     model.recompute_method=uniform;model.recompute_num_layers=1"
+      # At seq 4096 with TP=1/EP=8 that reaches micro_batch_size 6.
+      tensor_model_parallel_size: 2
+      expert_model_parallel_size: 4  # Critical for 26B MoE
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: true
+      use_megatron_fsdp: false
+      enable_primus_turbo: false
+
+      # Data
+      mock: true
+      data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null

examples/megatron_bridge/configs/MI300X/gemma4_26b_sft_posttrain.yaml

@@ -0,0 +1,54 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gemma4_26b_sft}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  post_trainer:
+    framework: megatron_bridge
+    config: sft_trainer.yaml
+
+    # Model to run
+    model: gemma4_26b.yaml
+
+    overrides:
+      stderr_sink_level: DEBUG
+
+      # Training configuration
+      train_iters: 100
+      global_batch_size: 128
+      micro_batch_size: 1
+      seq_length: ${PRIMUS_SEQ_LENGTH:2048}
+
+      # Nested overrides
+      log_interval: 1
+      eval_interval: 50
+      eval_iters: 10
+      skip_save: true
+
+      # Optimizer
+      finetune_lr: 1.0e-4
+      min_lr: 0.0
+      lr_warmup_iters: 10
+      lr_decay_iters: null
+
+      # Parallelism - MoE with expert parallelism
+      tensor_model_parallel_size: 1
+      expert_model_parallel_size: 8  # Critical for 26B MoE
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: false
+      use_megatron_fsdp: false
+      enable_primus_turbo: false
+
+      # PEFT configuration (LoRA)
+      peft: lora
+      peft_dim: 16
+      peft_alpha: 32
+
+      # Data
+      mock: true
+      data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null

examples/megatron_bridge/configs/MI300X/gemma4_31b_pretrain.yaml

@@ -0,0 +1,63 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gemma4_31b_pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron_bridge
+    config: pretrain_trainer.yaml
+
+    # Model to run
+    model: gemma4_31b.yaml
+
+    overrides:
+      stderr_sink_level: DEBUG
+
+      # Recipe override: use pretrain config instead of finetune
+      flavor: gemma4_31b_pretrain_config
+      dataset: null
+
+      # Training configuration
+      train_iters: 50
+      global_batch_size: 512
+      micro_batch_size: 1
+      seq_length: ${PRIMUS_SEQ_LENGTH:8192}
+
+      # Nested overrides
+      log_interval: 1
+      eval_interval: 500
+      eval_iters: 0
+      skip_save: true
+
+      # Optimizer
+      lr: 1.0e-5
+      min_lr: 0.0
+      lr_warmup_iters: 2
+      lr_decay_iters: null
+
+      # Parallelism - Dense model with TP and sequence parallelism
+      # TP=8 is the smallest degree that fits on a single 8xMI300X node;
+      # TP=4 leaves almost no headroom and TP=2 runs out of memory.
+      # Two opt-in patches under primus/backends/megatron_bridge/patches/gemma4
+      # each free activation memory and so allow a larger micro_batch_size, which
+      # is what throughput on this model is most sensitive to:
+      #   PRIMUS_GEMMA4_DENSE_ATTENTION_BACKEND=te  core attention on Transformer
+      #                                             Engine instead of LocalSpec
+      #   PRIMUS_GEMMA4_FUSED_NORMS=compile         fuse the RMSNorm chain and
+      #                                             drop its fp32 temporaries
+      # With both enabled at seq 4096, micro_batch_size 6 fits on one node.
+      tensor_model_parallel_size: 8
+      expert_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: true  # Recommended for activation memory
+      use_megatron_fsdp: false
+      enable_primus_turbo: false
+
+      # Data
+      mock: true
+      data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null

examples/megatron_bridge/configs/MI300X/gemma4_31b_sft_posttrain.yaml

@@ -0,0 +1,54 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:gemma4_31b_sft}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  post_trainer:
+    framework: megatron_bridge
+    config: sft_trainer.yaml
+
+    # Model to run
+    model: gemma4_31b.yaml
+
+    overrides:
+      stderr_sink_level: DEBUG
+
+      # Training configuration
+      train_iters: 100
+      global_batch_size: 128
+      micro_batch_size: 1
+      seq_length: ${PRIMUS_SEQ_LENGTH:2048}
+
+      # Nested overrides
+      log_interval: 1
+      eval_interval: 50
+      eval_iters: 10
+      skip_save: true
+
+      # Optimizer
+      finetune_lr: 1.0e-4
+      min_lr: 0.0
+      lr_warmup_iters: 10
+      lr_decay_iters: null
+
+      # Parallelism - Dense model with TP
+      tensor_model_parallel_size: 4  # For LoRA
+      expert_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: false
+      use_megatron_fsdp: false
+      enable_primus_turbo: false
+
+      # PEFT configuration (LoRA)
+      peft: lora
+      peft_dim: 16
+      peft_alpha: 32
+
+      # Data
+      mock: true
+      data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null}
+      train_data_path: null
+      valid_data_path: null
+      test_data_path: null