[megatron] Implement gemma4 - #1053
Draft
yeandy wants to merge 6 commits into
Draft
Conversation
Implement Megatron-Bridge recipe support for Gemma 4 models to enable SFT/LoRA post-training and pretraining workflows. Implementation: - Recipe module at primus/backends/megatron_bridge/recipes/gemma/ - gemma4.py with pretrain and finetune flavors for both models - Follows upstream Gemma 2 recipe pattern - Model configs: gemma4_26b.yaml and gemma4_31b.yaml - Example configs for SFT and pretraining on MI300X Architecture support: - 26B MoE: 30 layers, 2816 hidden, 128 experts (top-8) - Default: TP=1, EP=8, PP=1 - 31B Dense: 60 layers, 5376 hidden - Default: TP=2, EP=1, PP=1 with sequence parallelism Key features handled by AutoBridge from HF: - Sliding window attention (5 local + 1 global) - GeGLU activation - Dual RoPE timescales - Logit soft capping - MoE routing (26B) Bug fixes applied: - Fix DistributedDataParallelConfig import (megatron.bridge.training.config) - Fix duplicate hf_path parameters in finetune functions - Pop hf_path from user_kwargs to prevent parameter collision
The Gemma 4 recipe added in a7188bd calls AutoBridge.from_hf_pretrained, but the pinned Megatron-Bridge (9577b1280, Jan 2026) has no Gemma 4 bridge registered, so dispatch fails at runtime. Nothing caught this statically: the recipe imports only version-stable public API, and it was patterned on Gemma 2, which does exist at that pin. - Bump third_party/Megatron-Bridge to v0.6.0 (51885cf13). This adds gemma4_bridge.py, gemma4_provider.py and the VL variants, and moves the bundled megatron-core to 0.19.0. v0.5.0 is the first release containing Gemma 4; v0.6.0 additionally carries the GEMMA4_CONVERSION_MODE=text fix for Gemma 4 MoE checkpoints, which the 26B needs. - Pin transformers to 5.12.1 in both bridge hook requirement files. Gemma 4 does not exist anywhere on the 4.x line, and AutoBridge dispatches on the HF architecture class, so a 5.x floor is required rather than preferred. - Force GEMMA4_CONVERSION_MODE=text around both AutoBridge calls. The published Gemma 4 checkpoints are Gemma4ForConditionalGeneration, so the default "auto" dispatch selects Gemma4VLBridge and builds a Gemma4VLModel with vision and audio towers instead of the language model. Mirrors megatron.bridge.recipes.gemma.h100.gemma4. Still requires container validation: a mock-data smoke run, ROCm transformer-engine against megatron-core 0.19.0, and a revalidation pass over the other bridge models that share the transformers pin.
…v0.6.1 Bridge v0.6.1 deleted megatron.bridge.recipes.utils.finetune_utils and moved default_peft_config / default_squad_config into dataset_utils, so importing the Gemma 4 recipe module raised ModuleNotFoundError and no Gemma 4 flavor could be loaded at all. The failure was hard to place because _resolve_recipe() catches ImportError and continues to the next candidate module, so the real cause surfaced only as "Function 'gemma4_26b_pretrain_config' not found".
Gemma4VLBridge.provider_bridge consults GEMMA4_CONVERSION_MODE only on the dense branch. Both published checkpoints declare Gemma4ForConditionalGeneration, so the 26B-A4B (enable_moe_block=true) always fell through to Gemma4VLModelProvider and built vision + audio towers even when the recipe asked for text-only pretraining. Training then died 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 dereferenced None. Add the missing MoE branch as a runtime patch, reusing Bridge's own _build_moe_provider so the HF-to-provider field mapping is not duplicated. The upstream fix is two lines in gemma4_vl_bridge.py; patching here avoids writing to the third_party submodule, matching the approach already used for MLPerf. Also add an opt-in Transformer Engine core attention for the dense path, enabled with PRIMUS_GEMMA4_DENSE_ATTENTION_BACKEND=te. Bridge pins dense Gemma 4 to LocalSpecProvider, so it runs plain DotProductAttention with no flash attention and keeps the full score matrix for backward; TE frees enough activation memory to raise micro_batch_size. Only core_attention is swapped, since a full TESpecProvider would replace linear_qkv with a fused-layernorm variant that collides with the dense layer's separate input_layernorm.
Both example configs ran out of memory as written. 31B dense used TP=2, which cannot fit; TP=4 leaves almost no headroom, so the smallest workable degree is TP=8. 26B MoE used TP=1/EP=8 at seq 8192. Expert parallelism shards only the experts, so every rank still holds a full copy of the non-expert weights. TP=2/EP=4 fits; TP=1/EP=8 remains fine at shorter sequence lengths. Enable sequence parallelism alongside TP to keep activation memory down.
The file predates pre-commit being run on this branch and failed both hooks at HEAD. No functional change: import grouping per isort --profile black, plus black wrapping three call sites that exceeded the 110-column limit.
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Comment on lines
+53
to
+57
| from megatron.bridge.training.mixed_precision import ( | ||
| MixedPrecisionConfig, | ||
| bf16_mixed, | ||
| get_mixed_precision_config, | ||
| ) |
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
1 participant
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.