diff --git a/examples/megatron_bridge/configs/MI300X/gemma4_26b_pretrain.yaml b/examples/megatron_bridge/configs/MI300X/gemma4_26b_pretrain.yaml new file mode 100644 index 000000000..bf6aebf1f --- /dev/null +++ b/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 diff --git a/examples/megatron_bridge/configs/MI300X/gemma4_26b_sft_posttrain.yaml b/examples/megatron_bridge/configs/MI300X/gemma4_26b_sft_posttrain.yaml new file mode 100644 index 000000000..f51846918 --- /dev/null +++ b/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 diff --git a/examples/megatron_bridge/configs/MI300X/gemma4_31b_pretrain.yaml b/examples/megatron_bridge/configs/MI300X/gemma4_31b_pretrain.yaml new file mode 100644 index 000000000..60dceac26 --- /dev/null +++ b/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 diff --git a/examples/megatron_bridge/configs/MI300X/gemma4_31b_sft_posttrain.yaml b/examples/megatron_bridge/configs/MI300X/gemma4_31b_sft_posttrain.yaml new file mode 100644 index 000000000..1ae927733 --- /dev/null +++ b/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 diff --git a/primus/backends/megatron_bridge/patches/gemma4/__init__.py b/primus/backends/megatron_bridge/patches/gemma4/__init__.py new file mode 100644 index 000000000..fa5774c15 --- /dev/null +++ b/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"] diff --git a/primus/backends/megatron_bridge/patches/gemma4/gemma4_bridge_patches.py b/primus/backends/megatron_bridge/patches/gemma4/gemma4_bridge_patches.py new file mode 100644 index 000000000..603a125e8 --- /dev/null +++ b/primus/backends/megatron_bridge/patches/gemma4/gemma4_bridge_patches.py @@ -0,0 +1,195 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Runtime Megatron-Bridge patches for Gemma 4. + +1. ``gemma4.bridge.text_mode_for_moe`` — make the MoE path honor + ``GEMMA4_CONVERSION_MODE=text``. ``Gemma4VLBridge.provider_bridge`` only + consults the conversion mode on the dense path:: + + if not getattr(text_config, "enable_moe_block", False): + self._is_dense = True + if self._conversion_mode() == "text": + return self._build_dense_provider(text_config) + return self._build_dense_vl_provider(...) + + self._is_dense = False + # ... unconditionally builds Gemma4VLModelProvider + + The published MoE checkpoint (``google/gemma-4-26B-A4B``) sets + ``enable_moe_block=true`` and declares ``Gemma4ForConditionalGeneration``, so + it always lands on the VL provider and builds vision + audio towers. Training + then fails in the loss: the VL model returns the LLaVA-style + ``(loss, new_loss_mask)`` tuple, but ``gpt_step.forward_step`` never passes + ``loss_mask`` into the model, so ``masked_next_token_loss`` dereferences + ``None``. + +2. ``gemma4.dense.te_core_attention`` — opt-in Transformer Engine attention for + the dense path, which Bridge otherwise pins to ``LocalSpecProvider``. + +The upstream fix for (1) is two lines in ``gemma4_vl_bridge.py`` right after +``self._is_dense = False``; until that lands, patching here keeps +``third_party/Megatron-Bridge`` untouched. +""" + +from __future__ import annotations + +import copy +import os +from typing import Any, Optional + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + +_PATCHED_ATTR = "_primus_gemma4_text_moe_patched" +_TE_PATCHED_ATTR = "_primus_gemma4_dense_te_patched" + + +def _find_vl_bridge_class() -> Optional[type]: + """Locate ``Gemma4VLBridge`` without importing ``megatron.bridge.models.gemma_vl``. + + That package is stubbed by the Primus adapter whenever its NV-only deps + (megatron-energon / qwen-vl-utils) are missing, so importing it directly is + not safe. The class is reachable as a subclass of the always-real + ``Gemma4Bridge`` once Bridge's model registry has been imported. + """ + try: + from megatron.bridge.models.gemma.gemma4_bridge import Gemma4Bridge + except Exception: + return None + + pending = list(Gemma4Bridge.__subclasses__()) + while pending: + cls = pending.pop() + if cls.__name__ == "Gemma4VLBridge": + return cls + pending.extend(cls.__subclasses__()) + return None + + +@register_patch( + "gemma4.bridge.text_mode_for_moe", + backend="megatron_bridge", + phase="setup", + description="Honor GEMMA4_CONVERSION_MODE=text on the Gemma 4 MoE path, not just dense", +) +def patch_gemma4_text_mode_for_moe(ctx: PatchContext) -> None: + vl_bridge = _find_vl_bridge_class() + if vl_bridge is None: + return + + original = vl_bridge.__dict__.get("provider_bridge") + if original is None or getattr(original, _PATCHED_ATTR, False): + return + + def provider_bridge(self, hf_pretrained: Any): + hf_config = hf_pretrained.config + text_config = getattr(hf_config, "text_config", None) or hf_config + + wants_text_moe = ( + getattr(text_config, "enable_moe_block", False) + and self._conversion_mode() == "text" + and hasattr(self, "_build_moe_provider") + ) + if not wants_text_moe: + return original(self, hf_pretrained) + + self._is_dense = False + provider = self._build_moe_provider(text_config) + log_rank_0( + "[Patch:gemma4.bridge.text_mode_for_moe] GEMMA4_CONVERSION_MODE=text: " + f"built {type(provider).__name__} instead of Gemma4VLModelProvider" + ) + return provider + + setattr(provider_bridge, _PATCHED_ATTR, True) + vl_bridge.provider_bridge = provider_bridge + log_rank_0(f"[Patch:gemma4.bridge.text_mode_for_moe] Patched {vl_bridge.__name__}.provider_bridge") + + +def _build_dense_te_core_attention() -> Optional[type]: + """Build a TE core-attention class for the Gemma 4 *dense* layer spec. + + ``get_gemma4_layer_spec`` pins ``LocalSpecProvider``, so dense Gemma 4 runs on + ``DotProductAttention``: no flash attention, and the full [b, h, s, s] score + matrix is materialized and kept for backward. On 31B / seq 4096 that is ~94 GB + of the 161 GB peak, which caps the reachable micro-batch size. + + The MoE path already has ``Gemma4TEDotProductAttention``, but it keys the + sliding-vs-global decision off ``interleaved_attn_pattern`` (the MoE + convention). Dense uses ``_is_gemma4_sliding_layer`` instead, so it needs its + own subclass with the same window bookkeeping. + """ + try: + from megatron.bridge.models.gemma.modeling_gemma4 import ( + _is_gemma4_sliding_layer, + ) + from megatron.core.extensions.transformer_engine import TEDotProductAttention + except Exception: + return None + + class Gemma4DenseTEDotProductAttention(TEDotProductAttention): + """Dense Gemma 4 core attention on Transformer Engine (flash attention).""" + + def __init__( + self, config, layer_number, attn_mask_type, attention_type, attention_dropout=None, **kwargs + ): + config = copy.deepcopy(config) + # Global layers must see the full context. Gemma4DenseProvider already + # stores the TE-ready (left, right) span; only the MoE provider + # carries a bare int. + if not _is_gemma4_sliding_layer(config, layer_number): + config.window_size = None + elif not isinstance(config.window_size, tuple): + config.window_size = (config.window_size - 1, 0) + + super().__init__( + config=config, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type=attention_type, + attention_dropout=attention_dropout, + **kwargs, + ) + + return Gemma4DenseTEDotProductAttention + + +@register_patch( + "gemma4.dense.te_core_attention", + backend="megatron_bridge", + phase="setup", + description="Opt-in: run dense Gemma 4 core attention on Transformer Engine instead of LocalSpecProvider", +) +def patch_gemma4_dense_te_attention(ctx: PatchContext) -> None: + if os.environ.get("PRIMUS_GEMMA4_DENSE_ATTENTION_BACKEND", "").lower() != "te": + return + + try: + from megatron.bridge.models.gemma import gemma4_provider + except Exception: + return + + original = getattr(gemma4_provider, "get_gemma4_layer_spec", None) + if original is None or getattr(original, _TE_PATCHED_ATTR, False): + return + + te_core_attention = _build_dense_te_core_attention() + if te_core_attention is None: + log_rank_0( + "[Patch:gemma4.dense.te_core_attention] Transformer Engine unavailable; keeping local attention" + ) + return + + def get_gemma4_layer_spec(config=None): + spec = original(config) + spec.submodules.self_attention.submodules.core_attention = te_core_attention + return spec + + setattr(get_gemma4_layer_spec, _TE_PATCHED_ATTR, True) + gemma4_provider.get_gemma4_layer_spec = get_gemma4_layer_spec + log_rank_0("[Patch:gemma4.dense.te_core_attention] Dense Gemma 4 core attention -> TEDotProductAttention") diff --git a/primus/backends/megatron_bridge/recipes/gemma/__init__.py b/primus/backends/megatron_bridge/recipes/gemma/__init__.py new file mode 100644 index 000000000..fbb7e2bbc --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/gemma/__init__.py @@ -0,0 +1,21 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Gemma 4 (26B MoE and 31B Dense) recipe extensions for Megatron-Bridge.""" + +from .gemma4 import ( + gemma4_26b_finetune_config, + gemma4_26b_pretrain_config, + gemma4_31b_finetune_config, + gemma4_31b_pretrain_config, +) + +__all__ = [ + "gemma4_26b_pretrain_config", + "gemma4_31b_pretrain_config", + "gemma4_26b_finetune_config", + "gemma4_31b_finetune_config", +] diff --git a/primus/backends/megatron_bridge/recipes/gemma/gemma4.py b/primus/backends/megatron_bridge/recipes/gemma/gemma4.py new file mode 100644 index 000000000..e2bee2caa --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/gemma/gemma4.py @@ -0,0 +1,499 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Gemma 4 (26B MoE and 31B Dense) recipes for Megatron-Bridge. + +These recipes define model providers and training configurations for the +Gemma 4 family of models. The recipes follow the pattern established by +upstream Megatron-Bridge Gemma2 recipes. + +Architecture highlights: +- Sliding window attention (5 local + 1 global pattern) +- GeGLU activation (quick_geglu) +- Dual RoPE timescales (local: 10k, global: 1M) +- Logit soft capping (30.0) +- 26B: MoE with 128 experts, top-8 routing +- 31B: Dense model +""" + +import os +from contextlib import contextmanager +from typing import List, Optional, Union + +import torch +from megatron.bridge import AutoBridge +from megatron.bridge.peft.base import PEFT +from megatron.bridge.recipes.utils.dataset_utils import ( + default_peft_config, + default_squad_config, + get_blend_fields_from_data_paths, +) +from megatron.bridge.recipes.utils.optimizer_utils import ( + distributed_fused_adam_with_cosine_annealing, +) +from megatron.bridge.recipes.utils.tokenizer_utils import ( + DEFAULT_NULL_TOKENIZER_VOCAB_SIZE, +) +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ( + CheckpointConfig, + ConfigContainer, + DistributedDataParallelConfig, + GPTDatasetConfig, + LoggerConfig, + RNGConfig, + TokenizerConfig, + TrainingConfig, +) +from megatron.bridge.training.mixed_precision import ( + MixedPrecisionConfig, + bf16_mixed, + get_mixed_precision_config, +) +from typing_extensions import TypedDict, Unpack + + +@contextmanager +def _gemma4_text_conversion_mode(): + """Force AutoBridge onto the text-only Gemma 4 path. + + The published Gemma 4 checkpoints are ``Gemma4ForConditionalGeneration`` + (vision + audio towers alongside the language model), so AutoBridge's + default "auto" dispatch selects Gemma4VLBridge and builds a Gemma4VLModel. + These recipes train the language model only, so the mode is pinned to + "text" to get Gemma4DenseProvider / Gemma4ModelProvider instead. Mirrors + megatron.bridge.recipes.gemma.h100.gemma4. + """ + previous_mode = os.environ.get("GEMMA4_CONVERSION_MODE") + os.environ["GEMMA4_CONVERSION_MODE"] = "text" + try: + yield + finally: + if previous_mode is None: + os.environ.pop("GEMMA4_CONVERSION_MODE", None) + else: + os.environ["GEMMA4_CONVERSION_MODE"] = previous_mode + + +class Gemma4CommonKwargs(TypedDict, total=False): + """Typed options accepted by Gemma 4 recipe helper functions.""" + + # Core identifiers + hf_path: str + dir: Optional[str] + name: str + # Dataset configuration + data_paths: Optional[List[str]] + data_args_path: Optional[str] + train_data_path: Optional[List[str]] + valid_data_path: Optional[List[str]] + test_data_path: Optional[str] + per_split_data_args_path: Optional[str] + mock: bool + # Model configuration + tensor_model_parallel_size: int + expert_model_parallel_size: int + pipeline_model_parallel_size: int + pipeline_dtype: Optional[torch.dtype] + virtual_pipeline_model_parallel_size: Optional[int] + context_parallel_size: int + sequence_parallel: bool + use_megatron_fsdp: bool + # Training hyperparameters + train_iters: int + global_batch_size: int + micro_batch_size: int + seq_length: int + lr: float + min_lr: float + lr_warmup_iters: int + lr_decay_iters: Optional[int] + eval_interval: int + save_interval: int + use_null_tokenizer: bool + # Precision / overlap configs + precision_config: Optional[Union[MixedPrecisionConfig, str]] + comm_overlap_config: Optional[CommOverlapConfig] + + +class Gemma4FinetuneKwargs(TypedDict, total=False): + """Typed options accepted by Gemma 4 finetuning recipe helper functions.""" + + # Core identifiers + hf_path: str + dir: Optional[str] + name: str + + # Finetuning-specific + pretrained_checkpoint: Optional[str] + peft: Union[str, PEFT, None] + packed_sequence: bool + + # Training hyperparameters + train_iters: int + global_batch_size: Optional[int] + micro_batch_size: int + seq_length: Optional[int] + eval_interval: int + save_interval: int + + # Model configuration + tensor_model_parallel_size: int + expert_model_parallel_size: int + pipeline_model_parallel_size: int + pipeline_dtype: Optional[torch.dtype] + virtual_pipeline_model_parallel_size: Optional[int] + context_parallel_size: int + sequence_parallel: bool + + # Optimizer + finetune_lr: Optional[float] + min_lr: float + lr_warmup_iters: int + lr_decay_iters: Optional[int] + + # W&B logging + wandb_project: Optional[str] + wandb_entity: Optional[str] + wandb_exp_name: Optional[str] + + # Precision + precision_config: Optional[Union[MixedPrecisionConfig, str]] + + +# Pretrain Configs +def gemma4_26b_pretrain_config(**user_kwargs: Unpack[Gemma4CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Gemma 4 26B MoE. + + Architecture: 30 layers, 2816 hidden, 128 experts (top-8) + Default parallelism: TP=1, EP=8, PP=1 + """ + recommended_kwargs: Gemma4CommonKwargs = { + "hf_path": "google/gemma-4-26B-A4B", + "tensor_model_parallel_size": 1, + "expert_model_parallel_size": 8, # Critical for MoE + "pipeline_model_parallel_size": 1, + "pipeline_dtype": torch.bfloat16, + } + combined_kwargs: Gemma4CommonKwargs = {**recommended_kwargs, **user_kwargs} + return _gemma4_common(**combined_kwargs) + + +def gemma4_31b_pretrain_config(**user_kwargs: Unpack[Gemma4CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Gemma 4 31B Dense. + + Architecture: 60 layers, 5376 hidden, no MoE + Default parallelism: TP=2, EP=1, PP=1 with sequence parallelism + """ + recommended_kwargs: Gemma4CommonKwargs = { + "hf_path": "google/gemma-4-31B", + "tensor_model_parallel_size": 2, # Needed for 31B memory + "expert_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "pipeline_dtype": torch.bfloat16, + "sequence_parallel": True, # Recommended for activation memory + } + combined_kwargs: Gemma4CommonKwargs = {**recommended_kwargs, **user_kwargs} + return _gemma4_common(**combined_kwargs) + + +# Finetune Configs +def gemma4_26b_finetune_config(**user_kwargs: Unpack[Gemma4FinetuneKwargs]) -> ConfigContainer: + """Return a finetuning config for Gemma 4 26B MoE. + + Default configuration: 1 node, 8 GPUs + - LoRA/DoRA: TP=1, EP=8, PP=1, LR=1e-4 + - Full SFT: TP=1, EP=8, PP=1, LR=5e-6 + """ + # Remove hf_path from user_kwargs if present to avoid duplicate parameter + user_kwargs.pop("hf_path", None) + return _gemma4_finetune_common(hf_path="google/gemma-4-26B-A4B", **user_kwargs) + + +def gemma4_31b_finetune_config(**user_kwargs: Unpack[Gemma4FinetuneKwargs]) -> ConfigContainer: + """Return a finetuning config for Gemma 4 31B Dense. + + Default configuration: 2 nodes (SFT) or 1 node (LoRA), 8 GPUs per node + - LoRA/DoRA: TP=4, EP=1, PP=1, LR=1e-4 + - Full SFT: TP=8, EP=1, PP=2, LR=5e-6 + """ + peft_value = user_kwargs.get("peft", "lora") + is_full_sft = peft_value is None or (isinstance(peft_value, str) and peft_value.lower() == "none") + + if "tensor_model_parallel_size" not in user_kwargs: + user_kwargs["tensor_model_parallel_size"] = 8 if is_full_sft else 4 + if "pipeline_model_parallel_size" not in user_kwargs: + user_kwargs["pipeline_model_parallel_size"] = 2 if is_full_sft else 1 + + # Remove hf_path from user_kwargs if present to avoid duplicate parameter + user_kwargs.pop("hf_path", None) + return _gemma4_finetune_common(hf_path="google/gemma-4-31B", **user_kwargs) + + +def _gemma4_common( + hf_path: str, + dir: Optional[str] = None, + name: str = "default", + # Dataset configuration + data_paths: Optional[List[str]] = None, + data_args_path: Optional[str] = None, + train_data_path: Optional[List[str]] = None, + valid_data_path: Optional[List[str]] = None, + test_data_path: Optional[str] = None, + per_split_data_args_path: Optional[str] = None, + mock: bool = False, + # Model configuration + tensor_model_parallel_size: int = 1, + expert_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + pipeline_dtype: Optional[torch.dtype] = None, + virtual_pipeline_model_parallel_size: Optional[int] = None, + context_parallel_size: int = 1, + sequence_parallel: bool = False, + use_megatron_fsdp: bool = False, + # Training hyperparameters + train_iters: int = 300000, + global_batch_size: int = 32, + micro_batch_size: int = 2, + seq_length: int = 8192, + lr: float = 3e-4, + min_lr: float = 3e-5, + lr_warmup_iters: int = 500, + lr_decay_iters: Optional[int] = None, + eval_interval: int = 500, + save_interval: int = 500, + use_null_tokenizer: bool = False, + # Precision recipe + precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", + comm_overlap_config: Optional[CommOverlapConfig] = None, +) -> ConfigContainer: + """Create a pre-training configuration for Gemma 4 models.""" + + base_output_dir = dir if dir is not None else os.path.join(os.getcwd(), "nemo_experiments") + run_output_dir = os.path.join(base_output_dir, name) + checkpoint_dir = os.path.join(run_output_dir, "checkpoints") + tensorboard_dir = os.path.join(run_output_dir, "tb_logs") + + blend, blend_per_split, split = get_blend_fields_from_data_paths( + data_paths, + data_args_path, + train_data_path, + valid_data_path, + test_data_path, + per_split_data_args_path, + mock, + ) + + with _gemma4_text_conversion_mode(): + bridge = AutoBridge.from_hf_pretrained(hf_path) + model_cfg = bridge.to_megatron_provider(load_weights=False) + model_cfg.tensor_model_parallel_size = tensor_model_parallel_size + model_cfg.expert_model_parallel_size = expert_model_parallel_size + model_cfg.pipeline_model_parallel_size = pipeline_model_parallel_size + model_cfg.pipeline_dtype = pipeline_dtype + model_cfg.virtual_pipeline_model_parallel_size = virtual_pipeline_model_parallel_size + model_cfg.context_parallel_size = context_parallel_size + model_cfg.sequence_parallel = sequence_parallel + model_cfg.seq_length = seq_length + + opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=lr_warmup_iters, + lr_decay_iters=lr_decay_iters, + max_lr=lr, + min_lr=min_lr, + ) + + cfg = ConfigContainer( + model=model_cfg, + train=TrainingConfig( + train_iters=train_iters, + eval_interval=eval_interval, + eval_iters=32, + global_batch_size=global_batch_size, + micro_batch_size=micro_batch_size, + manual_gc=True, + manual_gc_interval=100, + manual_gc_eval=100, + ), + optimizer=opt_config, + scheduler=scheduler, + ddp=DistributedDataParallelConfig( + check_for_nan_in_grad=True, + use_distributed_optimizer=True, + use_megatron_fsdp=use_megatron_fsdp, + ), + dataset=GPTDatasetConfig( + random_seed=1234, + reset_attention_mask=False, + reset_position_ids=False, + eod_mask_loss=False, + seq_length=seq_length, + num_dataset_builder_threads=1, + blend=blend, + blend_per_split=blend_per_split, + split=split, + data_sharding=True, + dataloader_type="single", + skip_getting_attention_mask_from_dataset=True, + ), + logger=LoggerConfig( + log_interval=10, + tensorboard_dir=tensorboard_dir, + log_timers_to_tensorboard=True, + ), + tokenizer=TokenizerConfig( + tokenizer_type="NullTokenizer" if use_null_tokenizer else "HuggingFaceTokenizer", + tokenizer_model=hf_path if not use_null_tokenizer else None, + vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE if use_null_tokenizer else None, + ), + checkpoint=CheckpointConfig( + save_interval=save_interval, + save=checkpoint_dir, + load=checkpoint_dir, + ckpt_format="torch_dist", + fully_parallel_save=True, + ), + rng=RNGConfig(seed=1234), + comm_overlap=comm_overlap_config, + mixed_precision=( + get_mixed_precision_config(precision_config) + if isinstance(precision_config, str) + else precision_config + ), + ) + + return cfg + + +def _gemma4_finetune_common( + hf_path: str, + dir: Optional[str] = None, + name: str = "default", + # Core model configuration + tensor_model_parallel_size: int = 1, + expert_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + pipeline_dtype: Optional[torch.dtype] = None, + virtual_pipeline_model_parallel_size: Optional[int] = None, + context_parallel_size: int = 1, + sequence_parallel: bool = False, + # Finetuning-specific params + pretrained_checkpoint: Optional[str] = None, + peft: Union[str, PEFT, None] = "lora", + packed_sequence: bool = False, + # Training params + train_iters: int = 100, + global_batch_size: Optional[int] = None, + micro_batch_size: int = 1, + seq_length: Optional[int] = None, + eval_interval: int = 50, + save_interval: int = 50, + # Optimizer + finetune_lr: Optional[float] = None, + min_lr: float = 0.0, + lr_warmup_iters: int = 10, + lr_decay_iters: Optional[int] = None, + # W&B logging + wandb_project: Optional[str] = None, + wandb_entity: Optional[str] = None, + wandb_exp_name: Optional[str] = None, + # Precision + precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", +) -> ConfigContainer: + """Create a finetuning configuration for Gemma 4 models.""" + + base_output_dir = dir if dir is not None else os.path.join(os.getcwd(), "nemo_experiments") + run_output_dir = os.path.join(base_output_dir, name) + checkpoint_dir = os.path.join(run_output_dir, "checkpoints") + tensorboard_dir = os.path.join(run_output_dir, "tb_logs") + + with _gemma4_text_conversion_mode(): + bridge = AutoBridge.from_hf_pretrained(hf_path) + model_cfg = bridge.to_megatron_provider(load_weights=False) + model_cfg.tensor_model_parallel_size = tensor_model_parallel_size + model_cfg.expert_model_parallel_size = expert_model_parallel_size + model_cfg.pipeline_model_parallel_size = pipeline_model_parallel_size + model_cfg.pipeline_dtype = pipeline_dtype + model_cfg.virtual_pipeline_model_parallel_size = virtual_pipeline_model_parallel_size + model_cfg.context_parallel_size = context_parallel_size + model_cfg.sequence_parallel = sequence_parallel + + if seq_length is not None: + model_cfg.seq_length = seq_length + + # Auto-determine global batch size and learning rate based on PEFT mode + is_full_sft = peft is None or (isinstance(peft, str) and peft.lower() == "none") + + if global_batch_size is None: + global_batch_size = 8 if is_full_sft else 128 + + if finetune_lr is None: + finetune_lr = 5e-6 if is_full_sft else 1e-4 + + opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=lr_warmup_iters, + lr_decay_iters=lr_decay_iters if lr_decay_iters is not None else train_iters, + max_lr=finetune_lr, + min_lr=min_lr, + ) + + # Configure PEFT if enabled + peft_config = default_peft_config(peft) if peft else None + + # Dataset configuration for finetuning (SQuAD-style by default) + dataset_config = default_squad_config(packed_sequence=packed_sequence) + dataset_config.seq_length = model_cfg.seq_length if seq_length is None else seq_length + + cfg = ConfigContainer( + model=model_cfg, + train=TrainingConfig( + train_iters=train_iters, + eval_interval=eval_interval, + eval_iters=10, + global_batch_size=global_batch_size, + micro_batch_size=micro_batch_size, + manual_gc=True, + manual_gc_interval=100, + manual_gc_eval=100, + ), + optimizer=opt_config, + scheduler=scheduler, + ddp=DistributedDataParallelConfig( + check_for_nan_in_grad=True, + use_distributed_optimizer=not bool(peft_config), # PEFT doesn't use distributed optimizer + ), + dataset=dataset_config, + logger=LoggerConfig( + log_interval=1, + tensorboard_dir=tensorboard_dir, + log_timers_to_tensorboard=True, + wandb_project=wandb_project, + wandb_entity=wandb_entity, + wandb_exp_name=wandb_exp_name, + ), + tokenizer=TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=hf_path, + ), + checkpoint=CheckpointConfig( + save_interval=save_interval, + save=checkpoint_dir, + load=pretrained_checkpoint if pretrained_checkpoint else checkpoint_dir, + ckpt_format="torch_dist", + fully_parallel_save=True, + ), + rng=RNGConfig(seed=1234), + peft=peft_config, + mixed_precision=( + get_mixed_precision_config(precision_config) + if isinstance(precision_config, str) + else precision_config + ), + ) + + return cfg diff --git a/primus/configs/models/megatron_bridge/gemma4_26b.yaml b/primus/configs/models/megatron_bridge/gemma4_26b.yaml new file mode 100644 index 000000000..5d63a0a3d --- /dev/null +++ b/primus/configs/models/megatron_bridge/gemma4_26b.yaml @@ -0,0 +1,6 @@ +recipe: gemma.gemma4 +flavor: gemma4_26b_finetune_config +hf_path: google/gemma-4-26B-A4B + +dataset: + dataset_name: "rajpurkar/squad" diff --git a/primus/configs/models/megatron_bridge/gemma4_31b.yaml b/primus/configs/models/megatron_bridge/gemma4_31b.yaml new file mode 100644 index 000000000..af6d967e3 --- /dev/null +++ b/primus/configs/models/megatron_bridge/gemma4_31b.yaml @@ -0,0 +1,6 @@ +recipe: gemma.gemma4 +flavor: gemma4_31b_finetune_config +hf_path: google/gemma-4-31B + +dataset: + dataset_name: "rajpurkar/squad" diff --git a/runner/helpers/hooks/train/posttrain/megatron_bridge/requirements-megatron-bridge.txt b/runner/helpers/hooks/train/posttrain/megatron_bridge/requirements-megatron-bridge.txt index 6e7648a5e..b7ca204a7 100644 --- a/runner/helpers/hooks/train/posttrain/megatron_bridge/requirements-megatron-bridge.txt +++ b/runner/helpers/hooks/train/posttrain/megatron_bridge/requirements-megatron-bridge.txt @@ -2,7 +2,11 @@ # Extracted from third_party/Megatron-Bridge/pyproject.toml # Transformers (pinned version for compatibility) -transformers==4.57.6 +# Megatron-Bridge v0.6.0 requires transformers>=5.8,<=5.12.1. Gemma 4 does not +# exist on the 4.x line at all -- transformers gained models/gemma4 only around +# 5.6 -- and AutoBridge dispatches on the HF architecture class, so the 5.x +# floor is a hard requirement for Gemma4ForConditionalGeneration. +transformers==5.12.1 # Core dependencies from [project.dependencies] qwen-vl-utils diff --git a/runner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txt b/runner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txt index f9f717d76..2bc1cb896 100644 --- a/runner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txt +++ b/runner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txt @@ -3,7 +3,11 @@ # Extracted from third_party/Megatron-Bridge/pyproject.toml # Transformers (pinned version for compatibility) -transformers==4.57.6 +# Megatron-Bridge v0.6.0 requires transformers>=5.8,<=5.12.1. Gemma 4 does not +# exist on the 4.x line at all -- transformers gained models/gemma4 only around +# 5.6 -- and AutoBridge dispatches on the HF architecture class, so the 5.x +# floor is a hard requirement for Gemma4ForConditionalGeneration. +transformers==5.12.1 # Core dependencies from [project.dependencies] qwen-vl-utils diff --git a/third_party/Megatron-Bridge b/third_party/Megatron-Bridge index 9577b1280..51885cf13 160000 --- a/third_party/Megatron-Bridge +++ b/third_party/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 9577b1280eaadd60b9d7b0ce6df09ac80e87e323 +Subproject commit 51885cf132b2814188b6855c25a8588254274c2a