[megatron] Implement gemma4

examples/megatron_bridge/configs/MI300X/gemma4_26b_pretrain.yaml

@@ -0,0 +1,56 @@
+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.
+      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,58 @@
+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.
+      # Setting PRIMUS_GEMMA4_DENSE_ATTENTION_BACKEND=te moves core attention
+      # onto Transformer Engine, which frees enough activation memory to raise
+      # micro_batch_size (see primus/backends/megatron_bridge/patches/gemma4).
+      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

primus/backends/megatron_bridge/patches/gemma4/__init__.py

@@ -0,0 +1,15 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Megatron-Bridge Gemma 4 patches package.
+"""
+
+from primus.backends.megatron_bridge.patches.gemma4 import (  # noqa: F401
+    gemma4_bridge_patches,
+)
+
+__all__ = ["gemma4_bridge_patches"]