diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 357dd8d5e0..ab85a57ed6 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -716,14 +716,26 @@ def _load_megatron_model(args): ), **_hf_revision_kwargs(args.hf_revision), ) - model_provider = bridge.to_megatron_provider(load_weights=False) - model_provider.tensor_model_parallel_size = tp - model_provider.pipeline_model_parallel_size = pp - model_provider.expert_model_parallel_size = ep - model_provider.expert_tensor_parallel_size = etp - model_provider.pipeline_dtype = torch.bfloat16 - model_provider.finalize() - model_provider.initialize_model_parallel(seed=0) + if getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False): + model_config = bridge.get_model_config() + transformer = model_config.transformer + transformer.tensor_model_parallel_size = tp + transformer.pipeline_model_parallel_size = pp + transformer.expert_model_parallel_size = ep + transformer.expert_tensor_parallel_size = etp + transformer.pipeline_dtype = torch.bfloat16 + transformer.params_dtype = torch.bfloat16 + model_config.finalize() + bridge._get_or_initialize_pg_collection(transformer) + else: + model_provider = bridge.to_megatron_provider(load_weights=False) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0) megatron_model = bridge.load_megatron_model( args.megatron_model_path, mp_overrides={ @@ -744,14 +756,29 @@ def _load_megatron_model(args): ), **_hf_revision_kwargs(args.hf_revision), ) - model_provider = bridge.to_megatron_provider(load_weights=True) - model_provider.tensor_model_parallel_size = tp - model_provider.pipeline_model_parallel_size = pp - model_provider.expert_model_parallel_size = ep - model_provider.expert_tensor_parallel_size = etp - model_provider.pipeline_dtype = torch.bfloat16 - model_provider.finalize() - megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False) + if getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False): + model_config = bridge.get_model_config() + transformer = model_config.transformer + transformer.tensor_model_parallel_size = tp + transformer.pipeline_model_parallel_size = pp + transformer.expert_model_parallel_size = ep + transformer.expert_tensor_parallel_size = etp + transformer.pipeline_dtype = torch.bfloat16 + transformer.params_dtype = torch.bfloat16 + megatron_model = bridge.get_model( + model_config, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + else: + model_provider = bridge.to_megatron_provider(load_weights=True) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.finalize() + megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False) # Workaround: disable MTP for inference (causes hangs on NCCL collectives) for m in megatron_model: diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml new file mode 100644 index 0000000000..c41b1771a5 --- /dev/null +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -0,0 +1,418 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: muse_glimmer_30b +summary: > + Performance scope: pretrain_performance.H100 uses the tuned 32-H100 + fixed-shape mock-token dense-decoder recipe; functional training timing and + throughput remain sanity checks. This card covers the complete Muse Glimmer + 30B vision-language checkpoint through the builder-backed model-config, + model-builder, and model bridge path. The modeled architecture includes the + vision tower, adapter, projection, and a native Megatron-Core HybridModel text + decoder; it does not use a legacy model provider or GPTModel wrapper. Affine + decoder and vision norms use Transformer Engine when available and a + semantics-equivalent local fallback otherwise; the architecture's weightless + norms remain parameter-free. + All six modeling gates are verified: CPU and TP8 GPU conversion, exact + round-trip tensor equality, numerical forward parity, and deterministic + image-conditioned inference. Bounded multimodal pretraining, full SFT, + post-SFT export and inference, packed 8K SFT with CP2, LoRA PEFT, and + canonical dense-decoder performance are also verified. Checkpoint resume + remains unverified. +verification_index: + model_level: + verified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - megatron_to_hf_cpu + - megatron_to_hf_gpu + - manual_forward_pass + - inference + training: + H100: + verified: [pretrain, sft, sft_export_inference, sft_long_context, peft] + unverified: [checkpoint_resume] + performance: + H100: verified +model: + hf_id: meta-models/Muse-Glimmer-30B + hf_revision: f84ecc3a0ea984a4c04542a84269e3d065350a6e # pragma: allowlist secret + architecture: MuseGlimmerForConditionalGeneration + min_transformers_version: "5.15.0" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.07-py3 + bridge_commit: 923f894c09181215d38f0856d40e5ce5e8eee8ca # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu + --nodes 1 --hf-model meta-models/Muse-Glimmer-30B + --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e + --megatron-path work/model-verification/muse-glimmer-30b/cpu-megatron + --torch-dtype bfloat16 --overwrite + last_verified: 2026-08-14 + expected_result: > + On a CPU-only node with no CUDA runtime, the command builds a native + 52-layer HybridModel, completes all 1,228 mappings for 29,776,626,688 + parameters, and persists iter_0000000. The checkpoint reloads for CPU + export, whose exact audit covers all 1,436 source tensors and + 29,776,626,688 elements with zero missing, unexpected, shape, dtype, or + value mismatches. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 --hf-model meta-models/Muse-Glimmer-30B + --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e + --megatron-path work/model-verification/muse-glimmer-30b/gpu-megatron + --torch-dtype bfloat16 --tp 8 --overwrite + last_verified: 2026-08-14 + expected_result: > + The eight-GPU command builds the complete native HybridModel, completes + all 1,228 mappings, validates the TP8 distributed state, and persists + iter_0000000. The checkpoint reloads at TP8 for strict GPU export; its + complete round-trip audit is bitwise exact across all 1,436 tensors and + 29,776,626,688 elements. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu + --nodes 1 --hf-model meta-models/Muse-Glimmer-30B + --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e + --megatron-path work/model-verification/muse-glimmer-30b/cpu-megatron/iter_0000000 + --hf-path work/model-verification/muse-glimmer-30b/cpu-hf-export + --torch-dtype bfloat16 --overwrite + last_verified: 2026-08-14 + expected_result: > + Strict export exits successfully. All 1,436 BF16 tensors and + 29,776,626,688 elements match the pinned source exactly in keys, shapes, + dtypes, and values (maximum absolute difference 0). Transformers 5.15.0 + reloads the result as MuseGlimmerForConditionalGeneration with no missing, + unexpected, mismatched, or error keys and the same parameter count. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 --hf-model meta-models/Muse-Glimmer-30B + --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e + --megatron-path work/model-verification/muse-glimmer-30b/gpu-megatron/iter_0000000 + --hf-path work/model-verification/muse-glimmer-30b/gpu-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 8 --overwrite + last_verified: 2026-08-14 + expected_result: > + Strict distributed export exits successfully. Every language, vision, + adapter, and projection tensor matches the pinned source bitwise: 1,436 + tensors and 29,776,626,688 elements with zero missing, unexpected, + shape, dtype, or value mismatches. Transformers 5.15.0 strictly reloads + the result with the same class and parameter count. + + manual_forward_pass: + status: verified + precision: bf16 + command: > + ./scripts/inference/infer.sh --nodes 1 --gpus-per-node 8 + --task model-comparison + --hf_model_path meta-models/Muse-Glimmer-30B + --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e + --model_class MuseGlimmerForConditionalGeneration + --megatron_model_path + work/model-verification/muse-glimmer-30b/gpu-megatron/iter_0000000 + --image_path docs/images/tp1.png + --prompt "What type of neural network block is shown in this image?" + --tp 8 --pp 1 --ep 1 --etp 1 + last_verified: 2026-08-14 + expected_result: > + The repository image is 850 by 477 pixels with SHA-256 + a7a35c8abae1e6e88d4f38bd593a56e5c1f4694da9bec5cfeb0cf29cd49256b2. + Both implementations consume 584 padded input positions and a + 2040-by-1176 processed pixel tensor. The Hugging Face and Megatron + next-token predictions match at token 328 (' to'). + Logit cosine similarity is 0.999911, above the 0.99 gate; maximum and + mean absolute logit differences are 0.734375 and 0.180013, respectively. + + inference: + status: verified + precision: bf16 + command: > + ./scripts/inference/infer.sh --nodes 1 --gpus-per-node 8 + --task vlm-generation --hf_model_path meta-models/Muse-Glimmer-30B + --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e + --megatron_model_path + work/model-verification/muse-glimmer-30b/gpu-megatron/iter_0000000 + --image_path docs/images/tp1.png + --prompt "What type of neural network block is shown in this image?" + --max_new_tokens 64 --deterministic --tp 8 --pp 1 --ep 1 --etp 1 + last_verified: 2026-08-14 + expected_result: |- + Two independent greedy runs reload the complete TP8 HybridModel + checkpoint and use the same repository image (SHA-256 + a7a35c8abae1e6e88d4f38bd593a56e5c1f4694da9bec5cfeb0cf29cd49256b2). + Both produce finite logits and the same exact 64-token completion without + EOS. The response identifies the image as a "Pre-LayerNorm Transformer + block" and accurately describes its attention, feed-forward, and residual + paths. The literal completion, represented with JSON newline and quote + escapes, is " to=self\u003c|message|\u003eWhat type of neural network block + is shown in this image?\n\nThe image shows \u0022Pre-Layernorm Transformer + Block\u0022. Transformers. Shows Layer Norm before Self-Attention, then + Multi-Headed Attention, Linear, Dropout, residual add, then Layer Norm, + Linear, Activation, Linear, Dropout," + + pretrain: + H100: + status: verified + precision: bf16 + bridge_commit: 8be93c49130954e3257158db774d59075d2beaa6 # pragma: allowlist secret + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 4 --gpus-per-node 8 + --recipe muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config + --mode pretrain --max_steps 100 + --save_dir work/model-verification/muse-glimmer-30b/pretrain-reference + --save_interval 50 + checkpoint.load=null checkpoint.finetune=false + checkpoint.save_optim=true checkpoint.save_rng=true + checkpoint.async_save=false ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + logger.save_config_filepath=work/model-verification/muse-glimmer-30b/pretrain-reference-resolved.yaml + last_verified: 2026-08-12 + metrics: + initial_loss: 12.23977 + final_loss: 1.4103 + last_10_steps_step_time_ms_avg: 1035.34 + last_10_steps_model_tflops_per_gpu_avg: 42.49 + last_10_steps_tokens_per_second_per_gpu_avg: 247.262 + expected_result: > + On 32 H100 GPUs at TP8/PP2/DP2, the command trains the complete + randomly initialized HybridModel for exactly 100 optimizer steps on + the revision-pinned public CORD-v2 multimodal dataset. All losses and + gradient norms remain finite, with zero skipped and zero NaN + iterations. Loss decreases from 12.23977 to 1.410300; steps 91-100 + average 1,035.34 ms, 42.49 model TFLOP/s/GPU, and 247.262 token + slots/s/GPU. The persisted post-setup config records the exact runtime. + Complete step-50 and step-100 checkpoints each contain 32 nonempty + distributed shards and 1,550 state entries, including model, optimizer, + scheduler, and RNG state, plus metadata, train state, and the native + HybridModel run config. + + sft: + H100: + status: verified + precision: bf16 + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 4 --gpus-per-node 8 + --recipe muse_glimmer_30b_sft_32gpu_h100_bf16_config --mode sft + --pretrained_checkpoint + work/model-verification/muse-glimmer-30b/gpu-megatron/iter_0000000 + --max_steps 100 + --save_dir work/model-verification/muse-glimmer-30b/sft-checkpoints + --save_interval 100 checkpoint.load=null + checkpoint.save_optim=false checkpoint.save_rng=false + checkpoint.async_save=false ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + logger.save_config_filepath=work/model-verification/muse-glimmer-30b/sft-resolved.yaml + last_verified: 2026-08-12 + metrics: + initial_loss: 0.9772241 + final_loss: 0.02695329 + last_10_steps_step_time_ms_avg: 1826.64 + last_10_steps_model_tflops_per_gpu_avg: 94.03 + last_10_steps_tokens_per_second_per_gpu_avg: 560.592 + expected_result: > + On 32 H100 GPUs at TP8/PP2/DP2, the command loads the complete imported + HybridModel and performs exactly 100 full-model image-conditioned SFT + optimizer steps on the revision-pinned public CORD-v2 dataset. All + losses and gradient norms remain finite, with zero skipped and zero NaN + iterations. Loss decreases from 0.9772241 to 0.02695329; steps 91-100 + average 1,826.64 ms, 94.03 model TFLOP/s/GPU, and 560.592 token + slots/s/GPU. The complete 32-shard step-100 full-model checkpoint, + metadata, train state, and resolved config are saved and reloadable. + + sft_export_inference: + H100: + status: verified + precision: bf16 + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 8 --tp 8 --pp 2 --ep 1 --etp 1 + --hf-model meta-models/Muse-Glimmer-30B + --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e + --megatron-path + work/model-verification/muse-glimmer-30b/sft-checkpoints/iter_0000100 + --hf-path work/model-verification/muse-glimmer-30b/sft-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --overwrite + - > + ./scripts/inference/infer.sh --task hf-inference + --nodes 1 --gpus-per-node 1 + --hf-model work/model-verification/muse-glimmer-30b/sft-hf-export + --image docs/images/tp1.png + --prompt "What type of neural network block is shown in this image?" + --max-new-tokens 64 --chat-template + last_verified: 2026-08-12 + expected_result: |- + The first command completes all 1,228 MuseGlimmerBridge mappings and + writes a complete two-shard BF16 Hugging Face VLM. Transformers 5.15.0 + strictly reloads all 1,436 tensors with no missing, unexpected, + mismatched, or error keys. The second command performs one bounded + greedy image-conditioned generation and produces exactly 64 new tokens. + The literal completion, represented with JSON angle-bracket escapes, + is: " \u003ctransformer_block\u003e\u003ctransformer_block_price\u003e35\u003c/transformer_block_price\u003e\u003ctransformer_block_nm\u003ePre-Layernorm Transformer Block\u003c/transformer_block_nm\u003e\u003ctransformer_block_cnt\u003eNx\u003c/transformer_block_cnt\u003e\u003c/transformer_block\u003e\u003ctransformer_block\u003e\u003ctransformer_sub\u003e\u003ctransformer_nm\u003eLinear". + + sft_long_context: + H100: + status: verified + precision: bf16 + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret + enabled_features: + sequence_packing: in_batch + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 4 --gpus-per-node 8 + --recipe muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config + --mode sft --pretrained_checkpoint + work/model-verification/muse-glimmer-30b/gpu-megatron/iter_0000000 + --max_steps 100 + --save_dir work/model-verification/muse-glimmer-30b/long-sft-checkpoints + --save_interval 100 checkpoint.load=null + checkpoint.save_optim=false checkpoint.save_rng=false + checkpoint.async_save=false ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + logger.save_config_filepath=work/model-verification/muse-glimmer-30b/long-sft-resolved.yaml + last_verified: 2026-08-12 + metrics: + initial_loss: 1.034497 + final_loss: 0.03382061 + last_10_steps_step_time_ms_avg: 1419.91 + last_10_steps_model_tflops_per_gpu_avg: 16.07 + last_10_steps_tokens_per_second_per_gpu_avg: 1442.345 + expected_result: > + On 32 H100 GPUs, the command loads the complete imported HybridModel + and completes exactly 100 full-model SFT optimizer steps at sequence + length 8192 with in-batch packing, TP1/PP4/CP2, A2A context-parallel + communication, and selective core-attention + recompute. Every loss and gradient norm is finite, with zero skipped + and zero NaN iterations. Loss decreases from 1.034497 to 0.03382061; + steps 91-100 average 1,419.91 ms, 16.07 model TFLOP/s/GPU, and + 1,442.345 token slots/s/GPU. The complete step-100 full-model + checkpoint, metadata, train state, and resolved config are saved and + reloadable. + + peft: + H100: + status: verified + precision: bf16 + enabled_features: {} + bridge_commit: 8be93c49130954e3257158db774d59075d2beaa6 # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 + --recipe muse_glimmer_30b_peft_8gpu_h100_bf16_config --mode lora + --pretrained_checkpoint + work/model-verification/muse-glimmer-30b/gpu-megatron/iter_0000000 + --max_steps 100 + --save_dir work/model-verification/muse-glimmer-30b/peft-checkpoints + --save_interval 100 checkpoint.load=null + checkpoint.async_save=false ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + logger.save_config_filepath=work/model-verification/muse-glimmer-30b/peft-resolved.yaml + last_verified: 2026-08-12 + metrics: + initial_loss: 0.7263421 + final_loss: 0.03298394 + last_10_steps_step_time_ms_avg: 4823.07 + last_10_steps_model_tflops_per_gpu_avg: 283.36 + last_10_steps_tokens_per_second_per_gpu_avg: 1698.503 + expected_result: > + The command loads the complete imported native HybridModel checkpoint, + freezes the base model, and inserts LoRA dim 8 / alpha 16 adapters only + into the native linear_qkv and linear_proj attention projections. All + 100 optimizer-step rows are present exactly once with finite loss, zero + skipped iterations, and zero NaN iterations. Loss decreases from + 0.7263421 to 0.03298394; steps 91-100 average 4,823.07 ms, 283.36 model + TFLOP/s/GPU, and 1,698.503 token slots/s/GPU. The persisted post-setup + config records the exact runtime. The complete eight-shard + iter_0000100 checkpoint contains all 208 expected adapter tensors, its + run config, metadata, and train state. A direct reload restores model + and optimizer at step 100, completes finite step 101 with zero skipped + or NaN iterations, and saves a distinct complete eight-shard + checkpoint. + + checkpoint_resume: + H100: + status: unverified + precision: bf16 + depends_on: pretrain + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + last_10_steps_tokens_per_second_per_gpu_avg: null + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: null + expected_result: > + Resume the bounded pretraining run from an intermediate checkpoint, + restore model, optimizer, scheduler, data-order, and RNG state, match + declared sentinel losses within tolerance, and save a distinct + reloadable final checkpoint. + + pretrain_performance: + H100: + status: verified + precision: bf16 + bridge_commit: 8be93c49130954e3257158db774d59075d2beaa6 # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 4 --gpus-per-node 8 + --recipe muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config + --mode pretrain --max_steps 50 + checkpoint.async_save=false ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + logger.save_config_filepath=work/model-verification/muse-glimmer-30b/performance-resolved.yaml + last_verified: 2026-08-12 + metrics: + initial_loss: 12.26882 + final_loss: 2.442728 + last_10_steps_step_time_ms_avg: 9373.06 + last_10_steps_model_tflops_per_gpu_avg: 428.24 + last_10_steps_tokens_per_second_per_gpu_avg: 2621.983 + expected_result: > + The exact 32-H100 command completes all 50 fixed-shape mock-token dense + decoder steps at TP4/PP4/CP2, a 9|15|15|13 split of all 52 language + layers, sequence length 4096, MBS6/GBS192, BF16 model and gradient state, + and FP8 current-scaling Transformer GEMMs. The vision and projection + towers are frozen and excluded from the fixed-shape decoder workload. + Loss remains finite from 12.26882 to 2.442728 with zero skipped or NaN + iterations. Steps 41-50 average 9,373.06 ms, 428.24 model + TFLOP/s/GPU, and 2,621.983 token slots/s/GPU. The persisted post-setup + config records the exact runtime, and the process exits successfully + without checkpoint output. diff --git a/examples/models/muse_glimmer/README.md b/examples/models/muse_glimmer/README.md new file mode 100644 index 0000000000..332134443c --- /dev/null +++ b/examples/models/muse_glimmer/README.md @@ -0,0 +1,55 @@ +# Muse Glimmer 30B + +[Muse Glimmer 30B](https://huggingface.co/meta-models/Muse-Glimmer-30B) is a dense +vision-language model with a roughly 2B-parameter vision encoder and a 28B-parameter +decoder. Megatron Bridge models the complete checkpoint: the vision tower, adapter, +projection, and text decoder all participate in checkpoint conversion. + +## Modeling + +Muse Glimmer uses the builder-backed API. `MuseGlimmerBridge` translates the Hugging +Face configuration directly into `MuseGlimmerModelConfig`, and +`MuseGlimmerModelBuilder` constructs `MuseGlimmerModel` as a native Megatron-Core +`HybridModel`; no model provider or `GPTModel` wrapper is involved. + +The text decoder uses a native Megatron-Core `HybridStack` whose layer pattern selects +the Muse transformer-layer spec at every depth. It preserves GQA, +tensor/pipeline/context parallelism, 3:1 sliding/full attention, per-layer NoPE, +weightless Q/K RMSNorm, and the full-head sigmoid attention output gate. Muse-specific +residual-branch RMSNorms wrap the attention and MLP outputs without replacing +Megatron-Core's `TransformerLayer` execution path. The embedding norm, final scaled +RMSNorm, output multiplier, and tanh logit soft-cap reproduce the published decoder +math. + +The replicated vision path implements the published 50-layer encoder, half-pixel +bilinear learned-position lookup with zero padding, 2-D RoPE, 3:1 window/full attention, +2x2 pixel shuffle, two-layer GELU adapter, text-width projection, and weightless +perception RMSNorm. Its parameters are marked for tensor-parallel gradient averaging. + +## Builder-backed configuration + +```python +from megatron.bridge import AutoBridge + +bridge = AutoBridge.from_hf_pretrained( + "meta-models/Muse-Glimmer-30B", + revision="f84ecc3a0ea984a4c04542a84269e3d065350a6e", # pragma: allowlist secret +) +model_config = bridge.get_model_config() + +# Apply deployment-specific parallelism before construction. +model_config.tensor_model_parallel_size = 8 +models = bridge.get_model(model_config) +``` + +The published configuration needs Transformers 5.15 for native Hugging Face model +execution. Megatron Bridge includes only a configuration compatibility class so that +the lazy safetensors import and builder path continue to work with the repository's +currently pinned Transformers release; it does not vendor or execute the Hugging Face +model implementation. + +## Verification status + +See the [model verification card](../../model_verification_cards/muse-glimmer-30b/card.yaml) +for the pinned source revision and current evidence. The card keeps import, export, +forward correlation, inference, and training claims separate. diff --git a/pyproject.toml b/pyproject.toml index a27a89cccc..fc19a14a88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,7 @@ classifiers = [ "Topic :: Utilities", ] dependencies = [ - "transformers>=5.8,<=5.12.1", + "transformers>=5.15,<=5.15.0", "mistral-common>=1.10.0", "peft>=0.18.1", "datasets>=2.20.0", diff --git a/scripts/conversion/gpu_backend.py b/scripts/conversion/gpu_backend.py index 9ae1f49745..4e5889a617 100644 --- a/scripts/conversion/gpu_backend.py +++ b/scripts/conversion/gpu_backend.py @@ -183,6 +183,33 @@ def _configure_model_provider( model_provider.use_cpu_initialization = True +def _uses_model_builder(bridge: AutoBridge) -> bool: + """Return whether the selected bridge supports native builder construction.""" + return getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False) + + +def _configure_model_config( + model_config, + *, + tp: int, + pp: int, + ep: int, + etp: int, + dtype: torch.dtype, + use_cpu: bool = False, +) -> None: + """Apply distributed parallelism and dtype settings to a builder config.""" + transformer = model_config.transformer + transformer.tensor_model_parallel_size = tp + transformer.pipeline_model_parallel_size = pp + transformer.expert_model_parallel_size = ep + transformer.expert_tensor_parallel_size = etp + transformer.pipeline_dtype = dtype + transformer.params_dtype = dtype + if use_cpu: + transformer.use_cpu_initialization = True + + def _hf_tokenizer_kwargs(bridge: AutoBridge, *, trust_remote_code: bool) -> dict[str, object]: """Build tokenizer metadata for a saved Megatron checkpoint.""" tokenizer_kwargs: dict[str, object] = {} @@ -320,12 +347,22 @@ def import_checkpoint( torch_dtype=dtype, **revision_kwargs, ) - model_provider = bridge.to_megatron_provider(load_weights=True) - _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype) - _maybe_generate_pipeline_layout(bridge, model_provider, pp) - model_provider.finalize() - model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False) - megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False) + if _uses_model_builder(bridge): + model_config = bridge.get_model_config() + _configure_model_config(model_config, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype) + _maybe_generate_pipeline_layout(bridge, model_config, pp) + megatron_model = bridge.get_model( + model_config, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + else: + model_provider = bridge.to_megatron_provider(load_weights=True) + _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype) + _maybe_generate_pipeline_layout(bridge, model_provider, pp) + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False) + megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False) bridge.save_megatron_model( megatron_model, @@ -410,12 +447,20 @@ def export_checkpoint( # Preserve the reference wrapper's streaming state source and shard map while # exporting the checkpoint-derived architecture and vocabulary configuration. bridge.hf_pretrained.config = checkpoint_config_bridge.hf_pretrained - model_provider = bridge.to_megatron_provider(load_weights=False) - _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype, use_cpu=use_cpu) - _maybe_restore_pipeline_layout(bridge, model_provider, megatron_path, pp) - resolved_pipeline_layout = model_provider.pipeline_model_parallel_layout - model_provider.finalize() - model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False) + if _uses_model_builder(bridge): + model_config = bridge.get_model_config() + _configure_model_config(model_config, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype, use_cpu=use_cpu) + _maybe_restore_pipeline_layout(bridge, model_config, megatron_path, pp) + resolved_pipeline_layout = model_config.pipeline_model_parallel_layout + model_config.finalize() + bridge._get_or_initialize_pg_collection(model_config.transformer) + else: + model_provider = bridge.to_megatron_provider(load_weights=False) + _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype, use_cpu=use_cpu) + _maybe_restore_pipeline_layout(bridge, model_provider, megatron_path, pp) + resolved_pipeline_layout = model_provider.pipeline_model_parallel_layout + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False) model_parallel_overrides: dict[str, object] = { "tensor_model_parallel_size": tp, @@ -484,11 +529,21 @@ def roundtrip_checkpoint( torch_dtype=dtype, ) - model_provider = bridge.to_megatron_provider(load_weights=True) - _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype) - _maybe_generate_pipeline_layout(bridge, model_provider, pp) - model_provider.finalize() - model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False) - megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False) + if _uses_model_builder(bridge): + model_config = bridge.get_model_config() + _configure_model_config(model_config, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype) + _maybe_generate_pipeline_layout(bridge, model_config, pp) + megatron_model = bridge.get_model( + model_config, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + else: + model_provider = bridge.to_megatron_provider(load_weights=True) + _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype) + _maybe_generate_pipeline_layout(bridge, model_provider, pp) + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False) + megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False) _verify_roundtrip_weights(bridge, megatron_model) print_rank_0("GPU round-trip validation complete") diff --git a/scripts/inference/vlm_generation.py b/scripts/inference/vlm_generation.py index a3df69d512..482c6fb1bd 100644 --- a/scripts/inference/vlm_generation.py +++ b/scripts/inference/vlm_generation.py @@ -28,6 +28,7 @@ """ import argparse +import os import torch import torch.distributed as dist @@ -148,8 +149,65 @@ def _hf_revision_kwargs(revision: str | None) -> dict[str, str]: return {"revision": revision} if revision is not None else {} +def _uses_model_builder(bridge: AutoBridge) -> bool: + """Return whether the selected bridge supports native builder construction.""" + return getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False) + + +def _enable_deterministic_execution() -> None: + """Enable the deterministic settings required by distributed MCore inference.""" + required_environment = { + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_ALGO": "Ring", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + } + invalid = { + name: os.environ[name] + for name, expected in required_environment.items() + if name in os.environ and os.environ[name] != expected + } + if invalid: + expected = ", ".join(f"{name}={value}" for name, value in required_environment.items()) + raise ValueError(f"--deterministic requires {expected}; received {invalid}.") + os.environ.update(required_environment) + + torch.use_deterministic_algorithms(True) + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + + +def _configure_builder_model( + bridge: AutoBridge, + *, + tp: int, + pp: int, + ep: int, + etp: int, + pp_layout: str | None = None, + deterministic: bool = False, +): + """Configure a builder-backed model and initialize its process groups.""" + model_config = bridge.get_model_config() + transformer = model_config.transformer + transformer.tensor_model_parallel_size = tp + transformer.pipeline_model_parallel_size = pp + transformer.expert_model_parallel_size = ep + transformer.expert_tensor_parallel_size = etp + transformer.pipeline_dtype = torch.bfloat16 + transformer.params_dtype = torch.bfloat16 + transformer.deterministic_mode = deterministic + if pp_layout: + transformer.pipeline_model_parallel_layout = pp_layout + return model_config + + def main(args) -> None: """Run VLM inference with HuggingFace or Megatron checkpoints.""" + if args.deterministic: + _enable_deterministic_execution() maybe_initialize_distributed() tp = args.tp pp = args.pp @@ -187,17 +245,31 @@ def main(args) -> None: if args.megatron_model_path: print_rank_0(f"Loading Megatron model from: {args.megatron_model_path}") - model_provider = bridge.to_megatron_provider(load_weights=False) - model_provider.tensor_model_parallel_size = tp - model_provider.pipeline_model_parallel_size = pp - model_provider.expert_model_parallel_size = ep - model_provider.expert_tensor_parallel_size = etp - model_provider.pipeline_dtype = torch.bfloat16 - model_provider.init_model_with_meta_device = True - if args.pp_layout: - model_provider.pipeline_model_parallel_layout = args.pp_layout - model_provider.finalize() - model_provider.initialize_model_parallel(seed=0) + if _uses_model_builder(bridge): + model_config = _configure_builder_model( + bridge, + tp=tp, + pp=pp, + ep=ep, + etp=etp, + pp_layout=args.pp_layout, + deterministic=args.deterministic, + ) + model_config.finalize() + bridge._get_or_initialize_pg_collection(model_config.transformer) + else: + model_provider = bridge.to_megatron_provider(load_weights=False) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.deterministic_mode = args.deterministic + model_provider.init_model_with_meta_device = True + if args.pp_layout: + model_provider.pipeline_model_parallel_layout = args.pp_layout + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0) mp_overrides = { "tensor_model_parallel_size": tp, @@ -205,6 +277,7 @@ def main(args) -> None: "expert_model_parallel_size": ep, "expert_tensor_parallel_size": etp, "pipeline_dtype": torch.bfloat16, + "deterministic_mode": args.deterministic, } if args.pp_layout: mp_overrides["pipeline_model_parallel_layout"] = args.pp_layout @@ -215,15 +288,31 @@ def main(args) -> None: ) else: print_rank_0(f"Loading HuggingFace model from: {args.hf_model_path}") - model_provider = bridge.to_megatron_provider(load_weights=True) - model_provider.tensor_model_parallel_size = tp - model_provider.pipeline_model_parallel_size = pp - model_provider.expert_model_parallel_size = ep - model_provider.expert_tensor_parallel_size = etp - model_provider.pipeline_dtype = torch.bfloat16 - model_provider.finalize() - model_provider.initialize_model_parallel(seed=0) - model = model_provider.provide_distributed_model(wrap_with_ddp=False) + if _uses_model_builder(bridge): + model_config = _configure_builder_model( + bridge, + tp=tp, + pp=pp, + ep=ep, + etp=etp, + deterministic=args.deterministic, + ) + model = bridge.get_model( + model_config, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + else: + model_provider = bridge.to_megatron_provider(load_weights=True) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.deterministic_mode = args.deterministic + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0) + model = model_provider.provide_distributed_model(wrap_with_ddp=False) def _disable_mtp(m): m.config.mtp_num_layers = None @@ -439,6 +528,11 @@ def _disable_mtp(m): "--video_fps", type=float, default=2.0, help="Frames per second to sample from the video (default: 2.0)." ) parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code for HF model loading") + parser.add_argument( + "--deterministic", + action="store_true", + help="Require deterministic distributed kernels and seed all inference RNGs.", + ) args = parser.parse_args() main(args) diff --git a/scripts/training/recipe_metadata.py b/scripts/training/recipe_metadata.py index bc73425c27..08372bb325 100644 --- a/scripts/training/recipe_metadata.py +++ b/scripts/training/recipe_metadata.py @@ -29,6 +29,7 @@ ("qwen35_vl_", "qwen_vl"), ("deepseek_", "deepseek"), ("gpt_oss_", "gpt_oss"), + ("muse_glimmer_", "muse_glimmer"), ("nemotron", "nemotronh"), ("llama", "llama"), ("qwen", "qwen"), @@ -94,6 +95,7 @@ ("glm_45v_", "vlm_step"), ("kimi_k25_vl_", "vlm_step"), ("ministral3_", "vlm_step"), + ("muse_glimmer_", "vlm_step"), ("qwen25_vl_", "vlm_step"), ("flux_", "flux_step"), ("wan_", "wan_step"), diff --git a/skills/adding-model-support/SKILL.md b/skills/adding-model-support/SKILL.md index dc1a8c1f8e..73bccd145c 100644 --- a/skills/adding-model-support/SKILL.md +++ b/skills/adding-model-support/SKILL.md @@ -178,6 +178,36 @@ reusable across multiple model families. Shared conversion infrastructure provides hooks and base classes — subclass them locally rather than adding conditionals to shared code. +#### Prefer Transformer Engine normalization + +Before implementing a model-specific normalization with raw PyTorch operations, inspect the exact +Transformer Engine revision in `pyproject.toml`/`uv.lock`, its upstream module API, and Megatron-Core's +`TENorm` adapter. Production Bridge model code should assume Transformer Engine is installed. Use +`TENorm` or the corresponding TE module directly for LayerNorm and RMSNorm; do not add model-local +norm wrappers, copy an existing norm implementation, gate the norm on `HAVE_TE`, or provide a +`torch.nn` fallback. Preserve `eps`, parameter dtype/device initialization, sequence-parallel +metadata, and `zero_centered_gamma`; standard-gamma and zero-centered-gamma norms are not +interchangeable. The attention/linear backend and normalization choice are independent, so a local +attention spec should still use TE normalization. + +Treat a custom norm implementation as an architectural exception, not a compatibility pattern. It +is allowed only when the pinned TE API cannot represent the model's exact mathematical and parameter +contract, such as a genuinely parameter-free RMSNorm while TE requires an affine weight. Prefer a +framework-native module such as `torch.nn.RMSNorm(elementwise_affine=False)` or a direct functional +call over a model-local wrapper or copied formula. Document the missing TE capability and test that +the exception introduces no checkpoint parameter. Do not create a wrapper or fallback for ordinary +affine normalization. + +Do not introduce a dummy affine parameter merely to reach a fused kernel. For example, current TE +RMSNorm always owns a learnable `weight`, so a genuinely weightless/scaleless RMSNorm must retain a +parameter-free implementation until TE exposes matching semantics. Adding an all-ones frozen weight +would change state-dict keys, optimizer state, distributed-checkpoint schema, and conversion coverage. + +TE and reference framework kernels can differ numerically even when their architecture and weights +match. Keep exact HF↔Megatron weight round-trip as the conversion gate, then assess forward behavior +with the correlation criteria in the parity-testing skill rather than reverting to a slower norm only +to reproduce one framework's operation ordering. + #### Strategy 1: Create a local mapping subclass If the model has a layer whose weight layout doesn't match any existing mapping class, create a diff --git a/skills/adding-model-support/vlm-patterns.md b/skills/adding-model-support/vlm-patterns.md index c75525799b..f0fbf4061f 100644 --- a/skills/adding-model-support/vlm-patterns.md +++ b/skills/adding-model-support/vlm-patterns.md @@ -118,6 +118,18 @@ from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM ## VLM Model Class Patterns +### Normalization backend + +Thread the language model's `TransformerConfig` into custom Megatron vision modules so affine vision +LayerNorm/RMSNorm instances use Transformer Engine and the same initialization device as the rest of +the model. Assume TE is installed in production model code and instantiate `TENorm` or the TE module +directly. Do not add model-local wrappers, `HAVE_TE` branches, copied norm kernels, or `torch.nn` +fallbacks for ordinary affine normalization. When the HF architecture specifies a weightless norm, +first verify whether the pinned TE API supports that exact parameter-free contract; a custom +parameter-free path is an exception only when TE cannot express it. Prefer a framework-native +parameter-free module or direct functional call over a wrapper or copied formula, and do not add +checkpoint parameters to simulate it. + ### Option A: Megatron Vision Encoder (Qwen3.5 pattern) Both vision and language use Megatron modules. Full parallelism support. diff --git a/skills/parity-testing/SKILL.md b/skills/parity-testing/SKILL.md index a72c01ce4b..64bb136e32 100644 --- a/skills/parity-testing/SKILL.md +++ b/skills/parity-testing/SKILL.md @@ -112,6 +112,15 @@ Keep stricter cosine or absolute-difference thresholds as optional diagnostic targets when investigating arithmetic differences. Do not use them to downgrade a model-support card after the correlation gate passes. +### Transformer Engine normalization + +Using Transformer Engine LayerNorm/RMSNorm can change floating-point operation ordering relative to +the HF or PyTorch reference. Treat that as a forward-correlation question, not a weight-conversion +failure: state-dict round-trip must remain exact, while forward parity uses the cosine and next-token +gates above. Do not replace an available TE affine norm solely to reduce a report-only absolute logit +difference. Conversely, do not add affine parameters to weightless reference norms just to use TE; +that changes the architecture and invalidates exact mapping coverage. + ## Comparison Utilities These functions are useful when writing custom verification scripts or diff --git a/src/megatron/bridge/data/collators/registry.py b/src/megatron/bridge/data/collators/registry.py index f95e86a11a..89edb07378 100644 --- a/src/megatron/bridge/data/collators/registry.py +++ b/src/megatron/bridge/data/collators/registry.py @@ -54,6 +54,9 @@ class _ModelCollateSpec: ), "Glm4vProcessor": _ModelCollateSpec("megatron.bridge.models.glm_vl.data.collate_fn", "glm4v_collate_fn"), "KimiK25Processor": _ModelCollateSpec("megatron.bridge.models.kimi_vl.data.collate_fn", "kimi_k25_vl_collate_fn"), + "MuseGlimmerProcessor": _ModelCollateSpec( + "megatron.bridge.models.muse_glimmer.data.collate_fn", "muse_glimmer_collate_fn" + ), "deepseek-v4": _ModelCollateSpec( "megatron.bridge.models.deepseek.data.collate_fn", "deepseek_v4_collate_fn", diff --git a/src/megatron/bridge/models/__init__.py b/src/megatron/bridge/models/__init__.py index 12c409ce28..5ccb393965 100644 --- a/src/megatron/bridge/models/__init__.py +++ b/src/megatron/bridge/models/__init__.py @@ -131,6 +131,12 @@ from megatron.bridge.models.mistral import ( MistralModelProvider, ) +from megatron.bridge.models.muse_glimmer import ( + MuseGlimmerBridge, + MuseGlimmerModel, + MuseGlimmerModelBuilder, + MuseGlimmerModelConfig, +) from megatron.bridge.models.nemotron import ( NemotronBridge, ) @@ -257,6 +263,11 @@ "MiniMaxM3Bridge", "MiniMaxM3VLModel", "MiniMaxM3VLModelProvider", + # Muse Glimmer + "MuseGlimmerBridge", + "MuseGlimmerModel", + "MuseGlimmerModelBuilder", + "MuseGlimmerModelConfig", "OlMoEBridge", "OlMoEModelProvider", "NemotronHBridge", diff --git a/src/megatron/bridge/models/common/__init__.py b/src/megatron/bridge/models/common/__init__.py index 03e7396120..0a72385c24 100644 --- a/src/megatron/bridge/models/common/__init__.py +++ b/src/megatron/bridge/models/common/__init__.py @@ -16,9 +16,11 @@ BuildConfigT, ModelBuilder, ModelConfig, + ModelConfigOverrideMixin, ModelT, Serializable, compose_hooks, + deserialize_model_config, ) from megatron.bridge.models.common.unimodal import build_virtual_pipeline_stages, unimodal_build_distributed_models @@ -27,9 +29,11 @@ "BuildConfigT", "ModelBuilder", "ModelConfig", + "ModelConfigOverrideMixin", "ModelT", "Serializable", "compose_hooks", + "deserialize_model_config", "build_virtual_pipeline_stages", "unimodal_build_distributed_models", ] diff --git a/src/megatron/bridge/models/common/base.py b/src/megatron/bridge/models/common/base.py index fd091234f4..9c645d204b 100644 --- a/src/megatron/bridge/models/common/base.py +++ b/src/megatron/bridge/models/common/base.py @@ -14,7 +14,7 @@ from dataclasses import fields as dataclass_fields from dataclasses import is_dataclass -from typing import Any +from typing import Any, ClassVar from megatron.training.models.base import ( BuildConfigT, # noqa: F401 @@ -27,7 +27,87 @@ ModelConfig as _MegatronModelConfig, ) -from megatron.bridge.utils.instantiate_utils import _resolve_target, _validate_target_prefix +from megatron.bridge.models.transformer_config import TransformerConfig +from megatron.bridge.utils.instantiate_utils import _resolve_target, _validate_target_prefix, instantiate + + +def deserialize_model_config(data: dict[str, Any]) -> _MegatronModelConfig: + """Deserialize a builder config, including persisted enum and dtype values.""" + + def _restore_value(value: Any, full_key: str) -> Any: + if isinstance(value, list): + return [_restore_value(item, f"{full_key}.{index}") for index, item in enumerate(value)] + if not isinstance(value, dict): + return value + if "_target_" not in value: + return {key: _restore_value(item, f"{full_key}.{key}") for key, item in value.items()} + return _from_dict(value, full_key) + + def _from_dict(subdata: dict[str, Any], full_key: str) -> Any: + target = subdata.get("_target_") + if target is None: + raise ValueError("Cannot deserialize: missing '_target_' field") + if not isinstance(target, str): + raise ValueError(f"Cannot deserialize: '_target_' must be a string, got {type(target).__name__}") + + config_cls = _resolve_target(target, full_key=full_key, check_callable=False) + if not isinstance(config_cls, type) or not is_dataclass(config_cls): + return instantiate(subdata) + + valid_fields = {f.name for f in dataclass_fields(config_cls) if f.init} + filtered_data = { + key: _restore_value(value, f"{full_key}.{key}") + for key, value in subdata.items() + if key in valid_fields and not key.startswith("_") + } + return config_cls(**filtered_data) + + builder = data.get("_builder_") + if not isinstance(builder, str): + raise ValueError("Cannot deserialize: missing '_builder_' field") + _validate_target_prefix(target=builder, full_key="_builder_") + + result = _from_dict(data, full_key="_target_") + if not isinstance(result, _MegatronModelConfig): + raise ValueError(f"Cannot deserialize: outer target produced {type(result).__name__}, not ModelConfig") + result.builder = builder + return result + + +class ModelConfigOverrideMixin: + """Route flat overrides to their declared owner and reject unknown fields. + + Builder-backed model configs store Megatron-Core transformer settings in a + nested ``transformer`` dataclass while exposing flat assignment for recipe + compatibility. This mixin keeps that convenience without allowing typos to + create phantom configuration attributes. Subclasses declare the nested + config type through ``transformer_config_class``. + """ + + transformer_config_class: ClassVar[type[TransformerConfig]] + + def __setattr__(self, name: str, value: Any, /) -> None: + """Assign a declared outer or nested field and reject phantom fields.""" + try: + transformer = object.__getattribute__(self, "transformer") + except AttributeError: + object.__setattr__(self, name, value) + return + + model_fields = getattr(type(self), "__dataclass_fields__", {}) + transformer_fields = getattr(type(transformer), "__dataclass_fields__", {}) + descriptor = getattr(type(self), name, None) + + if name == "transformer" or name in model_fields: + object.__setattr__(self, name, value) + elif name in transformer_fields: + setattr(transformer, name, value) + elif name == "builder" or name.startswith("_") or hasattr(descriptor, "__set__"): + object.__setattr__(self, name, value) + else: + raise AttributeError( + f"Neither {type(self).__name__} nor {type(transformer).__name__} declares a field named {name!r}." + ) class ModelConfig(_MegatronModelConfig): @@ -41,37 +121,6 @@ def get_builder_cls(self) -> type: return builder_cls @classmethod - def from_dict(cls, data: dict[str, Any]) -> "ModelConfig": + def from_dict(cls, data: dict[str, Any]) -> _MegatronModelConfig: """Deserialize config from dictionary with Bridge target validation.""" - - def _from_dict(subdata: dict[str, Any], full_key: str) -> Any: - target = subdata.get("_target_") - if target is None: - raise ValueError("Cannot deserialize: missing '_target_' field") - if not isinstance(target, str): - raise ValueError(f"Cannot deserialize: '_target_' must be a string, got {type(target).__name__}") - - config_cls = _resolve_target(target, full_key=full_key) - if not isinstance(config_cls, type) or not is_dataclass(config_cls): - raise ValueError(f"Cannot deserialize: target '{target}' did not resolve to a dataclass type") - - valid_fields = {f.name for f in dataclass_fields(config_cls) if f.init} - filtered_data = {k: v for k, v in subdata.items() if k in valid_fields and not k.startswith("_")} - - subconfigs = {} - for k, v in filtered_data.items(): - if isinstance(v, dict) and "_target_" in v: - subconfigs[k] = _from_dict(v, full_key=f"{full_key}.{k}") - filtered_data.update(subconfigs) - - return config_cls(**filtered_data) - - builder = data.get("_builder_") - if not isinstance(builder, str): - raise ValueError("Cannot deserialize: missing '_builder_' field") - _validate_target_prefix(target=builder, full_key="_builder_") - - result = _from_dict(data, full_key="_target_") - result.builder = builder - - return result + return deserialize_model_config(data) diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index cfe790c4d4..2954b4a13c 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -1488,8 +1488,18 @@ def import_ckpt( model_context = nullcontext() if dist.is_initialized() else temporary_distributed_context(backend="gloo") with model_context: - # Convert to Megatron model - megatron_model = bridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True) + # Prefer the native ModelConfig/ModelBuilder path for migrated model + # families while preserving the provider path for legacy bridges. + if bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION: + model_config = bridge.get_model_config() + model_config.transformer.use_cpu_initialization = True + megatron_model = bridge.get_model( + model_config, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + else: + megatron_model = bridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True) # Save as Megatron checkpoint hf_tokenizer_kwargs = {} @@ -1871,6 +1881,8 @@ def get_model( pg_collection = self._get_or_initialize_pg_collection(transformer_config) kwargs.setdefault("data_parallel_random_init", False) models = builder.build_distributed_models(pg_collection=pg_collection, **kwargs) + for model in models: + model.model_config = model_config succeeded = True finally: transformer_config.perform_initialization = original_perform_initialization diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index 12886afce0..ade4bc0d7e 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -38,6 +38,7 @@ Type, TypeVar, Union, + cast, ) import torch @@ -52,6 +53,7 @@ from transformers.configuration_utils import PretrainedConfig from transformers.modeling_utils import PreTrainedModel +from megatron.bridge.models.common import ModelConfigOverrideMixin from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.param_mapping import ( MegatronParamMapping, @@ -434,11 +436,17 @@ def mapping_registry(self) -> MegatronMappingRegistry: # provider, such as one implementing MLA. PROVIDER_CLASS = None - # Builder-backed construction is rolling out incrementally by model family. Override these - # when the standard GPT or Transformer config cannot represent the model, and set - # MODEL_CONFIG_CLASS to None for families that have not yet migrated. + # Builder-backed construction is rolling out incrementally by model family. Override this + # when the standard GPT model config cannot represent the model, and set it to None for + # families that have not yet migrated. The selected model config owns its nested + # transformer config class. MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = BridgeGPTModelConfig - TRANSFORMER_CONFIG_CLASS: ClassVar[type[BridgeTransformerConfig]] = BridgeTransformerConfig + + # Conversion still uses the provider-backed construction path unless a model family has + # explicitly migrated its conversion lifecycle to ModelBuilder. Keeping this separate + # from MODEL_CONFIG_CLASS preserves get_model_config() for legacy families without + # silently changing how their checkpoints are constructed. + USE_MODEL_CONFIG_FOR_CONVERSION: ClassVar[bool] = False # Leave unset unless HF export must copy nonstandard files in addition to the usual artifacts, # for example ``["*reasoning_parser.py"]``. @@ -731,12 +739,23 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> ModelConfig: ) config_kwargs = self.hf_config_to_model_config_kwargs(hf_config) + if not issubclass(model_config_class, ModelConfigOverrideMixin): + raise TypeError(f"{model_config_class.__name__} must inherit {ModelConfigOverrideMixin.__name__}.") + model_config_with_overrides = cast(type[ModelConfigOverrideMixin], model_config_class) + transformer_config_class = model_config_with_overrides.transformer_config_class + if not isinstance(transformer_config_class, type) or not issubclass( + transformer_config_class, BridgeTransformerConfig + ): + raise TypeError( + f"{model_config_class.__name__}.transformer_config_class must be a " + f"{BridgeTransformerConfig.__name__} subclass." + ) model_kwargs, transformer_kwargs = self._partition_model_config_kwargs( config_kwargs, model_config_class, - self.TRANSFORMER_CONFIG_CLASS, + transformer_config_class, ) - transformer_config = self.TRANSFORMER_CONFIG_CLASS(**transformer_kwargs) + transformer_config = transformer_config_class(**transformer_kwargs) return model_config_class(transformer=transformer_config, **model_kwargs) # Set by @register_bridge decorator diff --git a/src/megatron/bridge/models/conversion/utils.py b/src/megatron/bridge/models/conversion/utils.py index 18bee479ee..a2f80025f6 100644 --- a/src/megatron/bridge/models/conversion/utils.py +++ b/src/megatron/bridge/models/conversion/utils.py @@ -256,9 +256,10 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0): # Create a copy to avoid modifying the original cleaned_obj = copy.copy(obj) - for attr_name in list(vars(cleaned_obj).keys()): - attr_value = getattr(cleaned_obj, attr_name) - + # Read stored attributes from ``__dict__`` directly. Configuration classes may + # deliberately reject dynamic attribute access for values whose meaning is + # layer-dependent, even though the raw value still needs to be copied for IPC. + for attr_name, attr_value in list(vars(cleaned_obj).items()): # Recursively clean attribute cleaned_value = remove_non_pickleables(attr_value, max_depth, current_depth + 1) diff --git a/src/megatron/bridge/models/gemma/gemma4_bridge.py b/src/megatron/bridge/models/gemma/gemma4_bridge.py index a258a76f24..0295413ec8 100644 --- a/src/megatron/bridge/models/gemma/gemma4_bridge.py +++ b/src/megatron/bridge/models/gemma/gemma4_bridge.py @@ -98,6 +98,31 @@ def _infer_attn_pattern(layer_types: list[str]) -> tuple[int, int] | list[str]: return (len(layer_types), 0) +def _attention_config_value( + hf_config: Any, + layer_type: str, + field_name: str, + default: Any, + *, + legacy_field_name: str | None = None, +) -> Any: + """Read an attention field from a concrete HF layer config when available.""" + layer_types = getattr(hf_config, "layer_types", None) + per_layer_config = getattr(hf_config, "per_layer_config", None) + if isinstance(layer_types, (list, tuple)) and per_layer_config is not None: + try: + layer_index = layer_types.index(layer_type) + except ValueError: + pass + else: + return getattr(per_layer_config[layer_index], field_name, default) + + # Transformers 5.15 rejects direct reads of fields that vary by layer. + # Older configs instead serialize separate flat local/global field names. + fallback_name = legacy_field_name or field_name + return getattr(hf_config, "__dict__", {}).get(fallback_name, default) + + def _layer_types_from_provider(provider: Gemma4ModelProvider | Gemma4DenseProvider) -> list[str]: """Reconstruct the Hugging Face per-layer attention pattern.""" if isinstance(provider, Gemma4DenseProvider): @@ -181,6 +206,11 @@ class Gemma4Bridge(MegatronModelBridge): """ _CONDITIONAL_MOE_FIELDS = frozenset({"num_moe_experts", "moe_router_topk", "moe_ffn_hidden_size"}) + CONFIG_MAPPING = [ + mapping + for mapping in MegatronModelBridge.CONFIG_MAPPING + if mapping[0] not in {"head_dim", "num_key_value_heads"} + ] def _should_map_hf_config_field(self, hf_config: Any, hf_name: str, megatron_name: str, value: Any) -> bool: if megatron_name in self._CONDITIONAL_MOE_FIELDS: @@ -213,11 +243,18 @@ def _build_dense_provider(self, hf_config) -> Gemma4DenseProvider: sliding_rope = rope_params.get("sliding_attention", {}) full_rope = rope_params.get("full_attention", {}) num_attention_heads = hf_config.num_attention_heads - num_query_groups = hf_config.num_key_value_heads - num_global_query_groups = getattr( + num_query_groups = _attention_config_value( hf_config, - "num_global_key_value_heads", + "sliding_attention", + "num_key_value_heads", + getattr(hf_config, "__dict__", {}).get("num_key_value_heads", 4), + ) + num_global_query_groups = _attention_config_value( + hf_config, + "full_attention", + "num_key_value_heads", num_query_groups, + legacy_field_name="num_global_key_value_heads", ) self._dense_num_attention_heads = num_attention_heads @@ -236,8 +273,14 @@ def _build_dense_provider(self, hf_config) -> Gemma4DenseProvider: ffn_hidden_size=hf_config.intermediate_size, num_attention_heads=num_attention_heads, num_query_groups=num_query_groups, - kv_channels=getattr(hf_config, "head_dim", 256), - global_kv_channels=getattr(hf_config, "global_head_dim", 512), + kv_channels=_attention_config_value(hf_config, "sliding_attention", "head_dim", 256), + global_kv_channels=_attention_config_value( + hf_config, + "full_attention", + "head_dim", + 512, + legacy_field_name="global_head_dim", + ), num_global_query_groups=num_global_query_groups, seq_length=hf_config.max_position_embeddings, vocab_size=hf_config.vocab_size, @@ -260,6 +303,12 @@ def _build_dense_provider(self, hf_config) -> Gemma4DenseProvider: def _build_moe_provider(self, hf_config) -> Gemma4ModelProvider: """Build a Gemma4ModelProvider from HF config (MoE path).""" provider_kwargs = self.hf_config_to_provider_kwargs(hf_config) + provider_kwargs["num_query_groups"] = _attention_config_value( + hf_config, + "sliding_attention", + "num_key_value_heads", + getattr(hf_config, "__dict__", {}).get("num_key_value_heads", 4), + ) provider = Gemma4ModelProvider(**provider_kwargs) provider.window_size = getattr(hf_config, "sliding_window", 1024) @@ -268,13 +317,25 @@ def _build_moe_provider(self, hf_config) -> Gemma4ModelProvider: rope_theta_from_hf(hf_config), ) - head_dim = getattr(hf_config, "head_dim", 256) + head_dim = _attention_config_value(hf_config, "sliding_attention", "head_dim", 256) provider.softmax_scale = 1.0 provider.kv_channels = head_dim provider.qk_layernorm = True - provider.global_head_dim = getattr(hf_config, "global_head_dim", 512) - provider.num_global_key_value_heads = getattr(hf_config, "num_global_key_value_heads", 2) + provider.global_head_dim = _attention_config_value( + hf_config, + "full_attention", + "head_dim", + 512, + legacy_field_name="global_head_dim", + ) + provider.num_global_key_value_heads = _attention_config_value( + hf_config, + "full_attention", + "num_key_value_heads", + 2, + legacy_field_name="num_global_key_value_heads", + ) provider.attention_k_eq_v = getattr(hf_config, "attention_k_eq_v", False) rope_params = getattr(hf_config, "rope_parameters", {}) @@ -320,12 +381,14 @@ def megatron_to_hf_config(cls, provider: Gemma4ModelProvider | Gemma4DenseProvid "enable_moe_block": is_moe, "final_logit_softcapping": provider.final_logit_softcapping, "global_head_dim": (provider.global_head_dim if is_moe else provider.global_kv_channels), + "head_dim": provider.kv_channels, "hidden_size_per_layer_input": getattr(provider, "per_layer_embed_dim", 0), "layer_types": _layer_types_from_provider(provider), "num_kv_shared_layers": getattr(provider, "num_kv_shared_layers", 0), "num_global_key_value_heads": ( provider.num_global_key_value_heads if is_moe else provider.num_global_query_groups ), + "num_key_value_heads": provider.num_query_groups, "rope_parameters": _rope_parameters_from_provider(provider), "sliding_window": mcore_to_hf_window_size(window_size), "use_double_wide_mlp": getattr(provider, "use_double_wide_mlp", False), @@ -406,16 +469,23 @@ def maybe_modify_loaded_hf_weight( text_config, "num_attention_heads", getattr(self, "_dense_num_attention_heads", 8) ) kv_head_dim = q_weight.shape[0] // num_q_heads - num_kv_heads = getattr(text_config, "num_key_value_heads", getattr(self, "_dense_num_query_groups", 2)) + num_kv_heads = _attention_config_value( + text_config, + "sliding_attention", + "num_key_value_heads", + getattr(self, "_dense_num_query_groups", 2), + ) layer_match = re.search(r"layers\.(\d+)\.", q_name) layer_types = getattr(text_config, "layer_types", None) if layer_match and layer_types: layer_idx = int(layer_match.group(1)) if layer_idx < len(layer_types) and layer_types[layer_idx] == "full_attention": - num_global_kv_heads = getattr( + num_global_kv_heads = _attention_config_value( text_config, - "num_global_key_value_heads", + "full_attention", + "num_key_value_heads", getattr(self, "_dense_num_global_query_groups", None), + legacy_field_name="num_global_key_value_heads", ) if num_global_kv_heads is not None: num_kv_heads = num_global_kv_heads diff --git a/src/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py b/src/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py index 4e92924c5c..031d34b632 100644 --- a/src/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py +++ b/src/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py @@ -49,6 +49,7 @@ ) from megatron.bridge.models.gemma.gemma4_bridge import ( Gemma4Bridge, + _attention_config_value, _Gemma4QKVMapping, _infer_attn_pattern, ) @@ -99,6 +100,12 @@ def provider_bridge( return self._build_moe_provider(text_config) provider_kwargs = self.hf_config_to_provider_kwargs(text_config) + provider_kwargs["num_query_groups"] = _attention_config_value( + text_config, + "sliding_attention", + "num_key_value_heads", + 4, + ) provider = Gemma4VLModelProvider(**provider_kwargs) provider.window_size = getattr(text_config, "sliding_window", 1024) @@ -107,13 +114,25 @@ def provider_bridge( rope_theta_from_hf(text_config), ) - head_dim = getattr(text_config, "head_dim", 256) + head_dim = _attention_config_value(text_config, "sliding_attention", "head_dim", 256) provider.softmax_scale = 1.0 provider.kv_channels = head_dim provider.qk_layernorm = True - provider.global_head_dim = getattr(text_config, "global_head_dim", 512) - provider.num_global_key_value_heads = getattr(text_config, "num_global_key_value_heads", 2) + provider.global_head_dim = _attention_config_value( + text_config, + "full_attention", + "head_dim", + 512, + legacy_field_name="global_head_dim", + ) + provider.num_global_key_value_heads = _attention_config_value( + text_config, + "full_attention", + "num_key_value_heads", + 2, + legacy_field_name="num_global_key_value_heads", + ) provider.attention_k_eq_v = getattr(text_config, "attention_k_eq_v", False) rope_params = getattr(text_config, "rope_parameters", {}) diff --git a/src/megatron/bridge/models/gpt/model_config.py b/src/megatron/bridge/models/gpt/model_config.py index 7cdb9e856d..1eb291f599 100644 --- a/src/megatron/bridge/models/gpt/model_config.py +++ b/src/megatron/bridge/models/gpt/model_config.py @@ -15,16 +15,17 @@ """Serializable Bridge extension of Megatron-LM's GPT model config.""" from dataclasses import dataclass -from typing import Any +from typing import Any, ClassVar -from megatron.bridge.models.common.base import ModelConfig +from megatron.bridge.models.common.base import ModelConfig, ModelConfigOverrideMixin from megatron.bridge.models.gpt.gpt_builder import GPTModelConfig +from megatron.bridge.models.transformer_config import TransformerConfig from megatron.bridge.utils.activation_map import callable_to_str, str_to_callable from megatron.bridge.utils.instantiate_utils import _resolve_target @dataclass(kw_only=True) -class BridgeGPTModelConfig(GPTModelConfig): +class BridgeGPTModelConfig(ModelConfigOverrideMixin, GPTModelConfig, ModelConfig): """GPT model config with strict overrides and callable serialization. Outer GPT build fields and nested transformer fields keep one declared @@ -37,29 +38,7 @@ class BridgeGPTModelConfig(GPTModelConfig): ``run_config.yaml`` without adding model-specific config fields. """ - def __setattr__(self, name: str, value: Any, /) -> None: - """Assign a declared outer or nested field and reject phantom fields.""" - try: - transformer = object.__getattribute__(self, "transformer") - except AttributeError: - object.__setattr__(self, name, value) - return - - model_fields = getattr(type(self), "__dataclass_fields__", {}) - transformer_fields = getattr(type(transformer), "__dataclass_fields__", {}) - - if name == "transformer": - object.__setattr__(self, name, value) - elif name in model_fields: - object.__setattr__(self, name, value) - elif name in transformer_fields: - setattr(transformer, name, value) - elif name == "builder" or name.startswith("_"): - object.__setattr__(self, name, value) - else: - raise AttributeError( - f"Neither {type(self).__name__} nor {type(transformer).__name__} declares a field named {name!r}." - ) + transformer_config_class: ClassVar[type[TransformerConfig]] = TransformerConfig def get_builder_cls(self) -> type: """Resolve the configured builder through Bridge's target allowlist.""" diff --git a/src/megatron/bridge/models/muse_glimmer/__init__.py b/src/megatron/bridge/models/muse_glimmer/__init__.py new file mode 100644 index 0000000000..8335b8838e --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Muse Glimmer model, builder, configuration, and bridge exports.""" + +from transformers import ( + MuseGlimmerConfig, + MuseGlimmerTextConfig, + MuseGlimmerVisionConfig, +) + +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import MuseGlimmerModel +from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import MuseGlimmerBridge +from megatron.bridge.models.muse_glimmer.muse_glimmer_builder import MuseGlimmerModelBuilder +from megatron.bridge.models.muse_glimmer.muse_glimmer_config import ( + MuseGlimmerModelConfig, + MuseGlimmerTransformerConfig, + MuseGlimmerVisionModelConfig, +) + + +__all__ = [ + "MuseGlimmerBridge", + "MuseGlimmerConfig", + "MuseGlimmerModel", + "MuseGlimmerModelBuilder", + "MuseGlimmerModelConfig", + "MuseGlimmerTextConfig", + "MuseGlimmerTransformerConfig", + "MuseGlimmerVisionConfig", + "MuseGlimmerVisionModelConfig", +] diff --git a/src/megatron/bridge/models/muse_glimmer/data/__init__.py b/src/megatron/bridge/models/muse_glimmer/data/__init__.py new file mode 100644 index 0000000000..66c104e04b --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/data/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Muse Glimmer data processing.""" + +from megatron.bridge.models.muse_glimmer.data.collate_fn import muse_glimmer_collate_fn + + +__all__ = ["muse_glimmer_collate_fn"] diff --git a/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py new file mode 100644 index 0000000000..cdc3ad5a89 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py @@ -0,0 +1,176 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Muse Glimmer multimodal SFT collator.""" + +from typing import Any + +import torch + +from megatron.bridge.data.collators.sequence import prepare_sequence_batch +from megatron.bridge.data.collators.sequence_padding import use_processor_right_padding +from megatron.bridge.data.collators.visual import THW_GRID_VISUAL_KEYS +from megatron.bridge.data.conversation_processing import ( + assistant_mask_boundary_config_from_markers, + build_assistant_loss_mask, + chat_template_kwargs_from_example, + shared_chat_template_kwargs_from_examples, +) +from megatron.bridge.data.datasets.utils import IGNORE_INDEX +from megatron.bridge.data.packing.in_batch import build_mcore_thd_sequence_batch_from_rows +from megatron.bridge.data.token_utils import extract_skipped_token_ids +from megatron.bridge.training.utils.visual_inputs import GenericVisualInputs + + +MUSE_GLIMMER_ASSISTANT_START = "<|start|>assistant to=user<|message|>" +MUSE_GLIMMER_TURN_END = "<|eot|>" + + +def muse_glimmer_collate_fn( + examples: list[dict[str, Any]], + processor: Any, + *, + sequence_length: int | None = None, + pad_to_max_length: bool = False, + pad_to_multiple_of: int = 128, + enable_in_batch_packing: bool = False, + in_batch_packing_pad_to_multiple_of: int = 1, +) -> dict[str, Any]: + """Collate Muse Glimmer conversations and their THW-grid media tensors.""" + skipped_tokens = extract_skipped_token_ids(processor) + boundary_config = assistant_mask_boundary_config_from_markers( + processor, + assistant_start=MUSE_GLIMMER_ASSISTANT_START, + assistant_end=MUSE_GLIMMER_TURN_END, + ) + + if enable_in_batch_packing: + sequence_rows = [] + visual_values: dict[str, list[torch.Tensor]] = {key: [] for key in THW_GRID_VISUAL_KEYS} + with use_processor_right_padding(processor): + for example in examples: + sample_batch = dict( + processor.apply_chat_template( + [example["conversation"]], + tokenize=True, + padding=False, + truncation=True, + return_tensors="pt", + return_dict=True, + **chat_template_kwargs_from_example(example), + ) + ) + input_ids = sample_batch["input_ids"][0] + attention_mask = sample_batch.get("attention_mask") + attention_mask = attention_mask[0] if attention_mask is not None else torch.ones_like(input_ids) + position_ids = sample_batch.get("position_ids") + position_ids = ( + position_ids[0] + if position_ids is not None + else torch.arange(input_ids.numel(), device=input_ids.device, dtype=torch.long) + ) + loss_mask = build_assistant_loss_mask( + example, + input_ids, + processor, + skipped_tokens, + boundary_config=boundary_config, + ).to(device=input_ids.device, dtype=torch.float32) + labels = torch.cat([input_ids[1:], input_ids.new_full((1,), IGNORE_INDEX)]) + if skipped_tokens.numel() > 0: + labels = labels.masked_fill( + torch.isin(labels, skipped_tokens.to(device=labels.device)), IGNORE_INDEX + ) + shifted_loss_mask = torch.cat([loss_mask[1:], loss_mask.new_zeros(1)]) + sequence_rows.append( + { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "labels": labels.masked_fill(shifted_loss_mask == 0, IGNORE_INDEX), + "loss_mask": shifted_loss_mask, + } + ) + for key in THW_GRID_VISUAL_KEYS: + value = sample_batch.get(key) + if isinstance(value, torch.Tensor): + visual_values[key].append(value) + + packed_batch = build_mcore_thd_sequence_batch_from_rows( + sequence_rows, + sequence_length=sequence_length, + pad_token_id=int(getattr(processor.tokenizer, "pad_token_id", 0) or 0), + ignore_index=IGNORE_INDEX, + pad_to_multiple_of=in_batch_packing_pad_to_multiple_of, + ) + visual_kwargs = {key: torch.cat(values, dim=0) for key, values in visual_values.items() if values} + packed_batch["visual_inputs"] = GenericVisualInputs(**visual_kwargs) if visual_kwargs else None + return packed_batch + + with use_processor_right_padding(processor): + batch = dict( + processor.apply_chat_template( + [example["conversation"] for example in examples], + tokenize=True, + padding=True, + truncation=True, + return_tensors="pt", + return_dict=True, + **shared_chat_template_kwargs_from_examples(examples), + ) + ) + + input_ids = batch["input_ids"] + if "position_ids" not in batch: + batch_size, seq_len = input_ids.shape + batch["position_ids"] = ( + torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1).clone().contiguous() + ) + + loss_mask = torch.stack( + [ + build_assistant_loss_mask( + example, + row_input_ids, + processor, + skipped_tokens, + boundary_config=boundary_config, + ) + for example, row_input_ids in zip(examples, input_ids, strict=True) + ] + ).to(device=input_ids.device, dtype=torch.float32) + labels = torch.cat( + [input_ids[:, 1:].contiguous(), input_ids.new_full((input_ids.size(0), 1), IGNORE_INDEX)], dim=1 + ) + if skipped_tokens.numel() > 0: + labels = labels.masked_fill(torch.isin(labels, skipped_tokens.to(device=labels.device)), IGNORE_INDEX) + loss_mask = torch.cat([loss_mask[:, 1:], loss_mask.new_zeros((loss_mask.size(0), 1))], dim=1) + batch["labels"] = labels.masked_fill(loss_mask == 0, IGNORE_INDEX) + batch["loss_mask"] = loss_mask + + visual_kwargs = {key: batch.pop(key) for key in THW_GRID_VISUAL_KEYS if key in batch} + batch["visual_inputs"] = GenericVisualInputs(**visual_kwargs) if visual_kwargs else None + prepare_sequence_batch( + batch, + sequence_length=sequence_length, + pad_to_max_length=pad_to_max_length, + pad_to_multiple_of=pad_to_multiple_of, + enable_in_batch_packing=False, + in_batch_packing_pad_to_multiple_of=1, + ignore_index=IGNORE_INDEX, + ) + return batch + + +__all__ = ["muse_glimmer_collate_fn"] diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py new file mode 100644 index 0000000000..bdac50530b --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -0,0 +1,755 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Megatron modeling components for Muse Glimmer.""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from megatron.core.extensions.transformer_engine import TENorm +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.models.backends import get_backend +from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.mappings import scatter_to_sequence_parallel_region +from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules +from megatron.core.transformer.utils import ensure_metadata_has_dp_cp_group +from megatron.core.utils import get_pg_rank +from torch import Tensor +from torch.utils.checkpoint import checkpoint as torch_checkpoint +from transformer_engine.pytorch import LayerNorm as TELayerNorm + +from megatron.bridge.models.gemma.modules import extend_instance +from megatron.bridge.models.muse_glimmer.muse_glimmer_config import ( + MuseGlimmerModelConfig, + MuseGlimmerVisionModelConfig, +) +from megatron.bridge.utils.common_utils import ( + hook_hf_module_setattr_for_tp_grad_sync, + slice_batch_for_context_parallel, +) + + +if TYPE_CHECKING: + from megatron.core.inference.contexts import BaseInferenceContext + from megatron.core.packed_seq_params import PackedSeqParams + + +class MuseGlimmerSelfAttention(SelfAttention): + """MCore self-attention with Muse's residual-branch post norm.""" + + def __init__(self, config: TransformerConfig, *args: Any, **kwargs: Any) -> None: + super().__init__(config, *args, **kwargs) + # TE RMSNorm always registers an affine weight, while Muse Q/K norms are + # parameter-free; use native RMSNorm to preserve the HF checkpoint contract. + self.q_layernorm = nn.RMSNorm( + self.hidden_size_per_attention_head, + eps=config.layernorm_epsilon, + elementwise_affine=False, + ) + self.k_layernorm = nn.RMSNorm( + self.hidden_size_per_attention_head, + eps=config.layernorm_epsilon, + elementwise_affine=False, + ) + self.post_layernorm = TENorm( + config, + config.hidden_size, + eps=float(getattr(config, "post_norm_epsilon", 1e-8)), + ) + + def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: + output, bias = super().forward(*args, **kwargs) + return self.post_layernorm(output), bias + + +class MuseGlimmerMLP(MLP): + """MCore gated MLP with Muse's residual-branch post norm.""" + + def __init__( + self, + config: TransformerConfig, + submodules: MLPSubmodules, + ffn_hidden_size: int | None = None, + **kwargs: Any, + ) -> None: + ffn_hidden_size = config.ffn_hidden_size if ffn_hidden_size is None else ffn_hidden_size + super().__init__( + config=config, + submodules=submodules, + ffn_hidden_size=ffn_hidden_size, + **kwargs, + ) + self.post_layernorm = TENorm( + config, + config.hidden_size, + eps=float(getattr(config, "post_norm_epsilon", 1e-8)), + ) + + def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: + output, bias = super().forward(*args, **kwargs) + return self.post_layernorm(output), bias + + +def get_muse_glimmer_hybrid_stack_spec(config: TransformerConfig) -> ModuleSpec: + """Build the native Hybrid stack spec used by the Muse decoder.""" + backend = get_backend(config.transformer_impl) + return ModuleSpec( + module=HybridStack, + submodules=HybridStackSubmodules( + attention_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=ModuleSpec( + module=MuseGlimmerSelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=SelfAttentionSubmodules( + linear_qkv=backend.column_parallel_linear(), + core_attention=backend.core_attention(), + linear_proj=backend.row_parallel_linear(), + ), + ), + self_attn_bda=get_bias_dropout_add, + pre_mlp_layernorm=TENorm, + mlp=functools.partial( + MuseGlimmerMLP.as_mlp_submodule, + submodules=MLPSubmodules( + linear_fc1=backend.column_parallel_linear(), + linear_fc2=backend.row_parallel_linear(), + ), + ), + mlp_bda=get_bias_dropout_add, + ), + ), + ), + ) + + +class MuseGlimmerOutputLayerMixin(nn.Module): + """Apply the Muse logit multiplier followed by tanh soft-capping.""" + + def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: + output, bias = super().forward(*args, **kwargs) + output = output * self.config.output_multiplier + softcap = self.config.final_logit_softcapping + if softcap: + output = softcap * torch.tanh(output / softcap) + return output, bias + + +def _vision_cu_seqlens(grid_thw: Tensor) -> Tensor: + lengths = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + return F.pad(lengths.cumsum(dim=0, dtype=torch.int32), (1, 0), value=0) + + +def _vision_position_ids(grid_thw: Tensor) -> Tensor: + positions = [] + for temporal, height, width in grid_thw.tolist(): + rows, columns = torch.meshgrid( + torch.arange(int(height), device=grid_thw.device), + torch.arange(int(width), device=grid_thw.device), + indexing="ij", + ) + positions.append(torch.stack((columns.flatten(), rows.flatten()), dim=-1).repeat(int(temporal), 1) + 1) + return torch.cat(positions, dim=0) + + +def _vision_window_index( + grid_thw: Tensor, + *, + window_size: int, +) -> tuple[Tensor, Tensor]: + indices = [] + cumulative_lengths = [0] + offset = 0 + for temporal, height, width in grid_thw.tolist(): + temporal, height, width = int(temporal), int(height), int(width) + index = torch.arange(temporal * height * width).reshape(temporal, height, width) + pad_height = (-height) % window_size + pad_width = (-width) % window_size + padded = F.pad(index, (0, pad_width, 0, pad_height), value=-1) + num_windows_height = (height + pad_height) // window_size + num_windows_width = (width + pad_width) // window_size + padded = padded.reshape( + temporal, + num_windows_height, + window_size, + num_windows_width, + window_size, + ) + padded = padded.permute(0, 1, 3, 2, 4).reshape(-1, window_size, window_size) + lengths = (padded != -1).sum(dim=(1, 2)) + flattened = padded.flatten() + indices.append(flattened[flattened != -1] + offset) + cumulative_lengths.extend((lengths.cumsum(0) + cumulative_lengths[-1]).tolist()) + offset += temporal * height * width + return ( + torch.cat(indices).to(device=grid_thw.device), + torch.unique_consecutive(torch.tensor(cumulative_lengths, device=grid_thw.device, dtype=torch.int32)), + ) + + +def _bilinear_indices_and_weights(grid_thw: Tensor, side: int) -> tuple[Tensor, Tensor]: + index_parts: list[list[Tensor]] = [[] for _ in range(4)] + weight_parts: list[list[Tensor]] = [[] for _ in range(4)] + for temporal, height, width in grid_thw.tolist(): + temporal, height, width = int(temporal), int(height), int(width) + height_grid = (torch.arange(height, device=grid_thw.device).float() + 0.5) * (side / height) - 0.5 + width_grid = (torch.arange(width, device=grid_thw.device).float() + 0.5) * (side / width) - 0.5 + height_floor = torch.floor(height_grid).long() + width_floor = torch.floor(width_grid).long() + height_ceil = height_floor + 1 + width_ceil = width_floor + 1 + height_fraction = height_grid - height_floor.float() + width_fraction = width_grid - width_floor.float() + height_floor_valid = (height_floor >= 0) & (height_floor < side) + height_ceil_valid = (height_ceil >= 0) & (height_ceil < side) + width_floor_valid = (width_floor >= 0) & (width_floor < side) + width_ceil_valid = (width_ceil >= 0) & (width_ceil < side) + height_floor = height_floor.clamp(0, side - 1) + height_ceil = height_ceil.clamp(0, side - 1) + width_floor = width_floor.clamp(0, side - 1) + width_ceil = width_ceil.clamp(0, side - 1) + corner_indices = [ + (height_floor[:, None] * side + width_floor[None, :]).flatten(), + (height_floor[:, None] * side + width_ceil[None, :]).flatten(), + (height_ceil[:, None] * side + width_floor[None, :]).flatten(), + (height_ceil[:, None] * side + width_ceil[None, :]).flatten(), + ] + corner_weights = [ + ( + (1 - height_fraction)[:, None] + * (1 - width_fraction)[None, :] + * (height_floor_valid[:, None] & width_floor_valid[None, :]) + ).flatten(), + ( + (1 - height_fraction)[:, None] + * width_fraction[None, :] + * (height_floor_valid[:, None] & width_ceil_valid[None, :]) + ).flatten(), + ( + height_fraction[:, None] + * (1 - width_fraction)[None, :] + * (height_ceil_valid[:, None] & width_floor_valid[None, :]) + ).flatten(), + ( + height_fraction[:, None] + * width_fraction[None, :] + * (height_ceil_valid[:, None] & width_ceil_valid[None, :]) + ).flatten(), + ] + for corner in range(4): + index_parts[corner].append(corner_indices[corner].repeat(temporal)) + weight_parts[corner].append(corner_weights[corner].repeat(temporal)) + return ( + torch.stack([torch.cat(part) for part in index_parts]), + torch.stack([torch.cat(part) for part in weight_parts]), + ) + + +class MuseGlimmerVisionPatchEmbedder(nn.Module): + """Linear patch embedding plus exact half-pixel bilinear position lookup.""" + + def __init__(self, config: MuseGlimmerVisionModelConfig) -> None: + super().__init__() + if config.pos_emb_height != config.pos_emb_width: + raise ValueError("Muse Glimmer currently requires a square learned vision position table.") + patch_dimension = config.patch_temporal * 3 * config.patch_size**2 + self.patch_embedding = nn.Linear(patch_dimension, config.hidden_size, bias=False) + self.position_embedding_table = nn.Embedding( + config.pos_emb_height * config.pos_emb_width, + config.hidden_size, + ) + self.num_grid_per_side = config.pos_emb_height + + def forward(self, pixel_values: Tensor, grid_thw: Tensor) -> Tensor: + embeddings = self.patch_embedding(pixel_values.to(self.patch_embedding.weight.dtype)).reshape( + pixel_values.shape[0], -1 + ) + indices, weights = _bilinear_indices_and_weights(grid_thw, self.num_grid_per_side) + position_embeddings = (self.position_embedding_table(indices) * weights[:, :, None]).sum(dim=0) + return embeddings + position_embeddings.to(embeddings.dtype) + + +def _rotate_half(hidden_states: Tensor) -> Tensor: + first, second = hidden_states.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +class MuseGlimmerVisionRotaryEmbedding(nn.Module): + """Muse 2-D RoPE with [W, H, W, H] frequency interleaving.""" + + def __init__(self, config: MuseGlimmerVisionModelConfig) -> None: + super().__init__() + head_dim = config.hidden_size // config.num_attention_heads + spatial_dim = head_dim // 2 + inverse_frequency = 1.0 / ( + config.rotary_base ** (torch.arange(0, spatial_dim, 2, dtype=torch.float32) / spatial_dim) + ) + self.register_buffer("inv_freq", inverse_frequency, persistent=False) + + def forward(self, hidden_states: Tensor, position_ids: Tensor) -> tuple[Tensor, Tensor]: + del hidden_states + width_frequency = torch.outer(position_ids[:, 0].float(), self.inv_freq.float()) + height_frequency = torch.outer(position_ids[:, 1].float(), self.inv_freq.float()) + frequency = torch.cat( + [width_frequency, height_frequency, width_frequency, height_frequency], + dim=-1, + ) + return frequency.cos(), frequency.sin() + + +class MuseGlimmerVisionAttention(nn.Module): + """Bidirectional packed vision attention matching the HF parameter layout.""" + + def __init__(self, config: MuseGlimmerVisionModelConfig) -> None: + super().__init__() + self.num_heads = config.num_attention_heads + self.head_dim = config.hidden_size // config.num_attention_heads + self.scaling = self.head_dim**-0.5 + self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True) + self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True) + self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True) + self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True) + + def forward( + self, + hidden_states: Tensor, + cu_seqlens: Tensor, + position_embeddings: tuple[Tensor, Tensor], + ) -> Tensor: + sequence_length = hidden_states.shape[0] + query = self.q_proj(hidden_states).reshape(sequence_length, self.num_heads, self.head_dim) + key = self.k_proj(hidden_states).reshape(sequence_length, self.num_heads, self.head_dim) + value = self.v_proj(hidden_states).reshape(sequence_length, self.num_heads, self.head_dim) + cosine, sine = position_embeddings + query_dtype, key_dtype = query.dtype, key.dtype + cosine = cosine[:, None, :].float() + sine = sine[:, None, :].float() + query = ((query.float() * cosine) + (_rotate_half(query.float()) * sine)).to(query_dtype) + key = ((key.float() * cosine) + (_rotate_half(key.float()) * sine)).to(key_dtype) + + outputs = [] + boundaries = cu_seqlens.tolist() + for start, end in zip(boundaries[:-1], boundaries[1:]): + q = query[start:end].transpose(0, 1).unsqueeze(0) + k = key[start:end].transpose(0, 1).unsqueeze(0) + v = value[start:end].transpose(0, 1).unsqueeze(0) + attention_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling + attention_weights = F.softmax(attention_weights, dim=-1, dtype=torch.float32).to(q.dtype) + output = torch.matmul(attention_weights, v) + outputs.append(output.squeeze(0).transpose(0, 1)) + return self.proj(torch.cat(outputs, dim=0).reshape(sequence_length, -1)) + + +class MuseGlimmerVisionMLP(nn.Module): + """Muse vision feed-forward network.""" + + def __init__(self, config: MuseGlimmerVisionModelConfig) -> None: + super().__init__() + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=True) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=True) + + def forward(self, hidden_states: Tensor) -> Tensor: + return self.fc2(F.gelu(self.fc1(hidden_states))) + + +class MuseGlimmerVisionEncoderLayer(nn.Module): + """Pre-normalized Muse vision transformer layer.""" + + def __init__(self, config: MuseGlimmerVisionModelConfig, transformer_config: TransformerConfig) -> None: + super().__init__() + norm_device = ( + "cpu" + if transformer_config.use_cpu_initialization + else "meta" + if transformer_config.init_model_with_meta_device + else torch.cuda.current_device() + ) + norm_kwargs = { + "normalized_shape": config.hidden_size, + "eps": config.layer_norm_epsilon, + "sequence_parallel": False, + "zero_centered_gamma": False, + "device": norm_device, + "dtype": transformer_config.params_dtype, + } + self.norm1 = TELayerNorm(**norm_kwargs) + self.attn = MuseGlimmerVisionAttention(config) + self.norm2 = TELayerNorm(**norm_kwargs) + self.mlp = MuseGlimmerVisionMLP(config) + + def forward( + self, + hidden_states: Tensor, + cu_seqlens: Tensor, + position_embeddings: tuple[Tensor, Tensor], + ) -> Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), cu_seqlens, position_embeddings) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class MuseGlimmerVisionModel(nn.Module): + """Muse vision tower with 3:1 window/full attention and 2x2 pixel shuffle.""" + + def __init__( + self, + config: MuseGlimmerVisionModelConfig, + transformer_config: TransformerConfig, + *, + recompute_layers: bool = False, + ) -> None: + super().__init__() + self.config = config + self.recompute_layers = recompute_layers + self.patch_embedder = MuseGlimmerVisionPatchEmbedder(config) + self.rotary_emb = MuseGlimmerVisionRotaryEmbedding(config) + norm_device = ( + "cpu" + if transformer_config.use_cpu_initialization + else "meta" + if transformer_config.init_model_with_meta_device + else torch.cuda.current_device() + ) + norm_kwargs = { + "normalized_shape": config.hidden_size, + "eps": config.layer_norm_epsilon, + "sequence_parallel": False, + "zero_centered_gamma": False, + "device": norm_device, + "dtype": transformer_config.params_dtype, + } + self.ln_pre = TELayerNorm(**norm_kwargs) + self.layers = nn.ModuleList( + [MuseGlimmerVisionEncoderLayer(config, transformer_config) for _ in range(config.num_hidden_layers)] + ) + self.ln_post = TELayerNorm(**norm_kwargs) + + def _pixel_shuffle(self, hidden_states: Tensor, grid_thw: Tensor) -> Tensor: + outputs = [] + offset = 0 + merge_size = self.config.merge_size + hidden_size = hidden_states.shape[-1] + for temporal, height, width in grid_thw.tolist(): + temporal, height, width = int(temporal), int(height), int(width) + if height % merge_size or width % merge_size: + raise ValueError( + f"Vision grid {(temporal, height, width)} must be divisible by merge_size={merge_size}." + ) + token_count = temporal * height * width + chunk = hidden_states[offset : offset + token_count] + permutation = torch.arange(height * width, device=hidden_states.device) + permutation = ( + permutation.reshape(height // merge_size, merge_size, width // merge_size, merge_size) + .permute(0, 2, 1, 3) + .reshape(-1) + ) + frame_offsets = torch.arange(temporal, device=hidden_states.device)[:, None] * height * width + permutation = (permutation[None, :] + frame_offsets).reshape(-1) + merged_tokens = temporal * (height // merge_size) * (width // merge_size) + chunk = chunk[permutation].reshape(merged_tokens, merge_size * merge_size, hidden_size) + outputs.append(chunk.permute(0, 2, 1).contiguous().reshape(merged_tokens, -1)) + offset += token_count + return torch.cat(outputs, dim=0) + + def forward(self, pixel_values: Tensor, grid_thw: Tensor) -> Tensor: + full_cu_seqlens = _vision_cu_seqlens(grid_thw) + window_index, window_cu_seqlens = _vision_window_index( + grid_thw, + window_size=self.config.pos_emb_height, + ) + hidden_states = self.ln_pre(self.patch_embedder(pixel_values, grid_thw)) + hidden_states = hidden_states[window_index] + position_ids = _vision_position_ids(grid_thw)[window_index] + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer_type, layer in zip(self.config.layer_types, self.layers): + cu_seqlens = full_cu_seqlens if layer_type == "full_attention" else window_cu_seqlens + if self.recompute_layers and self.training: + hidden_states = torch_checkpoint( + layer, + hidden_states, + cu_seqlens, + position_embeddings, + use_reentrant=False, + ) + else: + hidden_states = layer(hidden_states, cu_seqlens, position_embeddings) + hidden_states = hidden_states[torch.argsort(window_index)] + return self._pixel_shuffle(self.ln_post(hidden_states), grid_thw) + + +class MuseGlimmerVisionAdapter(nn.Module): + """Two-layer GELU adapter used before projection into the text width.""" + + def __init__(self, config: MuseGlimmerModelConfig) -> None: + super().__init__() + self.fc1 = nn.Linear(config.vision_output_size, config.projector_hidden_size, bias=False) + self.fc2 = nn.Linear(config.projector_hidden_size, config.projector_hidden_size, bias=False) + + def forward(self, hidden_states: Tensor) -> Tensor: + return F.gelu(self.fc2(F.gelu(self.fc1(hidden_states)))) + + +class MuseGlimmerModel(HybridModel): + """Native MCore Hybrid model with the Muse vision modules attached.""" + + _EMPTY_NORM_EXTRA_STATE_SUFFIXES = ( + "input_layernorm._extra_state", + "self_attention.post_layernorm._extra_state", + "pre_mlp_layernorm._extra_state", + "mlp.post_layernorm._extra_state", + "decoder.final_norm._extra_state", + "vision_tower.ln_pre._extra_state", + "vision_tower.ln_post._extra_state", + "norm1._extra_state", + "norm2._extra_state", + ) + _CENTERED_NORM_WEIGHT_SUFFIXES = ( + "input_layernorm.weight", + "self_attention.post_layernorm.weight", + "pre_mlp_layernorm.weight", + "mlp.post_layernorm.weight", + ) + + def __init__( + self, + config: MuseGlimmerModelConfig, + hybrid_stack_spec: ModuleSpec, + vocab_size: int, + pg_collection: ProcessGroupCollection, + *, + pre_process: bool, + post_process: bool, + vp_stage: int | None, + ) -> None: + super().__init__( + config=config.transformer, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=vocab_size, + max_sequence_length=config.seq_length, + hybrid_layer_pattern=config.hybrid_layer_pattern, + fp16_lm_cross_entropy=config.fp16_lm_cross_entropy, + parallel_output=config.parallel_output, + share_embeddings_and_output_weights=config.share_embeddings_and_output_weights, + position_embedding_type=config.position_embedding_type, + rotary_percent=config.rotary_percent, + rotary_base=config.rotary_base, + scatter_embedding_sequence_parallel=False, + seq_len_interpolation_factor=config.seq_len_interpolation_factor, + pre_process=pre_process, + post_process=post_process, + pg_collection=pg_collection, + vp_stage=vp_stage, + ) + self.model_config = config + if hasattr(self, "output_layer"): + extend_instance(self.output_layer, MuseGlimmerOutputLayerMixin) + if pre_process: + self.vision_tower = MuseGlimmerVisionModel( + config.vision, + config.transformer, + recompute_layers=config.recompute_vision_layers, + ) + self.vision_adapter = MuseGlimmerVisionAdapter(config) + self.vision_projection = nn.Linear( + config.projector_hidden_size, + config.transformer.hidden_size, + bias=False, + ) + for module in (self.vision_tower, self.vision_adapter, self.vision_projection): + module.to(dtype=config.transformer.params_dtype) + hook_hf_module_setattr_for_tp_grad_sync(module) + self.freeze( + freeze_language_model=config.freeze_language_model, + freeze_vision_model=config.freeze_vision_model, + freeze_vision_projection=config.freeze_vision_projection, + ) + + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: tuple[tuple[int, int, int], ...] = (), + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Return a backend-stable Muse checkpoint schema. + + Transformer Engine affine norms add empty ``_extra_state`` objects that + the explicit local fallback does not expose. Omitting only those empty + objects keeps checkpoints backend-interchangeable without relaxing + strict loading for parameters or stateful Transformer Engine modules. + """ + metadata = ensure_metadata_has_dp_cp_group(metadata) + sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata) + replica_id = (0, get_pg_rank(self.tp_group), get_pg_rank(metadata["dp_cp_group"])) + for key in list(sharded_state_dict): + if key.endswith(self._CENTERED_NORM_WEIGHT_SUFFIXES): + sharded_state_dict[key].replica_id = replica_id + elif key.endswith(self._EMPTY_NORM_EXTRA_STATE_SUFFIXES): + extra_state = sharded_state_dict.pop(key) + extra_state_data = getattr(extra_state, "data", None) + has_extra_state = ( + extra_state_data.numel() > 0 + if isinstance(extra_state_data, Tensor) + else extra_state_data is not None and bool(extra_state_data) + ) + if has_extra_state: + raise ValueError(f"Muse Transformer Engine norm extra state must be empty: {key}.") + return sharded_state_dict + + def freeze( + self, + *, + freeze_language_model: bool, + freeze_vision_model: bool, + freeze_vision_projection: bool, + ) -> None: + modules: list[nn.Module] = [] + if freeze_language_model: + modules.extend( + module + for module in ( + getattr(self, "embedding", None), + self.decoder, + getattr(self, "output_layer", None), + ) + if module is not None + ) + if freeze_vision_model and hasattr(self, "vision_tower"): + modules.append(self.vision_tower) + if freeze_vision_projection and hasattr(self, "vision_adapter"): + modules.extend([self.vision_adapter, self.vision_projection]) + for module in modules: + for parameter in module.parameters(): + parameter.requires_grad = False + + def _vision_features(self, pixel_values: Tensor, grid_thw: Tensor) -> Tensor: + features = self.vision_tower(pixel_values, grid_thw) + features = self.vision_projection(self.vision_adapter(features)) + return F.rms_norm( + features, + (features.shape[-1],), + eps=self.config.layernorm_epsilon, + ) + + @staticmethod + def _scatter_features(inputs_embeds: Tensor, input_ids: Tensor, features: Tensor, token_id: int) -> Tensor: + mask = (input_ids == token_id).unsqueeze(-1).expand_as(inputs_embeds) + if inputs_embeds[mask].numel() != features.numel(): + raise ValueError( + f"Muse media token count mismatch for token {token_id}: " + f"{mask[..., 0].sum().item()} text slots vs {features.shape[0]} vision features." + ) + return inputs_embeds.masked_scatter(mask, features.to(inputs_embeds.device, inputs_embeds.dtype)) + + def forward( + self, + input_ids: Tensor | None = None, + position_ids: Tensor | None = None, + attention_mask: Tensor | None = None, + inputs_embeds: Tensor | None = None, + pixel_values: Tensor | None = None, + image_grid_thw: Tensor | None = None, + pixel_values_videos: Tensor | None = None, + video_grid_thw: Tensor | None = None, + labels: Tensor | None = None, + inference_context: BaseInferenceContext | None = None, + runtime_gather_output: bool | None = None, + packed_seq_params: PackedSeqParams | None = None, + *, + inference_params: BaseInferenceContext | None = None, + loss_mask: Tensor | None = None, + padding_mask: Tensor | None = None, + ) -> Tensor | tuple[Tensor, Tensor]: + if self.pre_process: + if inputs_embeds is None: + if input_ids is None: + raise ValueError("input_ids is required when inputs_embeds is not supplied.") + multimodal_mask = (input_ids == self.model_config.image_token_id) | ( + input_ids == self.model_config.video_token_id + ) + llm_input_ids = input_ids.masked_fill(multimodal_mask, 0) + inputs_embeds = self.embedding(input_ids=llm_input_ids, position_ids=None) + inputs_embeds = F.rms_norm( + inputs_embeds, + (inputs_embeds.shape[-1],), + eps=self.config.layernorm_epsilon, + ) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + if pixel_values is not None: + if image_grid_thw is None or input_ids is None: + raise ValueError("image_grid_thw and input_ids are required with pixel_values.") + inputs_embeds = self._scatter_features( + inputs_embeds, + input_ids, + self._vision_features(pixel_values, image_grid_thw), + self.model_config.image_token_id, + ) + if pixel_values_videos is not None: + if video_grid_thw is None or input_ids is None: + raise ValueError("video_grid_thw and input_ids are required with pixel_values_videos.") + inputs_embeds = self._scatter_features( + inputs_embeds, + input_ids, + self._vision_features(pixel_values_videos, video_grid_thw), + self.model_config.video_token_id, + ) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + + inputs_embeds, labels, loss_mask, position_ids, attention_mask = slice_batch_for_context_parallel( + inputs_embeds=inputs_embeds, + labels=labels, + loss_mask=loss_mask, + position_ids=position_ids, + attention_mask=attention_mask, + packed_seq_params=packed_seq_params, + pg_collection=self.pg_collection, + ) + if self.config.sequence_parallel and inputs_embeds is not None: + inputs_embeds = scatter_to_sequence_parallel_region(inputs_embeds, group=self.pg_collection.tp) + + output = super().forward( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=inputs_embeds, + labels=labels, + inference_context=inference_context, + runtime_gather_output=runtime_gather_output, + packed_seq_params=packed_seq_params, + inference_params=inference_params, + loss_mask=loss_mask, + padding_mask=padding_mask, + ) + if loss_mask is not None: + return output, loss_mask + return output + + +__all__ = [ + "MuseGlimmerModel", + "MuseGlimmerVisionModel", + "get_muse_glimmer_hybrid_stack_spec", +] diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py new file mode 100644 index 0000000000..61097b3595 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -0,0 +1,404 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hugging Face ↔ Megatron conversion bridge for Muse Glimmer.""" + +from __future__ import annotations + +import math +from typing import Any, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PretrainedConfig + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + GatedMLPMapping, + MegatronParamMapping, + ReplicatedMapping, + RMSNorm2ZeroCenteredRMSNormMapping, + merge_qkv_weights, + split_qkv_weights, +) +from megatron.bridge.models.conversion.utils import remove_non_pickleables +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import MuseGlimmerModel +from megatron.bridge.models.muse_glimmer.muse_glimmer_config import ( + MuseGlimmerModelConfig, + MuseGlimmerTransformerConfig, + MuseGlimmerVisionModelConfig, +) + + +class MuseGlimmerQKVGMapping(MegatronParamMapping[dict[str, torch.Tensor]]): + """Fuse Muse's full-head Q/K/V/gate projections into MCore gated QKV.""" + + def __init__(self, megatron_param: str, *, q: str, k: str, v: str, gate: str) -> None: + super().__init__(megatron_param, {"q": q, "k": k, "v": v, "gate": gate}) + self._tp_mapping = AutoMapping(megatron_param, megatron_param) + + @staticmethod + def _combine_query_and_gate( + config: MuseGlimmerTransformerConfig, + query: torch.Tensor, + gate: torch.Tensor, + ) -> torch.Tensor: + if query.ndim != 2 or gate.ndim != 2: + raise ValueError("Muse Glimmer QKV/gate mapping supports weight matrices only.") + if query.shape != gate.shape: + raise ValueError(f"Muse query and gate shapes must match, got {query.shape} and {gate.shape}.") + head_dim = config.kv_channels or config.hidden_size // config.num_attention_heads + expected_rows = config.num_attention_heads * head_dim + if query.shape != (expected_rows, config.hidden_size): + raise ValueError( + "Unexpected Muse query/gate shape: " + f"expected {(expected_rows, config.hidden_size)}, got {tuple(query.shape)}." + ) + query = query.view(config.num_attention_heads, head_dim, config.hidden_size) + gate = gate.view(config.num_attention_heads, head_dim, config.hidden_size) + return torch.cat((query, gate), dim=1).reshape(-1, config.hidden_size) + + def hf_to_megatron( + self, + hf_weights: dict[str, torch.Tensor], + megatron_module: nn.Module, + ) -> torch.Tensor: + if self.tp_rank == 0: + config = self._get_config(megatron_module) + query_with_gate = self._combine_query_and_gate(config, hf_weights["q"], hf_weights["gate"]) + merged = merge_qkv_weights(config, query_with_gate, hf_weights["k"], hf_weights["v"]) + else: + merged = None + return self._tp_mapping.hf_to_megatron(merged, megatron_module) + + def megatron_to_hf( + self, + megatron_weights: torch.Tensor | None, + megatron_module: nn.Module | None, + ) -> dict[str, torch.Tensor]: + if megatron_weights is not None: + megatron_weights = self.maybe_dequantize(megatron_weights) + + if megatron_module is None: + config = self.broadcast_obj_from_pp_rank(None, "muse_qkvg_config") + else: + config = remove_non_pickleables(self._get_config(megatron_module), max_depth=3) + config = self.broadcast_obj_from_pp_rank(config, "muse_qkvg_config") + + packed = self._tp_mapping.megatron_to_hf(megatron_weights, megatron_module) + if not packed: + return {} + + query_with_gate, key, value = split_qkv_weights(config, next(iter(packed.values()))) + head_dim = config.kv_channels or config.hidden_size // config.num_attention_heads + query_with_gate = query_with_gate.view( + config.num_attention_heads, + 2 * head_dim, + config.hidden_size, + ) + query, gate = query_with_gate.split(head_dim, dim=1) + hf_param = cast(dict[str, str], self.hf_param) + return { + hf_param["q"]: query.reshape(-1, config.hidden_size), + hf_param["k"]: key, + hf_param["v"]: value, + hf_param["gate"]: gate.reshape(-1, config.hidden_size), + } + + def resolve(self, captures: tuple[str, ...]) -> "MuseGlimmerQKVGMapping": + megatron_param, hf_param = self._resolve_names(captures) + return type(self)( + megatron_param, + q=hf_param["q"], + k=hf_param["k"], + v=hf_param["v"], + gate=hf_param["gate"], + ) + + +@MegatronModelBridge.register_bridge( + source="MuseGlimmerForConditionalGeneration", + target=MuseGlimmerModel, + model_type="muse_glimmer", +) +class MuseGlimmerBridge(MegatronModelBridge): + """Builder-backed bridge for the complete Muse Glimmer multimodal model.""" + + MODEL_CONFIG_CLASS = MuseGlimmerModelConfig + USE_MODEL_CONFIG_FOR_CONVERSION = True + + @staticmethod + def _validate_architecture(hf_config: PretrainedConfig) -> tuple[Any, Any]: + text_config = getattr(hf_config, "text_config", None) + vision_config = getattr(hf_config, "vision_config", None) + if text_config is None or vision_config is None: + raise ValueError("Muse Glimmer requires both text_config and vision_config.") + + if getattr(text_config, "hidden_activation", "silu") != "silu": + raise ValueError("Muse Glimmer currently supports only the released SiLU text MLP.") + if bool(getattr(text_config, "attention_bias", False)): + raise ValueError("Muse Glimmer currently supports only the released bias-free text attention.") + if getattr(vision_config, "hidden_act", "gelu") != "gelu": + raise ValueError("Muse Glimmer currently supports only the released GELU vision MLP.") + if getattr(hf_config, "projector_hidden_act", "gelu") != "gelu": + raise ValueError("Muse Glimmer currently supports only the released GELU projector.") + + text_layers = list(text_config.layer_types) + rope_theta = list(text_config.layer_rope_theta) + vision_layers = list(vision_config.layer_types) + if len(text_layers) != text_config.num_hidden_layers or len(rope_theta) != text_config.num_hidden_layers: + raise ValueError("Muse text layer_types and layer_rope_theta must match num_hidden_layers.") + if len(vision_layers) != vision_config.num_hidden_layers: + raise ValueError("Muse vision layer_types must match num_hidden_layers.") + if any(layer not in {"sliding_attention", "full_attention"} for layer in text_layers): + raise ValueError(f"Unsupported Muse text layer pattern: {text_layers}.") + if any(layer not in {"window_attention", "full_attention"} for layer in vision_layers): + raise ValueError(f"Unsupported Muse vision layer pattern: {vision_layers}.") + + expected_vision_output = vision_config.hidden_size * int(vision_config.merge_size) ** 2 + if int(hf_config.out_hidden_size) != expected_vision_output: + raise ValueError( + f"out_hidden_size={hf_config.out_hidden_size} does not match pixel-shuffled vision width " + f"{expected_vision_output}." + ) + return text_config, vision_config + + def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerModelConfig: + """Translate the nested HF config into a native HybridModel config.""" + text_config, vision_config = self._validate_architecture(hf_config) + params_dtype = self.dtype_from_hf(hf_config, default=torch.bfloat16) + head_dim = int(text_config.head_dim) + rope_parameters = getattr(text_config, "rope_parameters", {}) + rotary_base = float(rope_parameters.get("rope_theta", 500_000.0)) + layer_types = list(text_config.layer_types) + layer_rope_theta = list(text_config.layer_rope_theta) + sliding_window = getattr(text_config, "sliding_window", None) + if not isinstance(sliding_window, int) or sliding_window <= 0: + raise ValueError("Muse Glimmer requires a positive integer sliding_window.") + + transformer = MuseGlimmerTransformerConfig( + num_layers=int(text_config.num_hidden_layers), + hidden_size=int(text_config.hidden_size), + ffn_hidden_size=int(text_config.intermediate_size), + num_attention_heads=int(text_config.num_attention_heads), + num_query_groups=int(text_config.num_key_value_heads), + kv_channels=head_dim, + normalization="RMSNorm", + layernorm_epsilon=float(text_config.rms_norm_eps), + post_norm_epsilon=float(text_config.post_norm_eps), + layernorm_zero_centered_gamma=True, + gated_linear_unit=True, + activation_func=F.silu, + add_bias_linear=False, + add_qkv_bias=bool(text_config.attention_bias), + hidden_dropout=0.0, + attention_dropout=float(text_config.attention_dropout), + qk_layernorm=True, + attention_output_gate=True, + softmax_scale=float(text_config.qk_scale_factor) / math.sqrt(head_dim), + attention_softmax_in_fp32=True, + window_size=(sliding_window - 1, 0), + window_attn_skip_freq=[layer == "sliding_attention" for layer in layer_types], + no_rope_freq=[float(theta) == 0.0 for theta in layer_rope_theta], + init_method_std=float(text_config.initializer_range), + params_dtype=params_dtype, + fp16=params_dtype == torch.float16, + bf16=params_dtype == torch.bfloat16, + transformer_impl="local", + output_multiplier=float(text_config.output_multiplier), + final_logit_softcapping=float(text_config.final_logit_softcapping), + ) + vision_rope = getattr(vision_config, "rope_parameters", {}) + vision = MuseGlimmerVisionModelConfig( + hidden_size=int(vision_config.hidden_size), + intermediate_size=int(vision_config.intermediate_size), + num_hidden_layers=int(vision_config.num_hidden_layers), + num_attention_heads=int(vision_config.num_attention_heads), + patch_size=int(vision_config.patch_size), + patch_temporal=int(vision_config.patch_temporal), + merge_size=int(vision_config.merge_size), + pos_emb_height=int(vision_config.pos_emb_height), + pos_emb_width=int(vision_config.pos_emb_width), + max_position_embeddings=int(vision_config.max_position_embeddings), + layer_norm_epsilon=float(vision_config.layer_norm_eps), + hidden_activation=str(vision_config.hidden_act), + rotary_base=float(vision_rope.get("rope_theta", 10_000.0)), + layer_types=list(vision_config.layer_types), + ) + return MuseGlimmerModelConfig( + transformer=transformer, + vocab_size=int(text_config.vocab_size), + make_vocab_size_divisible_by=self.make_vocab_size_divisible_by(int(text_config.vocab_size)), + should_pad_vocab=False, + seq_length=int(text_config.max_position_embeddings), + parallel_output=True, + share_embeddings_and_output_weights=bool( + getattr(hf_config, "tie_word_embeddings", text_config.tie_word_embeddings) + ), + hybrid_layer_pattern="*" * int(text_config.num_hidden_layers), + position_embedding_type="rope", + rotary_percent=1.0, + rotary_base=rotary_base, + vision=vision, + image_token_id=int(hf_config.image_token_id), + video_token_id=int(hf_config.video_token_id), + bos_token_id=text_config.bos_token_id, + eos_token_id=text_config.eos_token_id, + pad_token_id=text_config.pad_token_id, + vision_output_size=int(hf_config.out_hidden_size), + projector_hidden_size=int(hf_config.projector_hidden_size), + projector_hidden_activation=str(hf_config.projector_hidden_act), + ) + + @classmethod + def megatron_to_hf_config(cls, model_config: MuseGlimmerModelConfig) -> dict[str, Any]: + """Translate the serializable model config back to HF's nested layout.""" + if not isinstance(model_config, MuseGlimmerModelConfig): + raise TypeError(f"Expected MuseGlimmerModelConfig, got {type(model_config).__name__}.") + transformer = model_config.transformer + vision = model_config.vision + window_pattern = transformer.window_attn_skip_freq + if not isinstance(window_pattern, list): + raise ValueError("Muse Glimmer export requires an explicit per-layer window_attn_skip_freq list.") + no_rope_pattern = transformer.no_rope_freq + if not isinstance(no_rope_pattern, list): + raise ValueError("Muse Glimmer export requires an explicit per-layer no_rope_freq list.") + if transformer.window_size is None or transformer.kv_channels is None: + raise ValueError("Muse Glimmer export requires explicit window_size and kv_channels.") + + dtype = str(transformer.params_dtype).removeprefix("torch.") + layer_types = ["sliding_attention" if enabled else "full_attention" for enabled in window_pattern] + layer_rope_theta = [0 if disabled else float(model_config.rotary_base) for disabled in no_rope_pattern] + text_config = { + "model_type": "muse_glimmer_text", + "vocab_size": model_config.vocab_size, + "hidden_size": transformer.hidden_size, + "intermediate_size": transformer.ffn_hidden_size, + "num_hidden_layers": transformer.num_layers, + "num_attention_heads": transformer.num_attention_heads, + "num_key_value_heads": transformer.num_query_groups, + "head_dim": transformer.kv_channels, + "hidden_activation": cls.megatron_to_hf_activation(transformer.activation_func), + "max_position_embeddings": model_config.seq_length, + "rms_norm_eps": transformer.layernorm_epsilon, + "post_norm_eps": transformer.post_norm_epsilon, + "attention_bias": transformer.add_qkv_bias, + "attention_dropout": transformer.attention_dropout, + "initializer_range": transformer.init_method_std, + "sliding_window": transformer.window_size[0] + 1, + "layer_types": layer_types, + "layer_rope_theta": layer_rope_theta, + "rope_parameters": {"rope_theta": model_config.rotary_base, "rope_type": "default"}, + "qk_scale_factor": transformer.softmax_scale * math.sqrt(transformer.kv_channels), + "output_multiplier": transformer.output_multiplier, + "final_logit_softcapping": transformer.final_logit_softcapping, + "tie_word_embeddings": model_config.share_embeddings_and_output_weights, + "bos_token_id": model_config.bos_token_id, + "eos_token_id": model_config.eos_token_id, + "pad_token_id": model_config.pad_token_id, + } + vision_config = { + "model_type": "muse_glimmer_vision", + "hidden_size": vision.hidden_size, + "intermediate_size": vision.intermediate_size, + "num_hidden_layers": vision.num_hidden_layers, + "num_attention_heads": vision.num_attention_heads, + "patch_size": vision.patch_size, + "patch_temporal": vision.patch_temporal, + "merge_size": vision.merge_size, + "pos_emb_height": vision.pos_emb_height, + "pos_emb_width": vision.pos_emb_width, + "max_position_embeddings": vision.max_position_embeddings, + "layer_norm_eps": vision.layer_norm_epsilon, + "hidden_act": vision.hidden_activation, + "layer_types": list(vision.layer_types), + "rope_parameters": {"rope_theta": vision.rotary_base, "rope_type": "default"}, + } + return { + "architectures": ["MuseGlimmerForConditionalGeneration"], + "model_type": "muse_glimmer", + "dtype": dtype, + "text_config": text_config, + "vision_config": vision_config, + "image_token_id": model_config.image_token_id, + "video_token_id": model_config.video_token_id, + "out_hidden_size": model_config.vision_output_size, + "projector_hidden_size": model_config.projector_hidden_size, + "projector_hidden_act": model_config.projector_hidden_activation, + "tie_word_embeddings": model_config.share_embeddings_and_output_weights, + } + + def mapping_registry(self) -> MegatronMappingRegistry: + """Return complete text, vision, adapter, and projection mappings.""" + text_prefix = "model.language_model" + mappings: list[MegatronParamMapping[Any]] = [ + AutoMapping( + megatron_param="embedding.word_embeddings.weight", + hf_param=f"{text_prefix}.embed_tokens.weight", + ), + AutoMapping( + megatron_param="output_layer.weight", + hf_param="lm_head.weight", + ), + RMSNorm2ZeroCenteredRMSNormMapping( + megatron_param="decoder.final_norm.weight", + hf_param=f"{text_prefix}.norm.weight", + ), + ReplicatedMapping( + megatron_param="decoder.layers.*.input_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.input_layernorm.weight", + ), + ReplicatedMapping( + megatron_param="decoder.layers.*.self_attention.post_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.post_attention_layernorm.weight", + ), + ReplicatedMapping( + megatron_param="decoder.layers.*.pre_mlp_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.pre_feedforward_layernorm.weight", + ), + ReplicatedMapping( + megatron_param="decoder.layers.*.mlp.post_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.post_feedforward_layernorm.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.self_attention.linear_proj.weight", + hf_param=f"{text_prefix}.layers.*.self_attn.o_proj.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.mlp.linear_fc2.weight", + hf_param=f"{text_prefix}.layers.*.mlp.down_proj.weight", + ), + MuseGlimmerQKVGMapping( + "decoder.layers.*.self_attention.linear_qkv.weight", + q=f"{text_prefix}.layers.*.self_attn.q_proj.weight", + k=f"{text_prefix}.layers.*.self_attn.k_proj.weight", + v=f"{text_prefix}.layers.*.self_attn.v_proj.weight", + gate=f"{text_prefix}.layers.*.self_attn.gate_proj.weight", + ), + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.linear_fc1.weight", + gate=f"{text_prefix}.layers.*.mlp.gate_proj.weight", + up=f"{text_prefix}.layers.*.mlp.up_proj.weight", + ), + ReplicatedMapping(megatron_param="vision_tower.**", hf_param="model.vision_tower.**"), + ReplicatedMapping(megatron_param="vision_adapter.**", hf_param="model.vision_adapter.**"), + ReplicatedMapping(megatron_param="vision_projection.**", hf_param="model.vision_projection.**"), + ] + return MegatronMappingRegistry(*mappings) + + +__all__ = ["MuseGlimmerBridge", "MuseGlimmerQKVGMapping"] diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py new file mode 100644 index 0000000000..787d6bea97 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Model builder for Muse Glimmer.""" + +from __future__ import annotations + +from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.training.models.hybrid import HybridModelBuilder +from megatron.training.vocab_utils import calculate_padded_vocab_size + +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( + MuseGlimmerModel, + get_muse_glimmer_hybrid_stack_spec, +) +from megatron.bridge.models.muse_glimmer.muse_glimmer_config import MuseGlimmerModelConfig + + +class MuseGlimmerModelBuilder(HybridModelBuilder): + """Construct the complete Muse Glimmer model on MCore HybridModel.""" + + def __init__(self, model_config: MuseGlimmerModelConfig) -> None: + super().__init__(model_config) + + def build_model( + self, + pg_collection: ProcessGroupCollection, + pre_process: bool | None = None, + post_process: bool | None = None, + vp_stage: int | None = None, + ) -> MuseGlimmerModel: + """Build one Muse Glimmer pipeline stage. + + Args: + pg_collection: Process groups used for distributed construction. + pre_process: Whether this stage owns embeddings and the vision stack. + post_process: Whether this stage owns the language-model output head. + vp_stage: Optional virtual pipeline stage index. + + Returns: + A combined Muse Glimmer model stage. + """ + model_config = self._model_config + if not isinstance(model_config, MuseGlimmerModelConfig): + raise TypeError(f"Expected MuseGlimmerModelConfig, got {type(model_config).__name__}.") + + if model_config.vocab_size is None: + raise ValueError("Muse Glimmer vocab_size must be configured before model construction.") + if model_config.should_pad_vocab: + padded_vocab_size = calculate_padded_vocab_size( + model_config.vocab_size, + model_config.make_vocab_size_divisible_by, + model_config.transformer.tensor_model_parallel_size, + ) + else: + padded_vocab_size = model_config.vocab_size + + pre_process = pre_process if pre_process is not None else is_pp_first_stage(pg_collection.pp) + post_process = post_process if post_process is not None else is_pp_last_stage(pg_collection.pp) + return MuseGlimmerModel( + model_config, + get_muse_glimmer_hybrid_stack_spec(model_config.transformer), + padded_vocab_size, + pg_collection, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + + +__all__ = ["MuseGlimmerModelBuilder"] diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py new file mode 100644 index 0000000000..09006ca760 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py @@ -0,0 +1,149 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Native HybridModel configuration objects for Muse Glimmer.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, ClassVar + +from megatron.training.models.hybrid import HybridModelConfig + +from megatron.bridge.models.common import ModelConfigOverrideMixin, deserialize_model_config +from megatron.bridge.models.transformer_config import TransformerConfig +from megatron.bridge.utils.activation_map import callable_to_str, str_to_callable +from megatron.bridge.utils.instantiate_utils import _resolve_target + + +@dataclass +class MuseGlimmerTransformerConfig(TransformerConfig): + """Muse-specific decoder fields not represented by MCore's base config.""" + + post_norm_epsilon: float = 1e-8 + output_multiplier: float = 0.19611613513818404 + final_logit_softcapping: float = 20.0 + + +@dataclass +class MuseGlimmerVisionModelConfig: + """Serializable configuration for the replicated Muse vision encoder.""" + + hidden_size: int = 1_536 + intermediate_size: int = 8_960 + num_hidden_layers: int = 50 + num_attention_heads: int = 16 + patch_size: int = 14 + patch_temporal: int = 2 + merge_size: int = 2 + pos_emb_height: int = 32 + pos_emb_width: int = 32 + max_position_embeddings: int = 1_024 + layer_norm_epsilon: float = 1e-5 + hidden_activation: str = "gelu" + rotary_base: float = 10_000.0 + layer_types: list[str] = field(default_factory=list) + + @property + def depth(self) -> int: + """Expose the vision depth expected by Bridge's VLM FLOPs estimator.""" + return self.num_hidden_layers + + @property + def spatial_merge_size(self) -> int: + """Expose the spatial merge factor expected by VLM data utilities.""" + return self.merge_size + + +@dataclass(kw_only=True) +class MuseGlimmerModelConfig(ModelConfigOverrideMixin, HybridModelConfig): + """Complete builder configuration for Muse Glimmer.""" + + builder: ClassVar[str] = "megatron.bridge.models.muse_glimmer.MuseGlimmerModelBuilder" + transformer_config_class: ClassVar[type[TransformerConfig]] = MuseGlimmerTransformerConfig + hybrid_attention_layers_include_mlp: ClassVar[bool] = True + + vision: MuseGlimmerVisionModelConfig = field(default_factory=MuseGlimmerVisionModelConfig) + image_token_id: int = 200_092 + video_token_id: int = 200_091 + bos_token_id: int | None = 200_000 + eos_token_id: int | list[int] | None = 200_001 + pad_token_id: int | None = None + vision_output_size: int = 6_144 + projector_hidden_size: int = 4_096 + projector_hidden_activation: str = "gelu" + freeze_language_model: bool = False + freeze_vision_model: bool = False + freeze_vision_projection: bool = False + recompute_vision_layers: bool = False + + @property + def vision_config(self) -> MuseGlimmerVisionModelConfig: + """Expose the vision config to shared VLM training utilities.""" + return self.vision + + @property + def special_token_ids(self) -> dict[str, int]: + """Return media token IDs used by multimodal data pipelines.""" + return {"images": self.image_token_id, "videos": self.video_token_id} + + def get_builder_cls(self) -> type: + """Resolve the Muse builder through Bridge's target allowlist.""" + builder_cls = _resolve_target(self.builder, full_key="_builder_") + if not isinstance(builder_cls, type): + raise TypeError(f"Builder target '{self.builder}' did not resolve to a class.") + return builder_cls + + def as_dict(self) -> dict[str, Any]: + """Serialize the Hybrid config with a symbolic activation function.""" + data = super().as_dict() + transformer_data = data.get("transformer") + if not isinstance(transformer_data, dict): + raise TypeError("Serialized Muse Glimmer config must contain a transformer mapping.") + + activation_func = self.transformer.activation_func + if isinstance(activation_func, str): + str_to_callable(activation_func) + activation_name = activation_func + else: + activation_name = callable_to_str(activation_func) + if activation_name is None: + raise ValueError(f"Cannot serialize unregistered activation callable: {activation_func!r}.") + + transformer_data["activation_func"] = activation_name + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "MuseGlimmerModelConfig": + """Deserialize a Muse Hybrid config and restore its activation callable.""" + restored_data = dict(data) + transformer_data = restored_data.get("transformer") + if isinstance(transformer_data, dict): + restored_transformer = dict(transformer_data) + activation_name = restored_transformer.get("activation_func") + if isinstance(activation_name, str): + restored_transformer["activation_func"] = str_to_callable(activation_name) + restored_data["transformer"] = restored_transformer + + result = deserialize_model_config(restored_data) + if not isinstance(result, cls): + raise TypeError(f"Expected {cls.__name__}, got {type(result).__name__}.") + return result + + +__all__ = [ + "MuseGlimmerModelConfig", + "MuseGlimmerTransformerConfig", + "MuseGlimmerVisionModelConfig", +] diff --git a/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/rope.py b/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/rope.py index 9ba16f1bea..750b82dc27 100644 --- a/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/rope.py +++ b/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/rope.py @@ -17,11 +17,12 @@ import torch -def _get_feat_extract_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor: +def _get_feat_extract_output_lengths(input_lengths: torch.Tensor, n_window: int = 50) -> torch.Tensor: """Compute Qwen3-Omni thinker audio token lengths from feature lengths.""" - input_lengths_leave = input_lengths % 100 + chunk_len = n_window * 2 + input_lengths_leave = input_lengths % chunk_len feat_lengths = (input_lengths_leave - 1) // 2 + 1 - output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 + output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // chunk_len) * 13 return output_lengths diff --git a/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py b/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py index c42d0f7c62..00f787882d 100644 --- a/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py +++ b/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py @@ -28,7 +28,10 @@ Qwen3OmniMoeVisionEncoder as Qwen3OmniMoeVisionEncoderHF, ) -from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.rope import get_rope_index +from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.rope import ( + _get_feat_extract_output_lengths, + get_rope_index, +) from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.transformer_config import ( Qwen3OmniTransformerConfig, ) @@ -338,7 +341,10 @@ def get_audio_features( if expected_audio_token_counts is None: return audio_embeds - _, produced_output_lengths = self.audio_model._get_feat_extract_output_lengths(audio_feature_lengths) + produced_output_lengths = _get_feat_extract_output_lengths( + audio_feature_lengths, + getattr(self.audio_model.config, "n_window", 50), + ) expected_audio_token_counts = expected_audio_token_counts.to(produced_output_lengths.device) trimmed_embeds = [] diff --git a/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py b/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py new file mode 100644 index 0000000000..cd5cea1076 --- /dev/null +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Muse Glimmer performance recipes.""" + +from megatron.bridge.perf_recipes.muse_glimmer.h100 import ( + muse_glimmer_30b_pretrain_32gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config, +) + + +__all__ = [ + "muse_glimmer_30b_pretrain_32gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config", +] diff --git a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/__init__.py b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/__init__.py new file mode 100644 index 0000000000..b3410cde73 --- /dev/null +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""H100 performance recipes for Muse Glimmer.""" + +from megatron.bridge.perf_recipes.muse_glimmer.h100.muse_glimmer import ( + muse_glimmer_30b_pretrain_32gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config, +) + + +__all__ = [ + "muse_glimmer_30b_pretrain_32gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config", +] diff --git a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py new file mode 100644 index 0000000000..a4f5cf6124 --- /dev/null +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""H100 performance recipes for Muse Glimmer.""" + +from megatron.bridge.perf_recipes._common import _benchmark_common, _perf_precision +from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS +from megatron.bridge.recipes.muse_glimmer.h100 import ( + muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config, +) +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ConfigContainer, MockGPTDatasetConfig + + +def _muse_glimmer_30b_pretrain_32gpu_h100_config() -> ConfigContainer: + """Build the common Muse Glimmer dense-decoder H100 benchmark.""" + cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() + cfg.model.tensor_model_parallel_size = 4 + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.context_parallel_size = 2 + cfg.model.cp_comm_type = "all_gather" + cfg.model.hybrid_layer_pattern = "|".join("*" * layers for layers in (9, 15, 15, 13)) + cfg.model.freeze_vision_model = True + cfg.model.freeze_vision_projection = True + cfg.model.recompute_vision_layers = False + cfg.dataset = MockGPTDatasetConfig( + seq_length=4096, + random_seed=1234, + reset_attention_mask=False, + reset_position_ids=False, + eod_mask_loss=False, + num_dataset_builder_threads=1, + split="9999,8,2", + data_sharding=True, + dataloader_type="single", + skip_getting_attention_mask_from_dataset=True, + ) + cfg.train.global_batch_size = 96 + cfg.train.micro_batch_size = 3 + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.comm_overlap = CommOverlapConfig(tp_comm_bootstrap_backend="nccl", tp_comm_overlap=True) + + _benchmark_common(cfg, cross_entropy_impl="native") + return cfg + + +def muse_glimmer_30b_pretrain_32gpu_h100_bf16_config() -> ConfigContainer: + """Muse Glimmer dense decoder pretrain: 32× H100, BF16, TP=4 PP=4 CP=2.""" + cfg = _muse_glimmer_30b_pretrain_32gpu_h100_config() + cfg.mixed_precision = _perf_precision("bf16") + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + # Transformer Engine requires this stream ordering for TP overlap on H100. + "CUDA_DEVICE_MAX_CONNECTIONS": 1, + # CUDA allocator and graph-registration settings for the measured workload. + "NCCL_GRAPH_REGISTER": 0, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + # This recipe does not use NCCL user buffers. + "NCCL_NVLS_ENABLE": 0, + # Keep TE norm launch margins explicit for this measured workload. + "NVTE_BWD_LAYERNORM_SM_MARGIN": 0, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 0, + } + return cfg + + +def muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config() -> ConfigContainer: + """Muse Glimmer dense decoder pretrain: 32× H100, FP8 current scaling.""" + cfg = _muse_glimmer_30b_pretrain_32gpu_h100_config() + cfg.mixed_precision = _perf_precision("fp8_cs") + cfg.train.global_batch_size = 192 + cfg.train.micro_batch_size = 6 + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + # Transformer Engine requires this stream ordering for TP overlap on H100. + "CUDA_DEVICE_MAX_CONNECTIONS": 1, + # CUDA allocator and graph-registration settings for the measured workload. + "NCCL_GRAPH_REGISTER": 0, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + # This recipe does not use NCCL user buffers. + "NCCL_NVLS_ENABLE": 0, + # Keep TE norm launch margins explicit for this measured workload. + "NVTE_BWD_LAYERNORM_SM_MARGIN": 0, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 0, + } + return cfg + + +__all__ = [ + "muse_glimmer_30b_pretrain_32gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config", +] diff --git a/src/megatron/bridge/recipes/__init__.py b/src/megatron/bridge/recipes/__init__.py index 7de54cff18..e1bbaa5ae2 100644 --- a/src/megatron/bridge/recipes/__init__.py +++ b/src/megatron/bridge/recipes/__init__.py @@ -50,6 +50,8 @@ from megatron.bridge.recipes.ministral3.h100 import * from megatron.bridge.recipes.moonlight import * from megatron.bridge.recipes.moonlight.h100 import * +from megatron.bridge.recipes.muse_glimmer import * +from megatron.bridge.recipes.muse_glimmer.h100 import * from megatron.bridge.recipes.nemotron_omni import * from megatron.bridge.recipes.nemotron_omni.h100 import * from megatron.bridge.recipes.nemotron_vl import * diff --git a/src/megatron/bridge/recipes/muse_glimmer/__init__.py b/src/megatron/bridge/recipes/muse_glimmer/__init__.py new file mode 100644 index 0000000000..2b58dad7c0 --- /dev/null +++ b/src/megatron/bridge/recipes/muse_glimmer/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Muse Glimmer recipe aliases.""" + +from megatron.bridge.recipes.muse_glimmer.h100 import * # noqa: F403 +from megatron.bridge.recipes.muse_glimmer.h100 import __all__ diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py new file mode 100644 index 0000000000..ff7c457c45 --- /dev/null +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py @@ -0,0 +1,30 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""H100 recipes for Muse Glimmer.""" + +from megatron.bridge.recipes.muse_glimmer.h100.muse_glimmer import ( + muse_glimmer_30b_peft_8gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config, + muse_glimmer_30b_sft_32gpu_h100_bf16_config, + muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, +) + + +__all__ = [ + "muse_glimmer_30b_peft_8gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config", + "muse_glimmer_30b_sft_32gpu_h100_bf16_config", + "muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config", +] diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py new file mode 100644 index 0000000000..84d712e314 --- /dev/null +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -0,0 +1,335 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bounded H100 training recipes for Muse Glimmer 30B.""" + +from __future__ import annotations + +import torch + +from megatron.bridge import AutoBridge +from megatron.bridge.peft.base import PEFT +from megatron.bridge.recipes.common import _peft_common_vlm, _sft_common_vlm +from megatron.bridge.recipes.utils.dataset_utils import default_peft_config +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed + + +_MODEL_ID = "meta-models/Muse-Glimmer-30B" +_MODEL_REVISION = "f84ecc3a0ea984a4c04542a84269e3d065350a6e" # pragma: allowlist secret +_CORD_V2_REVISION = "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret + + +def muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() -> ConfigContainer: + """Return a 100-step random-init multimodal pretraining config on 32 H100 GPUs.""" + cfg = _sft_common_vlm() + + cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).get_model_config() + cfg.model.seq_length = 4096 + cfg.model.tensor_model_parallel_size = 8 + cfg.model.pipeline_model_parallel_size = 2 + cfg.model.pipeline_dtype = torch.bfloat16 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" + cfg.model.context_parallel_size = 1 + cfg.model.cp_comm_type = "p2p" + cfg.model.sequence_parallel = True + cfg.model.freeze_language_model = False + cfg.model.freeze_vision_model = False + cfg.model.freeze_vision_projection = False + cfg.model.recompute_vision_layers = True + cfg.model.transformer_impl = "transformer_engine" + cfg.model.attention_backend = "auto" + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["core_attn"] + + cfg.dataset.seq_length = 4096 + cfg.dataset.hf_processor_path = _MODEL_ID + cfg.dataset.hf_processor_kwargs = { + "revision": _MODEL_REVISION, + "max_image_tokens": 256, + } + cfg.dataset.source.split = "train" + cfg.dataset.source.load_kwargs = {"revision": _CORD_V2_REVISION} + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + cfg.dataset.pad_to_max_length = True + cfg.dataset.enable_in_batch_packing = False + # Worker-local RNG and prefetch queues are not part of the checkpoint. + # Main-process loading keeps the step-50 resume data stream reproducible. + cfg.dataset.num_workers = 0 + cfg.dataset.dataloader_type = "cyclic" + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + cfg.ddp.check_for_nan_in_grad = True + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + cfg.logger.log_interval = 1 + cfg.logger.log_throughput = True + cfg.rng.seed = 1234 + cfg.train.train_iters = 100 + # TP8/PP2 leaves two data-parallel replicas on 32 GPUs. One sample per + # replica keeps this bounded verification recipe to one microbatch per + # optimizer step while exercising the complete multimodal model. + cfg.train.global_batch_size = 2 + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=40, + max_lr=3e-4, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.optimizer.use_precision_aware_optimizer = True + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def muse_glimmer_30b_sft_32gpu_h100_bf16_config() -> ConfigContainer: + """Return a 100-step full multimodal SFT config on 32 H100 GPUs.""" + cfg = _sft_common_vlm() + + cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).get_model_config() + cfg.model.seq_length = 4096 + cfg.model.tensor_model_parallel_size = 8 + cfg.model.pipeline_model_parallel_size = 2 + cfg.model.pipeline_dtype = torch.bfloat16 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" + cfg.model.context_parallel_size = 1 + cfg.model.cp_comm_type = "p2p" + cfg.model.sequence_parallel = True + cfg.model.freeze_language_model = False + cfg.model.freeze_vision_model = False + cfg.model.freeze_vision_projection = False + cfg.model.recompute_vision_layers = True + cfg.model.transformer_impl = "transformer_engine" + cfg.model.attention_backend = "auto" + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["core_attn"] + + cfg.dataset.seq_length = 4096 + cfg.dataset.hf_processor_path = _MODEL_ID + cfg.dataset.hf_processor_kwargs = { + "revision": _MODEL_REVISION, + "max_image_tokens": 256, + } + cfg.dataset.source.split = "train" + cfg.dataset.source.load_kwargs = {"revision": _CORD_V2_REVISION} + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + cfg.dataset.num_workers = 4 + cfg.dataset.dataloader_type = "cyclic" + cfg.dataset.pad_to_max_length = True + cfg.dataset.enable_in_batch_packing = False + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + cfg.ddp.check_for_nan_in_grad = True + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + cfg.logger.log_interval = 1 + cfg.logger.log_throughput = True + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 8 + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=10, + max_lr=5e-6, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.save_optim = False + cfg.checkpoint.save_rng = False + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContainer: + """Return a 100-step packed 8K full SFT config with CP2 on 32 H100 GPUs.""" + cfg = _sft_common_vlm() + + cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).get_model_config() + cfg.model.seq_length = 8192 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.pipeline_dtype = torch.bfloat16 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.hybrid_layer_pattern = "|".join("*" * layers for layers in (9, 15, 15, 13)) + cfg.model.context_parallel_size = 2 + cfg.model.cp_comm_type = "a2a" + cfg.model.sequence_parallel = True + cfg.model.freeze_language_model = False + cfg.model.freeze_vision_model = False + cfg.model.freeze_vision_projection = False + cfg.model.recompute_vision_layers = True + cfg.model.transformer_impl = "transformer_engine" + cfg.model.attention_backend = "auto" + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["core_attn"] + + cfg.dataset.seq_length = 8192 + cfg.dataset.hf_processor_path = _MODEL_ID + cfg.dataset.hf_processor_kwargs = { + "revision": _MODEL_REVISION, + "max_image_tokens": 256, + } + cfg.dataset.source.split = "train" + cfg.dataset.source.load_kwargs = {"revision": _CORD_V2_REVISION} + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + cfg.dataset.num_workers = 4 + cfg.dataset.dataloader_type = "cyclic" + cfg.dataset.pad_to_max_length = False + cfg.dataset.enable_in_batch_packing = True + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + cfg.ddp.check_for_nan_in_grad = True + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + cfg.logger.log_interval = 1 + cfg.logger.log_throughput = True + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 8 + cfg.train.micro_batch_size = 2 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=10, + max_lr=5e-6, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.save_optim = False + cfg.checkpoint.save_rng = False + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def muse_glimmer_30b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> ConfigContainer: + """Return a 100-step attention-projection LoRA config on 8 H100 GPUs.""" + cfg = _peft_common_vlm() + + cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).get_model_config() + cfg.model.seq_length = 8192 + cfg.model.tensor_model_parallel_size = 8 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_dtype = None + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.hybrid_layer_pattern = "*" * 52 + cfg.model.context_parallel_size = 1 + cfg.model.cp_comm_type = "p2p" + cfg.model.sequence_parallel = True + cfg.model.freeze_language_model = False + cfg.model.freeze_vision_model = False + cfg.model.freeze_vision_projection = False + cfg.model.recompute_vision_layers = True + cfg.model.transformer_impl = "transformer_engine" + cfg.model.attention_backend = "auto" + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["core_attn"] + + cfg.dataset.seq_length = 8192 + cfg.dataset.hf_processor_path = _MODEL_ID + cfg.dataset.hf_processor_kwargs = { + "revision": _MODEL_REVISION, + "max_image_tokens": 256, + } + cfg.dataset.source.split = "train" + cfg.dataset.source.load_kwargs = {"revision": _CORD_V2_REVISION} + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + cfg.dataset.num_workers = 4 + cfg.dataset.dataloader_type = "cyclic" + cfg.dataset.pad_to_max_length = True + cfg.dataset.enable_in_batch_packing = False + + peft_cfg = default_peft_config(peft_scheme) + if isinstance(peft_scheme, str) and peft_scheme.lower() in {"lora", "dora"}: + peft_cfg.target_modules = ["linear_qkv", "linear_proj"] + peft_cfg.dim = 8 + peft_cfg.alpha = 16 + peft_cfg.dropout = 0.0 + cfg.peft = peft_cfg + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + cfg.ddp.check_for_nan_in_grad = True + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + cfg.logger.log_interval = 1 + cfg.logger.log_throughput = True + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 8 + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=10, + max_lr=1e-4, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +__all__ = [ + "muse_glimmer_30b_peft_8gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config", + "muse_glimmer_30b_sft_32gpu_h100_bf16_config", + "muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config", +] diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index cc2165eeed..849b47c743 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -23,8 +23,10 @@ import sys import threading from abc import ABC +from collections.abc import Mapping from dataclasses import dataclass, replace from enum import Enum, auto +from functools import partial from logging import getLogger from pathlib import Path from time import time @@ -162,6 +164,56 @@ DATALOADER_STATE_SUBDIR = "energon" +def _get_run_config_tp_pp(model_config: Mapping[str, Any]) -> tuple[int, int]: + """Read TP/PP sizes from a flat provider or nested HybridModel config.""" + transformer_config = model_config.get("transformer") + parallel_config = transformer_config if isinstance(transformer_config, Mapping) else model_config + return ( + parallel_config["tensor_model_parallel_size"], + parallel_config["pipeline_model_parallel_size"], + ) + + +class _CpuTorchDistSaveShardedStrategy(TorchDistSaveShardedStrategy): + """Run MCore's synchronous torch-dist writer without a CUDA staging barrier.""" + + @staticmethod + def _run_finalize(finalize_fn: Callable[[], None]) -> None: + """Use a backend-compatible device for MCore's failure-status tensor.""" + if torch.distributed.is_initialized() and torch.distributed.get_backend() == "nccl": + finalize_fn() + return + + current_device = torch.cuda.current_device + try: + torch.cuda.current_device = lambda: torch.device("cpu") + finalize_fn() + finally: + torch.cuda.current_device = current_device + + def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> None: + """Save CPU tensors synchronously without calling ``torch.cuda.synchronize``.""" + async_request = self.async_save(sharded_state_dict, checkpoint_dir, async_strategy="mcore") + preload_fn = async_request.preload_fn + if preload_fn is not None: + if not isinstance(preload_fn, partial): + raise TypeError(f"Expected a partial CPU preload callback, got {type(preload_fn).__name__}.") + bound = inspect.signature(preload_fn.func).bind_partial( + *preload_fn.args, + **(preload_fn.keywords or {}), + ) + if "non_blocking" not in bound.signature.parameters: + raise TypeError("MCore checkpoint preload callback does not accept non_blocking.") + bound.arguments["non_blocking"] = False + async_request = async_request._replace( + preload_fn=partial(preload_fn.func, *bound.args, **bound.kwargs), + ) + async_request = async_request._replace( + finalize_fns=[partial(self._run_finalize, finalize_fn) for finalize_fn in async_request.finalize_fns], + ) + async_request.execute_sync() + + # ============================================================================ # Checkpoint version and utilities # ============================================================================ @@ -2900,10 +2952,7 @@ def _load_checkpoint_from_path( tp_pp_match = True mismatch_msg = "" else: - ckpt_tp_pp = ( - run_config["model"]["tensor_model_parallel_size"], - run_config["model"]["pipeline_model_parallel_size"], - ) + ckpt_tp_pp = _get_run_config_tp_pp(run_config["model"]) run_tp_pp = ( cfg.model.tensor_model_parallel_size, cfg.model.pipeline_model_parallel_size, @@ -3023,10 +3072,7 @@ def _load_checkpoint_from_path( run_config_filename = get_checkpoint_run_config_filename(checkpoint_name) if file_exists(run_config_filename): run_config = read_run_config(run_config_filename) - ckpt_tp_pp = ( - run_config["model"]["tensor_model_parallel_size"], - run_config["model"]["pipeline_model_parallel_size"], - ) + ckpt_tp_pp = _get_run_config_tp_pp(run_config["model"]) run_tp_pp = ( cfg.model.tensor_model_parallel_size, cfg.model.pipeline_model_parallel_size, diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index cbba665735..8015a296a5 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/src/megatron/bridge/training/model_load_save.py @@ -30,7 +30,7 @@ from megatron.training.models.base import ModelConfig from megatron.bridge.models.model_provider import ModelParallelKwargs, ModelProviderMixin -from megatron.bridge.training.checkpointing import save_checkpoint +from megatron.bridge.training.checkpointing import _CpuTorchDistSaveShardedStrategy, save_checkpoint from megatron.bridge.training.config import CheckpointConfig, ConfigContainer, LoggerConfig from megatron.bridge.training.state import GlobalState from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer, build_tokenizer @@ -266,7 +266,7 @@ def load_model_config( read_run_config, ) from megatron.bridge.training.mlm_compat.arguments import _load_args_from_checkpoint, _transformer_config_from_args - from megatron.bridge.utils.instantiate_utils import instantiate + from megatron.bridge.utils.instantiate_utils import _resolve_target, instantiate run_config_filename = get_checkpoint_run_config_filename(checkpoint_path) @@ -291,7 +291,15 @@ def load_model_config( if mbridge_ckpt: if "_builder_" in run_config["model"]: - model_cfg = ModelConfig.from_dict(run_config["model"]) + model_dict = run_config["model"] + target = model_dict.get("_target_") + if isinstance(target, str): + model_config_cls = _resolve_target(target, full_key="model._target_") + if not isinstance(model_config_cls, type) or not issubclass(model_config_cls, ModelConfig): + raise TypeError(f"Builder model target must resolve to a ModelConfig class, got {target!r}.") + model_cfg = model_config_cls.from_dict(model_dict) + else: + model_cfg = ModelConfig.from_dict(model_dict) else: model_cfg = instantiate(run_config["model"]) else: @@ -571,16 +579,18 @@ def save_megatron_model( hf_tokenizer_kwargs=hf_tokenizer_kwargs or {}, ) - # Get model config from the first model instance - model_config = get_model_config(model[0]) + # Builder-backed models retain their complete outer ModelConfig for + # checkpoint reconstruction. Legacy provider models expose their provider + # through ``config`` as before. + model_config = getattr(model[0], "model_config", None) + if not isinstance(model_config, ModelConfig): + model_config = get_model_config(model[0]) - # Validate that the model config is a model provider - if not isinstance(model_config, ModelProviderMixin): + if not isinstance(model_config, (ModelConfig, ModelProviderMixin)): raise TypeError( - f"Expected model config to be an instance of ModelProviderMixin, " + "Expected model config to be an instance of ModelConfig or ModelProviderMixin, " f"but got {type(model_config).__name__}. " - f"Model configs must inherit from ModelProviderMixin to ensure proper " - f"model instantiation and configuration handling." + "Model configs must support builder- or provider-backed reconstruction." ) # Create global state for checkpointing @@ -631,6 +641,17 @@ def save_megatron_model( elif tokenizer_error is not None: raise tokenizer_error + runtime_config = getattr(model[0], "config", None) + use_cpu_save_strategy = getattr(runtime_config, "use_cpu_initialization", False) is True + if isinstance(model_config, ModelConfig): + transformer_config = getattr(model_config, "transformer", None) + use_cpu_save_strategy = ( + use_cpu_save_strategy or getattr(transformer_config, "use_cpu_initialization", False) is True + ) + checkpointing_context = None + if ckpt_format == "torch_dist" and use_cpu_save_strategy: + checkpointing_context = {"save_strategy": _CpuTorchDistSaveShardedStrategy("torch_dist", 1)} + if low_memory_save: # Low-memory save flow: process factories incrementally, freeing memory as we go import gc @@ -789,6 +810,7 @@ def _collect_factories(d): num_floating_point_operations_so_far=0, prebuilt_state_dict=state_dict, pg_collection=pg_collection, + checkpointing_context=checkpointing_context, callback_manager=None, ) else: @@ -799,6 +821,7 @@ def _collect_factories(d): optimizer=None, opt_param_scheduler=None, num_floating_point_operations_so_far=0, + checkpointing_context=checkpointing_context, callback_manager=None, ) diff --git a/src/megatron/bridge/training/utils/config_utils.py b/src/megatron/bridge/training/utils/config_utils.py index 0413d39928..87f12df03e 100644 --- a/src/megatron/bridge/training/utils/config_utils.py +++ b/src/megatron/bridge/training/utils/config_utils.py @@ -102,8 +102,12 @@ def _convert_pretrained_config_to_dict( if include_target and is_dataclass(value): field_names = {field.name for field in dataclass_fields(value) if not field.name.startswith("_")} + raw_values = vars(value) config_items = [ - (field.name, getattr(value, field.name)) + ( + field.name, + raw_values[field.name] if field.name in raw_values else object.__getattribute__(value, field.name), + ) for field in dataclass_fields(value) if not field.name.startswith("_") ] diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index f400664389..e9d148f322 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -1581,6 +1581,14 @@ def _compute_vit_flops(): patches_per_image = num_vision_patches / batch_size if batch_size > 0 else num_vision_patches return vit_flops(cfg, batch_size, patches_per_image) + # A Hybrid attention symbol normally describes only the attention work for + # FLOPs accounting. Some native HybridModel stacks, such as Muse Glimmer, + # use that symbol for a complete Transformer layer containing both attention + # and its MLP. The standard Transformer path already accounts for that full + # layer plus GQA, output gates, sliding attention, logits, and vision work. + if getattr(cfg.model, "hybrid_attention_layers_include_mlp", False): + return transformer_flops() + # Main entrypoint for FLOPs calculation. Mirror MCore's hybrid detection: # a physical hybrid pattern is sufficient to select hybrid accounting. if getattr(cfg.model, "is_hybrid_model", False) or getattr(cfg.model, "hybrid_layer_pattern", None): diff --git a/src/megatron/bridge/utils/common_utils.py b/src/megatron/bridge/utils/common_utils.py index f190e42416..a30aaf9056 100644 --- a/src/megatron/bridge/utils/common_utils.py +++ b/src/megatron/bridge/utils/common_utils.py @@ -289,16 +289,22 @@ def slice_batch_for_context_parallel( # MCore's THD path slices within each packed sequence rather than across them. if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": - if inputs_embeds is None: - raise ValueError("inputs_embeds is required for THD CP slicing") - - seq_len = inputs_embeds.size(1) + reference_tensor = inputs_embeds + if reference_tensor is None: + reference_tensor = next( + (tensor for tensor in (labels, loss_mask, position_ids) if tensor is not None), + None, + ) + if reference_tensor is None: + raise ValueError("At least one sequence tensor is required for THD CP slicing") + + seq_len = reference_tensor.size(1) index = get_packed_seq_cp_partition_indices( packed_seq_params, total_tokens=seq_len, cp_size=cp_size, cp_rank=cp_rank, - device=inputs_embeds.device, + device=reference_tensor.device, cp_group=pg_collection.cp, ) diff --git a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_conversion.py b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_conversion.py index f829304376..aa310b35e1 100644 --- a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_conversion.py +++ b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_conversion.py @@ -42,7 +42,9 @@ "num_attention_heads": 32, "num_experts_per_tok": 4, "num_hidden_layers": 2, - "num_key_value_heads": 4, + # DeepSeek MLA expands compressed KV latents to every attention head before + # dispatching to the HF attention backend, so it uses MHA rather than GQA. + "num_key_value_heads": 32, "num_nextn_predict_layers": 0, "q_lora_rank": 512, "topk_group": 2, diff --git a/tests/functional_tests/test_groups/models/nemotron_omni/test_nemotron_omni_conversion.py b/tests/functional_tests/test_groups/models/nemotron_omni/test_nemotron_omni_conversion.py index 8edc5dfd2c..c9f7ec1ba3 100644 --- a/tests/functional_tests/test_groups/models/nemotron_omni/test_nemotron_omni_conversion.py +++ b/tests/functional_tests/test_groups/models/nemotron_omni/test_nemotron_omni_conversion.py @@ -40,6 +40,10 @@ # 52-layer pattern, which keeps a representative mamba (M), MoE-MLP (E) and # attention (*) layer mix. _LLM_LAYER_TYPES = ["mamba", "moe", "mamba", "moe", "mamba", "attention"] +_LLM_LAYER_TYPE_ALIASES = { + "linear_attention": "mamba", + "full_attention": "attention", +} _LLM_OVERRIDES = { "hybrid_override_pattern": "MEMEM*", "layer_types": _LLM_LAYER_TYPES, @@ -63,6 +67,11 @@ def _apply_overrides(sub_config, overrides: dict) -> None: setattr(sub_config, key, value) +def _normalize_llm_layer_types(layer_types: list[str]) -> list[str]: + """Normalize equivalent remote-config names to the toy model's layer vocabulary.""" + return [_LLM_LAYER_TYPE_ALIASES.get(layer_type, layer_type) for layer_type in layer_types] + + def _fix_tied_weights_keys(model: nn.Module) -> None: """Convert _tied_weights_keys from list to dict for transformers 5.x compatibility.""" for module in model.modules(): @@ -214,11 +223,23 @@ def test_nemotron_omni_conversion_roundtrip(self, nemotron_omni_toy_model_path, with open(config_file) as f: saved_config = json.load(f) + source_config_file = Path(hf_model_id) / "config.json" + if source_config_file.exists(): + with open(source_config_file) as f: + source_config = json.load(f) + else: + source_config = AutoConfig.from_pretrained(hf_model_id, trust_remote_code=True).to_dict() + source_llm_config = source_config["llm_config"] + assert saved_config["architectures"][0] == "NemotronH_Nano_Omni_Reasoning_V3" assert saved_config["model_type"] == "NemotronH_Nano_Omni_Reasoning_V3" assert "llm_config" in saved_config assert "vision_config" in saved_config assert "sound_config" in saved_config assert saved_config["llm_config"]["num_hidden_layers"] == 6 - assert saved_config["llm_config"]["layer_types"] == _LLM_LAYER_TYPES - assert saved_config["llm_config"]["layers_block_type"] == _LLM_LAYER_TYPES + expected_layer_types = _normalize_llm_layer_types(source_llm_config["layer_types"]) + expected_block_types = _normalize_llm_layer_types(source_llm_config["layers_block_type"]) + assert expected_layer_types == _LLM_LAYER_TYPES + assert expected_block_types == _LLM_LAYER_TYPES + assert _normalize_llm_layer_types(saved_config["llm_config"]["layer_types"]) == expected_layer_types + assert _normalize_llm_layer_types(saved_config["llm_config"]["layers_block_type"]) == expected_block_types diff --git a/tests/unit_tests/conversion/launcher/test_gpu_backend.py b/tests/unit_tests/conversion/launcher/test_gpu_backend.py index 9fa262264a..ea790ed6c1 100644 --- a/tests/unit_tests/conversion/launcher/test_gpu_backend.py +++ b/tests/unit_tests/conversion/launcher/test_gpu_backend.py @@ -114,6 +114,9 @@ def provide_distributed_model(self, *args, **kwargs): class _FakeModelBridge: + MODEL_CONFIG_CLASS = None + USE_MODEL_CONFIG_FOR_CONVERSION = False + def get_hf_tokenizer_kwargs(self): return {"padding_side": "left"} @@ -242,6 +245,60 @@ def fake_from_hf_pretrained(*args, **kwargs): assert initialize_call[2] == {"seed": 0, "create_gloo_process_groups": False} assert prepared_outputs == [(("/ckpt",), {"overwrite": False, "source_paths": ["hf"]})] + def test_import_uses_builder_for_migrated_model(self, cli, monkeypatch): + calls = [] + transformer = types.SimpleNamespace() + model_config = types.SimpleNamespace(transformer=transformer, pipeline_model_parallel_layout=None) + + class BuilderModelBridge(_FakeModelBridge): + MODEL_CONFIG_CLASS = object + USE_MODEL_CONFIG_FOR_CONVERSION = True + + class FakeBridge: + _model_bridge = BuilderModelBridge() + hf_pretrained = _FakeHfPretrained() + hf_model_revision = None + + def get_model_config(self): + calls.append(("get_model_config", (), {})) + return model_config + + def get_model(self, *args, **kwargs): + calls.append(("get_model", args, kwargs)) + return ["hybrid-model"] + + def to_megatron_provider(self, *args, **kwargs): + raise AssertionError("migrated models must not use the provider path") + + def save_megatron_model(self, *args, **kwargs): + calls.append(("save_megatron_model", args, kwargs)) + + monkeypatch.setattr(cli, "_ensure_distributed_initialized", lambda timeout_minutes: None) + monkeypatch.setattr(cli, "_prepare_distributed_output", lambda *args, **kwargs: None) + monkeypatch.setattr(cli.AutoBridge, "from_hf_pretrained", lambda *args, **kwargs: FakeBridge()) + + cli.import_checkpoint.__wrapped__( + hf_model="hf", + hf_revision=None, + megatron_path="/ckpt", + tp=8, + pp=1, + ep=1, + etp=1, + torch_dtype="bfloat16", + trust_remote_code=False, + low_memory_save=False, + distributed_timeout_minutes=None, + overwrite=False, + ) + + assert transformer.tensor_model_parallel_size == 8 + get_model_call = next(call for call in calls if call[0] == "get_model") + assert get_model_call[1] == (model_config,) + assert get_model_call[2] == {"wrap_with_ddp": False, "mixed_precision_wrapper": None} + save_call = next(call for call in calls if call[0] == "save_megatron_model") + assert save_call[1] == (["hybrid-model"], "/ckpt") + class TestExportMegatronToHf: def test_export_uses_checkpoint_config_and_does_not_move_loaded_model_to_cuda(self, cli, monkeypatch, tmp_path): diff --git a/tests/unit_tests/models/gemma/test_gemma4_bridge.py b/tests/unit_tests/models/gemma/test_gemma4_bridge.py index 21d7d58b71..8ce379d2d0 100644 --- a/tests/unit_tests/models/gemma/test_gemma4_bridge.py +++ b/tests/unit_tests/models/gemma/test_gemma4_bridge.py @@ -219,6 +219,10 @@ def test_runtime_preserves_nonperiodic_attention_schedule(self, bridge): pretrained.config = hf_config provider = bridge.provider_bridge(pretrained) + assert provider.kv_channels == 256 + assert provider.global_head_dim == 512 + assert provider.num_query_groups == hf_config.per_layer_config[1].num_key_value_heads + assert provider.num_global_key_value_heads == hf_config.per_layer_config[0].num_key_value_heads runtime_layer_types = [ "sliding_attention" if _is_local_attn_layer(layer_number, provider.interleaved_attn_pattern) diff --git a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py index 0cf273827d..98ee198ac6 100644 --- a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py +++ b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py @@ -80,7 +80,7 @@ def test_hf_config_ignores_upstream_num_experts_default() -> None: provider_kwargs = GLM5Bridge().hf_config_to_provider_kwargs(hf_config) - assert hf_config.num_experts == 256 + assert getattr(hf_config, "num_experts", None) in (None, 256) assert provider_kwargs["num_moe_experts"] == hf_config.n_routed_experts == 8 diff --git a/tests/unit_tests/models/muse_glimmer/__init__.py b/tests/unit_tests/models/muse_glimmer/__init__.py new file mode 100644 index 0000000000..3e4afde2e2 --- /dev/null +++ b/tests/unit_tests/models/muse_glimmer/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. diff --git a/tests/unit_tests/models/muse_glimmer/test_collate.py b/tests/unit_tests/models/muse_glimmer/test_collate.py new file mode 100644 index 0000000000..7de7391f64 --- /dev/null +++ b/tests/unit_tests/models/muse_glimmer/test_collate.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import torch + +from megatron.bridge.data.collators.registry import resolve_model_collate +from megatron.bridge.models.muse_glimmer.data.collate_fn import muse_glimmer_collate_fn + + +class _Tokenizer: + pad_token_id = 0 + padding_side = "left" + all_special_ids = [] + + def __call__(self, text, **kwargs): + del kwargs + return {"input_ids": [ord(char) for char in text]} + + +class _MuseGlimmerProcessor: + def __init__(self) -> None: + self.tokenizer = _Tokenizer() + + def apply_chat_template(self, conversations, **kwargs): + assert kwargs["tokenize"] is True + assert kwargs["padding"] in {True, False} + assert self.tokenizer.padding_side == "right" + batch_size = len(conversations) + length = 4 if kwargs["padding"] else 3 + len(conversations[0]) + return { + "input_ids": torch.arange(11, 11 + length).unsqueeze(0).expand(batch_size, -1).clone(), + "attention_mask": torch.ones(batch_size, length, dtype=torch.long), + "pixel_values": torch.arange(24, dtype=torch.float32).reshape(2, 12), + "image_grid_thw": torch.tensor([[1, 2, 4]]), + } + + +def test_muse_glimmer_processor_resolves_family_collator(): + assert resolve_model_collate("MuseGlimmerProcessor") is muse_glimmer_collate_fn + + +def test_muse_glimmer_collator_builds_shifted_labels_and_visual_inputs(monkeypatch): + monkeypatch.setattr( + "megatron.bridge.models.muse_glimmer.data.collate_fn.build_assistant_loss_mask", + lambda *args, **kwargs: torch.tensor([0.0, 0.0, 1.0, 1.0]), + ) + processor = _MuseGlimmerProcessor() + batch = muse_glimmer_collate_fn( + [{"conversation": [{"role": "user", "content": "image"}, {"role": "assistant", "content": "ok"}]}], + processor, + pad_to_multiple_of=1, + ) + + assert processor.tokenizer.padding_side == "left" + assert torch.equal(batch["labels"], torch.tensor([[-100, 13, 14, -100]])) + assert torch.equal(batch["loss_mask"], torch.tensor([[0.0, 1.0, 1.0, 0.0]])) + assert torch.equal(batch["visual_inputs"].pixel_values, torch.arange(24, dtype=torch.float32).reshape(2, 12)) + assert torch.equal(batch["visual_inputs"].image_grid_thw, torch.tensor([[1, 2, 4]])) + assert "pixel_values" not in batch + + +def test_muse_glimmer_collator_packs_sequences_and_concatenates_visual_inputs(monkeypatch): + monkeypatch.setattr( + "megatron.bridge.models.muse_glimmer.data.collate_fn.build_assistant_loss_mask", + lambda _example, input_ids, *_args, **_kwargs: torch.cat([torch.zeros(1), torch.ones(input_ids.numel() - 1)]), + ) + processor = _MuseGlimmerProcessor() + examples = [ + {"conversation": [{"role": "user", "content": "image"}]}, + { + "conversation": [ + {"role": "user", "content": "image"}, + {"role": "assistant", "content": "ok"}, + ] + }, + ] + + batch = muse_glimmer_collate_fn( + examples, + processor, + sequence_length=16, + enable_in_batch_packing=True, + in_batch_packing_pad_to_multiple_of=4, + ) + + assert processor.tokenizer.padding_side == "left" + assert batch["input_ids"].shape == (1, 12) + assert torch.equal(batch["cu_seqlens_q"], torch.tensor([0, 4, 9], dtype=torch.int32)) + assert torch.equal(batch["cu_seqlens_q_padded"], torch.tensor([0, 4, 12], dtype=torch.int32)) + assert batch["visual_inputs"].pixel_values.shape == (4, 12) + assert torch.equal(batch["visual_inputs"].image_grid_thw, torch.tensor([[1, 2, 4], [1, 2, 4]])) diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py new file mode 100644 index 0000000000..a9e79505dd --- /dev/null +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -0,0 +1,623 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for the builder-backed Muse Glimmer implementation.""" + +from __future__ import annotations + +import math +import os +from collections.abc import Iterator +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TEDotProductAttention, TENorm +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.tensor_parallel.layers import ColumnParallelLinear +from megatron.core.transformer.dot_product_attention import DotProductAttention +from megatron.training.models.hybrid import HybridModelConfig +from transformer_engine.pytorch import LayerNorm as TELayerNorm +from transformer_engine.pytorch import RMSNorm as TERMSNorm + +from megatron.bridge import AutoBridge +from megatron.bridge.models.conversion.param_mapping import ( + RMSNorm2ZeroCenteredRMSNormMapping, + merge_qkv_weights, + split_qkv_weights, +) +from megatron.bridge.models.muse_glimmer import ( + MuseGlimmerConfig, + MuseGlimmerModel, + MuseGlimmerModelBuilder, + MuseGlimmerModelConfig, + MuseGlimmerTextConfig, + MuseGlimmerTransformerConfig, + MuseGlimmerVisionConfig, +) +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( + MuseGlimmerVisionAttention, + MuseGlimmerVisionModel, + get_muse_glimmer_hybrid_stack_spec, +) +from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import ( + MuseGlimmerBridge, + MuseGlimmerQKVGMapping, +) +from megatron.bridge.training.utils.flop_utils import num_floating_point_operations, vit_flops_from_grid_thw + + +pytestmark = pytest.mark.unit + +_NVTE_ATTENTION_ENV_VARS = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN") + + +@pytest.fixture(scope="module", autouse=True) +def restore_nvte_attention_environment() -> Iterator[None]: + """Keep MCore attention-backend selection local to this test module.""" + original_values = {name: os.environ.get(name) for name in _NVTE_ATTENTION_ENV_VARS} + yield + for name, value in original_values.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _tiny_hf_config() -> MuseGlimmerConfig: + text_config = MuseGlimmerTextConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=256, + sliding_window=32, + layer_types=["sliding_attention", "sliding_attention", "sliding_attention", "full_attention"], + layer_rope_theta=[500_000.0, 500_000.0, 500_000.0, 0.0], + rope_parameters={"rope_theta": 500_000.0, "rope_type": "default"}, + bos_token_id=100, + eos_token_id=101, + pad_token_id=0, + ) + vision_config = MuseGlimmerVisionConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=4, + patch_size=2, + patch_temporal=2, + merge_size=2, + pos_emb_height=4, + pos_emb_width=4, + max_position_embeddings=16, + layer_types=["window_attention", "window_attention", "window_attention", "full_attention"], + rope_parameters={"rope_theta": 10_000.0, "rope_type": "default"}, + ) + return MuseGlimmerConfig( + text_config=text_config, + vision_config=vision_config, + image_token_id=120, + video_token_id=121, + out_hidden_size=64, + projector_hidden_size=32, + projector_hidden_act="gelu", + architectures=["MuseGlimmerForConditionalGeneration"], + tie_word_embeddings=False, + torch_dtype=torch.bfloat16, + ) + + +def test_config_conversion_uses_model_config_path_only() -> None: + bridge = MuseGlimmerBridge() + hf_config = _tiny_hf_config() + + with ( + patch.object(bridge, "provider_bridge", side_effect=AssertionError("legacy provider path used")), + patch.object( + bridge, + "hf_config_to_provider_kwargs", + side_effect=AssertionError("legacy provider kwargs path used"), + ), + ): + model_config = bridge.hf_config_to_model_config(hf_config) + + assert isinstance(model_config, MuseGlimmerModelConfig) + assert isinstance(model_config, HybridModelConfig) + assert model_config.get_builder_cls() is MuseGlimmerModelBuilder + assert model_config.transformer_impl == "local" + assert model_config.num_layers == 4 + assert model_config.hidden_size == 64 + assert model_config.ffn_hidden_size == 128 + assert model_config.num_query_groups == 2 + assert model_config.kv_channels == 8 + assert model_config.softmax_scale == pytest.approx(3.87 / math.sqrt(8)) + assert model_config.window_size == (31, 0) + assert model_config.window_attn_skip_freq == [True, True, True, False] + assert model_config.no_rope_freq == [False, False, False, True] + assert model_config.attention_output_gate is True + assert model_config.qk_layernorm is True + assert model_config.hybrid_layer_pattern == "****" + assert model_config.special_token_ids == {"images": 120, "videos": 121} + assert model_config.vision_config is model_config.vision + assert model_config.vision_config.depth == 4 + assert model_config.vision_config.spatial_merge_size == 2 + + +def test_model_config_owns_transformer_type_and_routes_flat_overrides() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + + assert MuseGlimmerModelConfig.transformer_config_class is MuseGlimmerTransformerConfig + assert "TRANSFORMER_CONFIG_CLASS" not in MuseGlimmerBridge.__dict__ + assert MuseGlimmerBridge.USE_MODEL_CONFIG_FOR_CONVERSION is True + + model_config.tensor_model_parallel_size = 8 + model_config.seq_length = 512 + + assert model_config.transformer.tensor_model_parallel_size == 8 + assert "tensor_model_parallel_size" not in model_config.__dict__ + assert model_config.seq_length == 512 + assert "seq_length" in model_config.__dict__ + + with pytest.raises(AttributeError, match="declares a field"): + model_config.tesnor_model_parallel_size = 4 + + +def test_config_conversion_uses_top_level_embedding_sharing_contract() -> None: + """The VLM config, not its nested text config, owns HF weight tying.""" + hf_config = _tiny_hf_config() + hf_config.tie_word_embeddings = True + hf_config.text_config.tie_word_embeddings = False + + model_config = MuseGlimmerBridge().hf_config_to_model_config(hf_config) + + assert model_config.share_embeddings_and_output_weights is True + + +def test_config_conversion_falls_back_to_nested_embedding_sharing_contract() -> None: + """Released Muse configs without a top-level field retain text weight tying.""" + hf_config = _tiny_hf_config() + del hf_config.tie_word_embeddings + hf_config.text_config.tie_word_embeddings = True + + model_config = MuseGlimmerBridge().hf_config_to_model_config(hf_config) + + assert model_config.share_embeddings_and_output_weights is True + + +def test_vision_config_contributes_to_runtime_flops() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + cfg = SimpleNamespace(model=model_config) + + vision_flops = vit_flops_from_grid_thw(cfg, torch.tensor([[1, 4, 4]], dtype=torch.int64)) + + assert vision_flops > 0 + + +def test_hybrid_attention_layers_use_full_transformer_flops() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + cfg = SimpleNamespace(model=model_config) + batch_size = 8 + sequence_length = model_config.seq_length + + muse_flops = num_floating_point_operations( + cfg, + batch_size=batch_size, + seqlen_sum=batch_size * sequence_length, + seqlen_squared_sum=batch_size * sequence_length**2, + ) + + model_config.hybrid_attention_layers_include_mlp = False + model_config.hybrid_layer_pattern = None + transformer_flops = num_floating_point_operations( + cfg, + batch_size=batch_size, + seqlen_sum=batch_size * sequence_length, + seqlen_squared_sum=batch_size * sequence_length**2, + ) + + assert muse_flops == transformer_flops + + +@pytest.mark.parametrize( + ("transformer_impl", "expected_qkv", "expected_attention", "expected_norm"), + [ + ("local", ColumnParallelLinear, DotProductAttention, TENorm), + ("transformer_engine", TEColumnParallelLinear, TEDotProductAttention, TENorm), + ], +) +def test_decoder_stack_honors_transformer_backend( + transformer_impl: str, + expected_qkv: type, + expected_attention: type, + expected_norm: type, +) -> None: + config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer + config.transformer_impl = transformer_impl + + stack_spec = get_muse_glimmer_hybrid_stack_spec(config) + attention = stack_spec.submodules.attention_layer.submodules.self_attention.submodules + + assert attention.linear_qkv is expected_qkv + assert attention.core_attention is expected_attention + assert stack_spec.submodules.attention_layer.submodules.input_layernorm is expected_norm + assert stack_spec.submodules.attention_layer.submodules.pre_mlp_layernorm is expected_norm + + +def test_autobridge_selects_string_registration_and_serializes_config() -> None: + auto_bridge = AutoBridge.from_hf_config(_tiny_hf_config()) + model_config = auto_bridge.get_model_config() + + assert isinstance(auto_bridge._model_bridge, MuseGlimmerBridge) + assert isinstance(model_config, MuseGlimmerModelConfig) + assert model_config.builder == "megatron.bridge.models.muse_glimmer.MuseGlimmerModelBuilder" + + restored = MuseGlimmerModelConfig.from_dict(model_config.as_dict()) + assert isinstance(restored, MuseGlimmerModelConfig) + assert restored.vision.layer_types == model_config.vision.layer_types + assert restored.transformer.softmax_scale == model_config.transformer.softmax_scale + assert restored.get_builder_cls() is MuseGlimmerModelBuilder + + +def test_model_config_restores_persisted_dtype_and_enum_values() -> None: + serialized = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).as_dict() + serialized["transformer"]["params_dtype"] = { + "_target_": "torch.bfloat16", + "_call_": False, + } + serialized["transformer"]["attention_backend"] = { + "_target_": "megatron.core.transformer.enums.AttnBackend", + "_call_": True, + "_args_": [5], + "_name_": "auto", + } + + restored = MuseGlimmerModelConfig.from_dict(serialized) + + assert restored.transformer.params_dtype is torch.bfloat16 + assert restored.transformer.attention_backend.name == "auto" + + +def test_config_export_preserves_nested_muse_architecture() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + + exported = MuseGlimmerBridge.megatron_to_hf_config(model_config) + + assert exported["architectures"] == ["MuseGlimmerForConditionalGeneration"] + assert exported["model_type"] == "muse_glimmer" + assert exported["text_config"]["layer_types"] == [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ] + assert exported["text_config"]["layer_rope_theta"] == [500_000.0, 500_000.0, 500_000.0, 0] + assert exported["vision_config"]["layer_types"][-1] == "full_attention" + assert exported["out_hidden_size"] == 64 + + +def test_full_head_gate_qkv_layout_round_trips() -> None: + config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer + query = torch.arange(32 * 64, dtype=torch.float32).reshape(32, 64) + gate = query + 10_000 + key = torch.arange(16 * 64, dtype=torch.float32).reshape(16, 64) + 20_000 + value = key + 30_000 + + query_with_gate = MuseGlimmerQKVGMapping._combine_query_and_gate(config, query, gate) + packed = merge_qkv_weights(config, query_with_gate, key, value) + restored_query_with_gate, restored_key, restored_value = split_qkv_weights(config, packed) + restored_query, restored_gate = restored_query_with_gate.view(4, 16, 64).split(8, dim=1) + + torch.testing.assert_close(restored_query.reshape_as(query), query) + torch.testing.assert_close(restored_gate.reshape_as(gate), gate) + torch.testing.assert_close(restored_key, key) + torch.testing.assert_close(restored_value, value) + + +def test_affine_norms_use_transformer_engine_directly() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + config = model_config.transformer + config.use_cpu_initialization = True + config.params_dtype = torch.float32 + + layer_submodules = get_muse_glimmer_hybrid_stack_spec(config).submodules.attention_layer.submodules + centered_norm = layer_submodules.input_layernorm(config, config.hidden_size, eps=config.layernorm_epsilon) + vision_model = MuseGlimmerVisionModel(model_config.vision, config) + + assert layer_submodules.input_layernorm is TENorm + assert layer_submodules.pre_mlp_layernorm is TENorm + assert isinstance(centered_norm, TERMSNorm) + assert centered_norm.zero_centered_gamma is True + assert isinstance(vision_model.ln_pre, TELayerNorm) + assert isinstance(vision_model.ln_post, TELayerNorm) + assert all(isinstance(layer.norm1, TELayerNorm) for layer in vision_model.layers) + assert all(isinstance(layer.norm2, TELayerNorm) for layer in vision_model.layers) + assert set(vision_model.ln_pre.state_dict()) == {"weight", "bias", "_extra_state"} + + registry = MuseGlimmerBridge().mapping_registry() + assert isinstance( + registry.megatron_to_hf_lookup("decoder.final_norm.weight"), + RMSNorm2ZeroCenteredRMSNormMapping, + ) + for name, _ in vision_model.named_parameters(): + assert registry.megatron_to_hf_lookup(f"vision_tower.{name}") is not None + + +def test_mapping_registry_covers_complete_checkpoint() -> None: + registry = MuseGlimmerBridge().mapping_registry() + + qkvg = registry.megatron_to_hf_lookup("decoder.layers.2.self_attention.linear_qkv.weight") + assert isinstance(qkvg, MuseGlimmerQKVGMapping) + assert qkvg.hf_param["gate"] == "model.language_model.layers.2.self_attn.gate_proj.weight" + assert ( + registry.megatron_to_hf_lookup("vision_tower.layers.1.attn.q_proj.bias").hf_param + == "model.vision_tower.layers.1.attn.q_proj.bias" + ) + assert registry.megatron_to_hf_lookup("vision_adapter.fc2.weight").hf_param == "model.vision_adapter.fc2.weight" + + +@pytest.mark.run_only_on("GPU") +def test_tiny_vision_model_preserves_expected_token_count() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + model_config.transformer_impl = "local" + model_config.params_dtype = torch.float32 + vision_config = model_config.vision + pixel_values = torch.randn(16, 24, device="cuda") + grid_thw = torch.tensor([[1, 4, 4]], device="cuda") + + model = MuseGlimmerVisionModel(vision_config, model_config.transformer).cuda() + output = model(pixel_values, grid_thw) + + assert output.shape == (4, 64) + + +def test_vision_attention_matches_eager_reference_with_finite_backward() -> None: + vision_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).vision + attention = MuseGlimmerVisionAttention(vision_config) + hidden_states = torch.randn(5, vision_config.hidden_size, requires_grad=True) + cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32) + cosine = torch.ones(5, attention.head_dim) + sine = torch.zeros_like(cosine) + real_softmax = torch.nn.functional.softmax + + with ( + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.F.softmax", + wraps=real_softmax, + ) as softmax, + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.F.scaled_dot_product_attention", + side_effect=AssertionError("Muse vision attention must use the HF eager path"), + ), + ): + output = attention(hidden_states, cu_seqlens, (cosine, sine)) + + query = attention.q_proj(hidden_states).reshape(5, attention.num_heads, attention.head_dim) + key = attention.k_proj(hidden_states).reshape(5, attention.num_heads, attention.head_dim) + value = attention.v_proj(hidden_states).reshape(5, attention.num_heads, attention.head_dim) + expected_chunks = [] + for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()): + q = query[start:end].transpose(0, 1).unsqueeze(0) + k = key[start:end].transpose(0, 1).unsqueeze(0) + v = value[start:end].transpose(0, 1).unsqueeze(0) + weights = torch.matmul(q, k.transpose(-2, -1)) * attention.scaling + weights = real_softmax(weights, dim=-1, dtype=torch.float32).to(q.dtype) + expected_chunks.append(torch.matmul(weights, v).squeeze(0).transpose(0, 1)) + expected = attention.proj(torch.cat(expected_chunks, dim=0).reshape(5, -1)) + + assert softmax.call_count == 2 + assert all(call.kwargs["dtype"] is torch.float32 for call in softmax.call_args_list) + torch.testing.assert_close(output, expected) + output.square().mean().backward() + assert torch.isfinite(hidden_states.grad).all() + assert all( + parameter.grad is not None and torch.isfinite(parameter.grad).all() for parameter in attention.parameters() + ) + + +@pytest.mark.run_only_on("GPU") +def test_vision_layer_recomputation_is_training_only() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + model_config.transformer_impl = "local" + model_config.params_dtype = torch.float32 + vision_config = model_config.vision + pixel_values = torch.randn(16, 24, device="cuda") + grid_thw = torch.tensor([[1, 4, 4]], device="cuda") + + with patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.torch_checkpoint", + side_effect=lambda function, *args, **kwargs: function(*args), + ) as checkpoint: + model = MuseGlimmerVisionModel(vision_config, model_config.transformer, recompute_layers=True).cuda() + training_output = model(pixel_values, grid_thw) + model.eval() + inference_output = model(pixel_values, grid_thw) + + assert checkpoint.call_count == vision_config.num_hidden_layers + torch.testing.assert_close(training_output, inference_output) + + +@pytest.fixture(scope="module") +def tiny_hybrid_model() -> Iterator[MuseGlimmerModel]: + auto_bridge = AutoBridge.from_hf_config(_tiny_hf_config()) + model_config = auto_bridge.get_model_config() + model_config.use_cpu_initialization = True + model_config.transformer_impl = "local" + model_config.params_dtype = torch.float32 + model_config.bf16 = False + model_config.bias_activation_fusion = False + model_config.masked_softmax_fusion = False + model_config.persist_layer_norm = False + model_config.bias_dropout_fusion = False + model_config.apply_rope_fusion = False + model_config.gradient_accumulation_fusion = False + model_config.cross_entropy_loss_fusion = False + + models = auto_bridge.get_model( + model_config, + load_weights=False, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + assert len(models) == 1 + model = models[0] + try: + yield model + finally: + from megatron.core import parallel_state + + if parallel_state.is_initialized(): + parallel_state.destroy_model_parallel() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def test_builder_constructs_native_hybrid_model_on_cpu(tiny_hybrid_model: MuseGlimmerModel) -> None: + model = tiny_hybrid_model + + assert isinstance(model, MuseGlimmerModel) + assert isinstance(model, HybridModel) + assert not isinstance(model, GPTModel) + assert model.hybrid_layer_pattern == "****" + assert model.decoder.layer_type_list == ["*", "*", "*", "*"] + names = dict(model.named_parameters()) + assert "vision_tower.patch_embedder.patch_embedding.weight" in names + assert "decoder.layers.0.self_attention.post_layernorm.weight" in names + assert "decoder.layers.0.mlp.post_layernorm.weight" in names + assert not any(name.startswith("language_model.") for name in names) + assert isinstance(model.decoder.layers[0].self_attention.q_layernorm, torch.nn.RMSNorm) + assert isinstance(model.decoder.layers[0].self_attention.k_layernorm, torch.nn.RMSNorm) + assert dict(model.decoder.layers[0].self_attention.q_layernorm.named_parameters()) == {} + assert dict(model.decoder.layers[0].self_attention.k_layernorm.named_parameters()) == {} + assert isinstance(model.decoder.final_norm, TERMSNorm) + assert model.decoder.final_norm.zero_centered_gamma is True + assert isinstance(model.vision_tower.ln_pre, TELayerNorm) + assert isinstance(model.vision_tower.layers[0].norm1, TELayerNorm) + assert names["vision_tower.patch_embedder.patch_embedding.weight"].dtype == torch.float32 + assert all( + getattr(parameter, "average_gradients_across_tp_domain", False) + for name, parameter in names.items() + if name.startswith(("vision_tower.", "vision_adapter.", "vision_projection.")) + ) + + inputs_embeds = torch.randn(2, 5, 64) + input_ids = torch.randint(0, 100, (2, 5)) + expected_output = torch.randn(5, 2, 64) + with ( + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.slice_batch_for_context_parallel", + side_effect=lambda **kwargs: ( + kwargs["inputs_embeds"], + kwargs["labels"], + kwargs["loss_mask"], + kwargs["position_ids"], + kwargs["attention_mask"], + ), + ), + patch.object(HybridModel, "forward", return_value=expected_output) as hybrid_forward, + ): + output = model(input_ids=input_ids, inputs_embeds=inputs_embeds) + + assert output is expected_output + torch.testing.assert_close(hybrid_forward.call_args.kwargs["decoder_input"], inputs_embeds.transpose(0, 1)) + + media_input_ids = input_ids.clone() + media_input_ids[0, 1] = model.model_config.image_token_id + media_input_ids[1, 3] = model.model_config.video_token_id + embedding_output = torch.randn(5, 2, 64) + with ( + patch.object(model.embedding, "forward", return_value=embedding_output) as embedding_forward, + patch.object(HybridModel, "forward", return_value=expected_output), + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.slice_batch_for_context_parallel", + side_effect=lambda **kwargs: ( + kwargs["inputs_embeds"], + kwargs["labels"], + kwargs["loss_mask"], + kwargs["position_ids"], + kwargs["attention_mask"], + ), + ), + ): + model(input_ids=media_input_ids) + + expected_embedding_ids = media_input_ids.clone() + expected_embedding_ids[expected_embedding_ids == model.model_config.image_token_id] = 0 + expected_embedding_ids[expected_embedding_ids == model.model_config.video_token_id] = 0 + torch.testing.assert_close(embedding_forward.call_args.kwargs["input_ids"], expected_embedding_ids) + + +def test_checkpoint_schema_omits_only_empty_backend_norm_state(tiny_hybrid_model: MuseGlimmerModel) -> None: + norm_weight = SimpleNamespace(replica_id=(0, 0, 0)) + state = { + "decoder.layers.0.input_layernorm._extra_state": SimpleNamespace(data=torch.empty(0)), + "decoder.layers.0.self_attention.post_layernorm._extra_state": SimpleNamespace(data=None), + "decoder.layers.0.pre_mlp_layernorm._extra_state": SimpleNamespace(data=None), + "decoder.layers.0.mlp.post_layernorm._extra_state": SimpleNamespace(data=None), + "decoder.final_norm._extra_state": SimpleNamespace(data=None), + "vision_tower.ln_pre._extra_state": SimpleNamespace(data=None), + "vision_tower.layers.0.norm1._extra_state": SimpleNamespace(data=None), + "decoder.layers.0.self_attention.linear_qkv._extra_state": SimpleNamespace(data={"scale": 1}), + "decoder.layers.0.input_layernorm.weight": norm_weight, + } + + with ( + patch.object(HybridModel, "sharded_state_dict", return_value=state), + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.get_pg_rank", + side_effect=[3, 4], + ), + ): + result = tiny_hybrid_model.sharded_state_dict() + + assert "decoder.layers.0.input_layernorm._extra_state" not in result + assert "decoder.layers.0.self_attention.post_layernorm._extra_state" not in result + assert "decoder.layers.0.pre_mlp_layernorm._extra_state" not in result + assert "decoder.layers.0.mlp.post_layernorm._extra_state" not in result + assert "decoder.final_norm._extra_state" not in result + assert "vision_tower.ln_pre._extra_state" not in result + assert "vision_tower.layers.0.norm1._extra_state" not in result + assert "decoder.layers.0.self_attention.linear_qkv._extra_state" in result + assert "decoder.layers.0.input_layernorm.weight" in result + assert norm_weight.replica_id == (0, 3, 4) + + +def test_checkpoint_schema_rejects_stateful_backend_norm_state(tiny_hybrid_model: MuseGlimmerModel) -> None: + state = { + "decoder.layers.0.input_layernorm._extra_state": SimpleNamespace(data=torch.ones(1)), + } + + with ( + patch.object(HybridModel, "sharded_state_dict", return_value=state), + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.get_pg_rank", + side_effect=[3, 4], + ), + pytest.raises(ValueError, match="Transformer Engine norm extra state must be empty"), + ): + tiny_hybrid_model.sharded_state_dict() + + +def test_qkvg_mapping_executes_against_hybrid_qkv_module(tiny_hybrid_model: MuseGlimmerModel) -> None: + model = tiny_hybrid_model + registry = MuseGlimmerBridge().mapping_registry() + for parameter_name, _ in model.named_parameters(): + assert registry.megatron_to_hf_lookup(parameter_name) is not None + + mapping = registry.megatron_to_hf_lookup("decoder.layers.0.self_attention.linear_qkv.weight") + qkv_module = model.decoder.layers[0].self_attention.linear_qkv + query = torch.arange(32 * 64, dtype=torch.float32).reshape(32, 64) + gate = query + 10_000 + key = torch.arange(16 * 64, dtype=torch.float32).reshape(16, 64) + 20_000 + value = key + 30_000 + + packed = mapping.hf_to_megatron({"q": query, "k": key, "v": value, "gate": gate}, qkv_module) + exported = mapping.megatron_to_hf(packed, qkv_module) + + torch.testing.assert_close(exported["model.language_model.layers.0.self_attn.q_proj.weight"], query) + torch.testing.assert_close(exported["model.language_model.layers.0.self_attn.k_proj.weight"], key) + torch.testing.assert_close(exported["model.language_model.layers.0.self_attn.v_proj.weight"], value) + torch.testing.assert_close(exported["model.language_model.layers.0.self_attn.gate_proj.weight"], gate) diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index 5dd342a29c..a8111a9b5b 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -719,6 +719,7 @@ def run_original_hook(models): models = bridge.get_model(model_config, pg_collection=pg_sentinel) assert models == [model_sentinel] + assert model_sentinel.model_config is model_config assert call_order == ["load", "original"] assert build_kwargs["data_parallel_random_init"] is False assert model_config.transformer.perform_initialization is True @@ -1920,6 +1921,7 @@ def test_import_ckpt_basic(self, mock_from_hf_pretrained, mock_to_megatron_model # Setup mocks mock_bridge = Mock(spec=AutoBridge) mock_from_hf_pretrained.return_value = mock_bridge + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = False mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model @@ -1951,6 +1953,7 @@ def test_import_ckpt_with_kwargs(self, mock_from_hf_pretrained, mock_to_megatron # Setup mocks mock_bridge = Mock(spec=AutoBridge) mock_from_hf_pretrained.return_value = mock_bridge + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = False mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model @@ -1995,6 +1998,7 @@ def test_import_ckpt_with_low_memory_save( """Test import_ckpt low-memory save forwarding.""" mock_bridge = Mock(spec=AutoBridge) mock_from_hf_pretrained.return_value = mock_bridge + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = False mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model mock_bridge.save_megatron_model = Mock() @@ -2023,6 +2027,40 @@ def test_import_ckpt_with_low_memory_save( low_memory_save=True, ) + @patch.object(AutoBridge, "from_hf_pretrained") + def test_import_ckpt_uses_builder_for_migrated_model(self, mock_from_hf_pretrained): + """Migrated families import through ModelConfig without a legacy provider.""" + mock_bridge = Mock(spec=AutoBridge) + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = True + model_config = Mock() + model_config.transformer = Mock(use_cpu_initialization=False) + mock_bridge.get_model_config.return_value = model_config + megatron_model = [Mock()] + mock_bridge.get_model.return_value = megatron_model + mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} + mock_from_hf_pretrained.return_value = mock_bridge + + with patch( + "megatron.bridge.training.model_load_save.temporary_distributed_context", + return_value=nullcontext(), + ): + AutoBridge.import_ckpt("hf/model", "/checkpoint") + + assert model_config.transformer.use_cpu_initialization is True + mock_bridge.get_model.assert_called_once_with( + model_config, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + mock_bridge.to_megatron_model.assert_not_called() + mock_bridge.save_megatron_model.assert_called_once_with( + megatron_model, + "/checkpoint", + hf_tokenizer_path="hf/model", + hf_tokenizer_kwargs={}, + low_memory_save=False, + ) + @patch("megatron.bridge.training.model_load_save.temporary_distributed_context") @patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=False) @patch.object(AutoBridge, "from_hf_pretrained") @@ -2034,6 +2072,7 @@ def test_import_ckpt_scopes_standalone_cpu_state_to_gloo_context( ): """Standalone CPU import uses the shared temporary Gloo lifecycle.""" mock_bridge = Mock(spec=AutoBridge) + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = False mock_bridge.to_megatron_model.return_value = [Mock()] mock_bridge.save_megatron_model = Mock() mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} @@ -2057,6 +2096,7 @@ def test_import_ckpt_preserves_existing_distributed_context( ): """Import reuses distributed state owned by its caller.""" mock_bridge = Mock(spec=AutoBridge) + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = False mock_bridge.to_megatron_model.return_value = [Mock()] mock_bridge.save_megatron_model = Mock() mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} diff --git a/tests/unit_tests/models/test_autobridge_registration_matrix.py b/tests/unit_tests/models/test_autobridge_registration_matrix.py index f4030f041d..5c71bfc969 100644 --- a/tests/unit_tests/models/test_autobridge_registration_matrix.py +++ b/tests/unit_tests/models/test_autobridge_registration_matrix.py @@ -59,6 +59,9 @@ "MiniMaxM2ForCausalLM": "megatron.bridge.models.minimax_m2.minimax_m2_bridge.MiniMaxM2Bridge", "MiniMaxM3SparseForConditionalGeneration": ("megatron.bridge.models.minimax_m3.minimax_m3_bridge.MiniMaxM3Bridge"), "Mistral3ForConditionalGeneration": ("megatron.bridge.models.ministral3.ministral3_bridge.Ministral3Bridge"), + "MuseGlimmerForConditionalGeneration": ( + "megatron.bridge.models.muse_glimmer.muse_glimmer_bridge.MuseGlimmerBridge" + ), "NemotronHForCausalLM": "megatron.bridge.models.nemotronh.nemotron_h_bridge.NemotronHBridge", "NemotronH_Nano_Omni_Reasoning_V3": ( "megatron.bridge.models.nemotron_omni.nemotron_omni_bridge.NemotronOmniBridge" @@ -114,6 +117,7 @@ "MiMoV2FlashForCausalLM", "MiniMaxM2ForCausalLM", "MiniMaxM3SparseForConditionalGeneration", + "MuseGlimmerForConditionalGeneration", "NemotronHForCausalLM", "NemotronH_Nano_Omni_Reasoning_V3", "NemotronH_Super_Omni_Reasoning_V3", diff --git a/tests/unit_tests/models/test_conversion_utils.py b/tests/unit_tests/models/test_conversion_utils.py index e232eb067d..fc2a773d3c 100644 --- a/tests/unit_tests/models/test_conversion_utils.py +++ b/tests/unit_tests/models/test_conversion_utils.py @@ -14,7 +14,7 @@ import pytest -from megatron.bridge.models.conversion.utils import mcore_to_hf_window_size +from megatron.bridge.models.conversion.utils import mcore_to_hf_window_size, remove_non_pickleables @pytest.mark.parametrize( @@ -33,3 +33,23 @@ def test_mcore_to_hf_window_size(window_size, expected): def test_mcore_to_hf_window_size_rejects_malformed_pair(): with pytest.raises(ValueError, match="two-element MCore window"): mcore_to_hf_window_size([2047]) + + +def test_remove_non_pickleables_reads_raw_config_attributes(): + class HeterogeneousConfig: + def __init__(self): + self.num_key_value_heads = 8 + self.callback = lambda: None + + def __getattribute__(self, name): + if name == "num_key_value_heads": + raise RuntimeError("attribute must be read from the per-layer config") + return super().__getattribute__(name) + + original = HeterogeneousConfig() + + cleaned = remove_non_pickleables(original) + + assert vars(cleaned)["num_key_value_heads"] == 8 + assert cleaned.callback is None + assert vars(original)["callback"] is not None diff --git a/tests/unit_tests/models/test_gpt_model_config.py b/tests/unit_tests/models/test_gpt_model_config.py index f5da796c4b..24f94559d9 100644 --- a/tests/unit_tests/models/test_gpt_model_config.py +++ b/tests/unit_tests/models/test_gpt_model_config.py @@ -13,6 +13,7 @@ # limitations under the License. import inspect +from dataclasses import dataclass from unittest.mock import Mock, patch import pytest @@ -103,10 +104,47 @@ def test_model_config_mapping_rejects_unknown_fields() -> None: ) +def test_model_config_class_selects_nested_transformer_type(monkeypatch: pytest.MonkeyPatch) -> None: + @dataclass + class SpecializedTransformerConfig(TransformerConfig): + specialized_field: bool = False + + class SpecializedModelConfig(BridgeGPTModelConfig): + transformer_config_class = SpecializedTransformerConfig + + class SpecializedBridge(LlamaBridge): + MODEL_CONFIG_CLASS = SpecializedModelConfig + + bridge = SpecializedBridge() + monkeypatch.setattr( + bridge, + "hf_config_to_model_config_kwargs", + lambda _hf_config: { + "num_layers": 2, + "hidden_size": 128, + "num_attention_heads": 4, + "ffn_hidden_size": 256, + "activation_func": F.silu, + "specialized_field": True, + "vocab_size": 256, + }, + ) + + model_config = bridge.hf_config_to_model_config(object()) + + assert isinstance(model_config, SpecializedModelConfig) + assert isinstance(model_config.transformer, SpecializedTransformerConfig) + assert model_config.specialized_field is True + + def test_gpt_model_config_is_the_bridge_default() -> None: assert MegatronModelBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig assert LlamaBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig assert "MODEL_CONFIG_CLASS" not in LlamaBridge.__dict__ + assert MegatronModelBridge.USE_MODEL_CONFIG_FOR_CONVERSION is False + assert LlamaBridge.USE_MODEL_CONFIG_FOR_CONVERSION is False + assert BridgeGPTModelConfig.transformer_config_class is TransformerConfig + assert not hasattr(MegatronModelBridge, "TRANSFORMER_CONFIG_CLASS") def test_bridge_can_explicitly_disable_model_config() -> None: diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py new file mode 100644 index 0000000000..93fecf530e --- /dev/null +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -0,0 +1,203 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Focused tests for Muse Glimmer verification recipes.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from megatron.bridge.peft.lora import LoRA +from megatron.bridge.perf_recipes.muse_glimmer.h100 import ( + muse_glimmer_30b_pretrain_32gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config, +) +from megatron.bridge.recipes.muse_glimmer.h100 import muse_glimmer as muse_glimmer_recipes +from megatron.bridge.recipes.muse_glimmer.h100 import ( + muse_glimmer_30b_peft_8gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config, + muse_glimmer_30b_sft_32gpu_h100_bf16_config, + muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, +) +from megatron.bridge.training.config import ConfigContainer, MockGPTDatasetConfig +from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies + + +pytestmark = pytest.mark.unit + + +@pytest.fixture(autouse=True) +def _offline_recipe_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: + patch_recipe_construction_dependencies(monkeypatch) + + +@pytest.mark.parametrize( + ( + "recipe", + "seq_length", + "tensor_parallel_size", + "pipeline_parallel_size", + "context_parallel_size", + "global_batch_size", + ), + [ + (muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config, 4096, 8, 2, 1, 2), + (muse_glimmer_30b_sft_32gpu_h100_bf16_config, 4096, 8, 2, 1, 8), + (muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, 8192, 1, 4, 2, 8), + (muse_glimmer_30b_peft_8gpu_h100_bf16_config, 8192, 8, 1, 1, 8), + ], +) +def test_muse_glimmer_recipe_contracts( + recipe, + seq_length, + tensor_parallel_size, + pipeline_parallel_size, + context_parallel_size, + global_batch_size, +) -> None: + cfg = recipe() + + assert isinstance(cfg, ConfigContainer) + assert cfg.model.tensor_model_parallel_size == tensor_parallel_size + assert cfg.model.pipeline_model_parallel_size == pipeline_parallel_size + assert cfg.model.pipeline_dtype is (torch.bfloat16 if pipeline_parallel_size > 1 else None) + expected_pattern = { + 1: "*" * 52, + 2: f"{'*' * 20}|{'*' * 32}", + 4: "|".join("*" * layers for layers in (9, 15, 15, 13)), + }[pipeline_parallel_size] + assert getattr(cfg.model, "hybrid_layer_pattern", "*" * 52) == expected_pattern + assert cfg.model.recompute_vision_layers is True + assert cfg.dataset.hf_processor_kwargs == { + "revision": muse_glimmer_recipes._MODEL_REVISION, + "max_image_tokens": 256, + } + assert getattr(cfg.model, "pipeline_model_parallel_layout", None) is None + assert cfg.model.context_parallel_size == context_parallel_size + assert getattr(cfg.model, "cp_comm_type", "p2p") == ("a2a" if context_parallel_size > 1 else "p2p") + assert cfg.model.sequence_parallel is True + assert cfg.model.recompute_granularity == "selective" + assert cfg.model.recompute_modules == ["core_attn"] + assert cfg.model.seq_length == seq_length + assert cfg.dataset.seq_length == seq_length + assert cfg.dataset.source.load_kwargs == { + "revision": "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret + } + is_long_context = recipe is muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config + assert cfg.dataset.pad_to_max_length is not is_long_context + assert cfg.dataset.enable_in_batch_packing is is_long_context + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == global_batch_size + assert cfg.train.micro_batch_size == (2 if is_long_context else 1) + assert cfg.validation.eval_iters == 0 + assert cfg.validation.eval_interval == 0 + assert cfg.logger.log_throughput is True + + +def test_muse_glimmer_pretrain_owns_resume_checkpoint_contract() -> None: + cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() + + assert cfg.rng.seed == 1234 + assert cfg.scheduler.lr_warmup_iters == 40 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.optimizer.lr == pytest.approx(3e-4) + assert cfg.optimizer.use_precision_aware_optimizer is True + assert cfg.dataset.num_workers == 0 + assert cfg.checkpoint.save_interval == 50 + assert cfg.checkpoint.load is None + + +def test_muse_glimmer_long_context_uses_te_supported_packed_cp_layout() -> None: + cfg = muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() + + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.context_parallel_size == 2 + assert cfg.model.cp_comm_type == "a2a" + assert cfg.model.hybrid_layer_pattern == "|".join("*" * layers for layers in (9, 15, 15, 13)) + assert cfg.train.micro_batch_size == 2 + + +def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: + cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_config() + + assert isinstance(cfg.dataset, MockGPTDatasetConfig) + assert cfg.dataset.seq_length == 4096 + assert cfg.model.hybrid_layer_pattern == "|".join("*" * layers for layers in (9, 15, 15, 13)) + assert cfg.model.freeze_language_model is False + assert cfg.model.freeze_vision_model is True + assert cfg.model.freeze_vision_projection is True + assert cfg.model.recompute_vision_layers is False + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.context_parallel_size == 2 + assert cfg.model.cp_comm_type == "all_gather" + assert cfg.train.train_iters == 50 + assert cfg.train.global_batch_size == 96 + assert cfg.train.micro_batch_size == 3 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 50 + assert cfg.ddp.overlap_grad_reduce is True + assert cfg.ddp.overlap_param_gather is True + assert cfg.comm_overlap.tp_comm_bootstrap_backend == "nccl" + assert cfg.comm_overlap.tp_comm_overlap is True + assert cfg.checkpoint.save is None + assert cfg.checkpoint.load is None + assert cfg.env_vars == { + "CUDA_DEVICE_MAX_CONNECTIONS": 1, + "NCCL_GRAPH_REGISTER": 0, + "NCCL_NVLS_ENABLE": 0, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 0, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 0, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + "TORCH_NCCL_HIGH_PRIORITY": 1, + } + + +def test_muse_glimmer_fp8_performance_recipe_uses_current_scaling() -> None: + cfg = muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config() + + assert cfg.mixed_precision.fp8 == "hybrid" + assert cfg.mixed_precision.fp8_recipe == "tensorwise" + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.context_parallel_size == 2 + assert cfg.train.global_batch_size == 192 + assert cfg.train.micro_batch_size == 6 + + +def test_muse_glimmer_lora_targets_native_attention_projections() -> None: + cfg = muse_glimmer_30b_peft_8gpu_h100_bf16_config() + + assert isinstance(cfg.peft, LoRA) + assert cfg.peft.target_modules == ["linear_qkv", "linear_proj"] + assert cfg.peft.dim == 8 + assert cfg.peft.alpha == 16 + assert cfg.peft.dropout == 0.0 + assert cfg.rng.seed == 5678 + assert cfg.optimizer.lr == pytest.approx(1e-4) + + +def test_muse_glimmer_recipes_use_builder_model_config_api(monkeypatch: pytest.MonkeyPatch) -> None: + class _Bridge: + def get_model_config(self): + return SimpleNamespace() + + def to_megatron_provider(self, *args, **kwargs): + del args, kwargs + raise AssertionError("legacy GPT provider path used") + + class _AutoBridge: + @staticmethod + def from_hf_pretrained(*args, **kwargs): + del args, kwargs + return _Bridge() + + monkeypatch.setattr(muse_glimmer_recipes, "AutoBridge", _AutoBridge) + + cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() + + assert cfg.model is not None diff --git a/tests/unit_tests/scripts/inference/test_vlm_generation.py b/tests/unit_tests/scripts/inference/test_vlm_generation.py new file mode 100644 index 0000000000..720b8148e4 --- /dev/null +++ b/tests/unit_tests/scripts/inference/test_vlm_generation.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Focused tests for deterministic VLM generation configuration.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + + +_SCRIPT_PATH = Path(__file__).parents[4] / "scripts" / "inference" / "vlm_generation.py" +_SPEC = importlib.util.spec_from_file_location("vlm_generation_under_test", _SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +vlm_generation = importlib.util.module_from_spec(_SPEC) +sys.path.insert(0, str(_SCRIPT_PATH.parent)) +try: + _SPEC.loader.exec_module(vlm_generation) +finally: + sys.path.pop(0) + + +pytestmark = pytest.mark.unit + + +def test_model_builder_requires_explicit_conversion_migration() -> None: + legacy_bridge = SimpleNamespace(_model_bridge=SimpleNamespace(MODEL_CONFIG_CLASS=object)) + migrated_bridge = SimpleNamespace( + _model_bridge=SimpleNamespace(MODEL_CONFIG_CLASS=object, USE_MODEL_CONFIG_FOR_CONVERSION=True) + ) + + assert vlm_generation._uses_model_builder(legacy_bridge) is False + assert vlm_generation._uses_model_builder(migrated_bridge) is True + + +def test_deterministic_execution_sets_required_process_environment(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "CUBLAS_WORKSPACE_CONFIG", + "CUDA_DEVICE_MAX_CONNECTIONS", + "NCCL_ALGO", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", + ): + monkeypatch.delenv(name, raising=False) + + with ( + patch.object(vlm_generation.torch, "use_deterministic_algorithms"), + patch.object(vlm_generation.torch, "manual_seed"), + patch.object(vlm_generation.torch.cuda, "manual_seed_all"), + ): + vlm_generation._enable_deterministic_execution() + + assert vlm_generation.os.environ["CUBLAS_WORKSPACE_CONFIG"] == ":4096:8" + assert vlm_generation.os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] == "1" + assert vlm_generation.os.environ["NCCL_ALGO"] == "Ring" + assert vlm_generation.os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" + + +def test_deterministic_execution_rejects_conflicting_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CUBLAS_WORKSPACE_CONFIG", ":16:8") + + with pytest.raises(ValueError, match="CUBLAS_WORKSPACE_CONFIG=:4096:8"): + vlm_generation._enable_deterministic_execution() + + +def test_deterministic_execution_seeds_and_configures_torch(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + monkeypatch.setenv("NCCL_ALGO", "Ring") + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + monkeypatch.setattr(vlm_generation.torch.backends.cudnn, "benchmark", True) + monkeypatch.setattr(vlm_generation.torch.backends.cudnn, "deterministic", False) + + with ( + patch.object(vlm_generation.torch, "use_deterministic_algorithms") as use_deterministic, + patch.object(vlm_generation.torch, "manual_seed") as manual_seed, + patch.object(vlm_generation.torch.cuda, "manual_seed_all") as cuda_manual_seed_all, + ): + vlm_generation._enable_deterministic_execution() + + use_deterministic.assert_called_once_with(True) + manual_seed.assert_called_once_with(0) + cuda_manual_seed_all.assert_called_once_with(0) + assert vlm_generation.torch.backends.cudnn.benchmark is False + assert vlm_generation.torch.backends.cudnn.deterministic is True + + +def test_builder_configuration_propagates_deterministic_mode() -> None: + transformer = SimpleNamespace() + model_config = SimpleNamespace(transformer=transformer) + bridge = SimpleNamespace(get_model_config=MagicMock(return_value=model_config)) + + result = vlm_generation._configure_builder_model( + bridge, + tp=8, + pp=1, + ep=1, + etp=1, + deterministic=True, + ) + + assert result is model_config + assert transformer.deterministic_mode is True + assert transformer.tensor_model_parallel_size == 8 diff --git a/tests/unit_tests/scripts/training/test_recipe_metadata.py b/tests/unit_tests/scripts/training/test_recipe_metadata.py index 336a376e76..f2401d3d7e 100644 --- a/tests/unit_tests/scripts/training/test_recipe_metadata.py +++ b/tests/unit_tests/scripts/training/test_recipe_metadata.py @@ -134,6 +134,7 @@ def test_benchmark_recipe_metadata_selects_task_and_step(recipe_name, task, step ("glm_45v_sft_config", "vlm_step"), ("kimi_k25_vl_sft_config", "vlm_step"), ("ministral3_8b_sft_config", "vlm_step"), + ("muse_glimmer_30b_sft_32gpu_h100_bf16_config", "vlm_step"), ("step37_sft_flickr8k_config", "step37_flickr8k_step"), ("flux_12b_pretrain_config", "flux_step"), ("wan_14b_pretrain_config", "wan_step"), diff --git a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py index efc9463a47..34decc235a 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py @@ -56,6 +56,11 @@ ("moonlight-16b-a3b", "sft_long_context", "H100"): (8192, 128, 8), ("moonlight-16b-a3b", "peft", "H100"): (2048, 32, 4), ("moonlight-16b-a3b", "checkpoint_resume", "H100"): (4096, 1024, 16), + ("muse-glimmer-30b", "pretrain", "H100"): (4096, 2, 32), + ("muse-glimmer-30b", "sft", "H100"): (4096, 8, 32), + ("muse-glimmer-30b", "sft_long_context", "H100"): (8192, 8, 32), + ("muse-glimmer-30b", "peft", "H100"): (8192, 8, 8), + ("muse-glimmer-30b", "pretrain_performance", "H100"): (4096, 192, 32), ("nemotron-3-nano-4b", "pretrain", "H100"): (4096, 1024, 8), ("nemotron-3-nano-4b", "sft", "H100"): (2048, 32, 8), ("nemotron-3-nano-4b", "sft_long_context", "H100"): (32768, 8, 8), diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index 02962f787a..6261964bf3 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -261,6 +261,134 @@ def test_megatron_forward_activates_inference_mode(self): mock_context.__exit__.assert_called_once() mock_forward.assert_called_once_with(forward_only=True) + def test_checkpoint_load_uses_builder_config_when_available(self): + """Hybrid checkpoints initialize model parallelism without a provider shim.""" + transformer = SimpleNamespace() + model_config = SimpleNamespace(transformer=transformer, finalize=MagicMock()) + model = MagicMock() + bridge = MagicMock() + bridge._model_bridge = SimpleNamespace(USE_MODEL_CONFIG_FOR_CONVERSION=True) + bridge.get_model_config.return_value = model_config + bridge.load_megatron_model.return_value = [model] + args = SimpleNamespace( + hf_model_path="hf/model", + hf_revision="revision", + trust_remote_code=False, + megatron_model_path="/checkpoint", + tp=8, + pp=1, + ep=1, + etp=1, + enable_debug_hooks=False, + ) + + with ( + patch.object(compare.AutoBridge, "from_hf_pretrained", return_value=bridge), + patch.object(compare, "is_safe_repo", return_value=False), + patch.object(compare, "disable_mtp_for_inference"), + ): + compare._load_megatron_model(args) + + assert transformer.tensor_model_parallel_size == 8 + assert transformer.pipeline_model_parallel_size == 1 + assert transformer.expert_model_parallel_size == 1 + assert transformer.expert_tensor_parallel_size == 1 + assert transformer.pipeline_dtype is torch.bfloat16 + assert transformer.params_dtype is torch.bfloat16 + assert not hasattr(model_config, "tensor_model_parallel_size") + model_config.finalize.assert_called_once_with() + bridge._get_or_initialize_pg_collection.assert_called_once_with(transformer) + bridge.to_megatron_provider.assert_not_called() + bridge.load_megatron_model.assert_called_once_with( + "/checkpoint", + mp_overrides={ + "tensor_model_parallel_size": 8, + "pipeline_model_parallel_size": 1, + "expert_model_parallel_size": 1, + "expert_tensor_parallel_size": 1, + }, + wrap_with_ddp=False, + ) + + def test_hf_conversion_uses_nested_builder_transformer_config(self): + """Hybrid HF conversion passes nested parallel settings to the builder.""" + transformer = SimpleNamespace() + model_config = SimpleNamespace(transformer=transformer) + model = MagicMock() + bridge = MagicMock() + bridge._model_bridge = SimpleNamespace(USE_MODEL_CONFIG_FOR_CONVERSION=True) + bridge.get_model_config.return_value = model_config + bridge.get_model.return_value = [model] + args = SimpleNamespace( + hf_model_path="hf/model", + hf_revision="revision", + trust_remote_code=False, + megatron_model_path=None, + tp=8, + pp=2, + ep=1, + etp=1, + enable_debug_hooks=False, + ) + + with ( + patch.object(compare.AutoBridge, "from_hf_pretrained", return_value=bridge), + patch.object(compare, "is_safe_repo", return_value=False), + patch.object(compare, "disable_mtp_for_inference"), + ): + compare._load_megatron_model(args) + + assert transformer.tensor_model_parallel_size == 8 + assert transformer.pipeline_model_parallel_size == 2 + assert transformer.expert_model_parallel_size == 1 + assert transformer.expert_tensor_parallel_size == 1 + assert transformer.pipeline_dtype is torch.bfloat16 + assert transformer.params_dtype is torch.bfloat16 + bridge.to_megatron_provider.assert_not_called() + bridge.get_model.assert_called_once_with( + model_config, + wrap_with_ddp=False, + mixed_precision_wrapper=None, + ) + + def test_hf_conversion_keeps_legacy_provider_path(self): + """A legacy bridge with a generic ModelConfig still converts through its provider.""" + model = MagicMock() + provider = MagicMock() + provider.provide_distributed_model.return_value = [model] + bridge = MagicMock() + bridge._model_bridge = SimpleNamespace(USE_MODEL_CONFIG_FOR_CONVERSION=False) + bridge.to_megatron_provider.return_value = provider + args = SimpleNamespace( + hf_model_path="hf/model", + hf_revision="revision", + trust_remote_code=False, + megatron_model_path=None, + tp=2, + pp=1, + ep=1, + etp=1, + enable_debug_hooks=False, + ) + + with ( + patch.object(compare.AutoBridge, "from_hf_pretrained", return_value=bridge), + patch.object(compare, "is_safe_repo", return_value=False), + patch.object(compare, "disable_mtp_for_inference"), + ): + compare._load_megatron_model(args) + + bridge.get_model_config.assert_not_called() + bridge.get_model.assert_not_called() + bridge.to_megatron_provider.assert_called_once_with(load_weights=True) + assert provider.tensor_model_parallel_size == 2 + assert provider.pipeline_model_parallel_size == 1 + assert provider.expert_model_parallel_size == 1 + assert provider.expert_tensor_parallel_size == 1 + assert provider.pipeline_dtype is torch.bfloat16 + provider.finalize.assert_called_once_with() + provider.provide_distributed_model.assert_called_once_with(wrap_with_ddp=False) + def test_tp_logits_skip_gather_when_runtime_output_is_already_full(self): """Test that runtime-gathered text logits are not gathered a second time.""" full_logits = torch.randn(1, 1, 128) diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index c35f1b4228..7deb0b061e 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -17,12 +17,15 @@ import pickle import tempfile from contextlib import ExitStack +from functools import partial from pathlib import Path from unittest.mock import Mock, mock_open, patch import numpy as np import pytest import torch +from megatron.core.dist_checkpointing.strategies.async_utils import AsyncRequest +from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy from megatron.core.msc_utils import MultiStorageClientFeature from megatron.bridge.training.checkpointing import ( @@ -35,9 +38,11 @@ DefaultCheckpointManager, _build_auto_bridge_for_save, _clear_auto_bridge_cache, + _CpuTorchDistSaveShardedStrategy, _extract_megatron_lm_args_from_state_dict, _get_checkpoint_format, _get_non_persistent_iteration, + _get_run_config_tp_pp, _has_global_non_persistent_checkpoint, _load_base_checkpoint, _load_checkpoint_from_path, @@ -95,6 +100,17 @@ def __reduce__(self): class TestCheckpointUtilities: """Test utility functions for checkpoint management.""" + @pytest.mark.parametrize( + "model_config", + [ + {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 2}, + {"transformer": {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 2}}, + ], + ) + def test_get_run_config_tp_pp_supports_flat_and_hybrid_configs(self, model_config): + """Checkpoint parallelism checks support provider and HybridModel layouts.""" + assert _get_run_config_tp_pp(model_config) == (4, 2) + @pytest.mark.parametrize( "checkpoints_path,iteration,release,expected", [ @@ -776,6 +792,7 @@ def test_save_checkpoint_global( mock_get_rng.assert_called_once() else: mock_get_rng.assert_not_called() + assert mock_gen_state.call_args.args[4] is (mock_get_rng.return_value if save_rng else None) # Verify that the tracker file was written with the correct iteration tracker_calls = [ @@ -793,6 +810,54 @@ def test_save_checkpoint_global( written_content = "".join([str(call[0][0]) for call in write_calls if len(call[0]) > 0]) assert "1000" in written_content, f"Expected '1000' in written content, got: {written_content}" + def test_cpu_torch_dist_strategy_uses_blocking_preload(self, tmp_path): + """CPU checkpoint staging must not synchronize an unavailable CUDA device.""" + preload_modes = [] + written_buckets = [] + finalize_devices = [] + + def preload(write_buckets, non_blocking=True): + preload_modes.append(non_blocking) + return write_buckets + + def write(_rank, write_buckets, _results_queue): + written_buckets.extend(write_buckets) + + request = AsyncRequest( + async_fn=write, + async_fn_args=(0, None, None), + finalize_fns=[lambda: finalize_devices.append(torch.cuda.current_device())], + preload_fn=partial(preload, ["cpu-tensor"], True), + ) + strategy = _CpuTorchDistSaveShardedStrategy() + original_current_device = torch.cuda.current_device + + with ( + patch.object(TorchDistSaveShardedStrategy, "async_save", return_value=request), + patch("torch.distributed.barrier"), + ): + strategy.save({}, tmp_path) + + assert preload_modes == [False] + assert written_buckets == ["cpu-tensor"] + assert finalize_devices == [torch.device("cpu")] + assert torch.cuda.current_device is original_current_device + + def test_cpu_torch_dist_strategy_preserves_cuda_finalize_for_nccl(self): + """NCCL cannot broadcast MCore's final status tensor from the CPU.""" + finalize_devices = [] + + with ( + patch("torch.distributed.is_initialized", return_value=True), + patch("torch.distributed.get_backend", return_value="nccl"), + patch("torch.cuda.current_device", return_value=3), + ): + _CpuTorchDistSaveShardedStrategy._run_finalize( + lambda: finalize_devices.append(torch.cuda.current_device()) + ) + + assert finalize_devices == [3] + def test_async_retention_keeps_tracker_checkpoint_until_finalize(self, tmp_path, save_checkpoint_fixtures): """The tracker-selected checkpoint must survive until its async replacement is durable.""" old_checkpoint = tmp_path / "iter_0000500" diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index 6cc7d7e497..8773a3f7c9 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/tests/unit_tests/training/test_model_load_save.py @@ -21,6 +21,7 @@ import pytest import torch from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout +from megatron.training.models.base import ModelConfig from megatron.bridge.models.gpt.gpt_builder import GPTModelConfig from megatron.bridge.models.gpt_provider import GPTModelProvider @@ -945,9 +946,47 @@ def finalize(self) -> None: optimizer=None, opt_param_scheduler=None, num_floating_point_operations_so_far=0, + checkpointing_context=None, callback_manager=None, ) + @patch("megatron.bridge.training.model_load_save.save_checkpoint") + @patch("megatron.bridge.training.model_load_save.get_model_config") + @patch("megatron.bridge.training.model_load_save.GlobalState") + @patch("megatron.bridge.training.model_load_save.ConfigContainer") + @patch("megatron.bridge.training.model_load_save.OptimizerConfig") + @patch("megatron.bridge.training.model_load_save.LoggerConfig") + @patch("megatron.bridge.training.model_load_save.CheckpointConfig") + def test_save_megatron_model_accepts_builder_config( + self, + mock_ckpt_config, + mock_logger_config, + mock_opt_config, + mock_config_container, + mock_global_state, + mock_get_model_config, + mock_save_checkpoint, + ): + """Builder-backed checkpoints serialize the complete outer model config.""" + mock_model = Mock() + model_config = Mock(spec=ModelConfig) + model_config.transformer = SimpleNamespace(use_cpu_initialization=True) + mock_model.model_config = model_config + mock_state = Mock() + mock_global_state.return_value = mock_state + + save_megatron_model([mock_model], "/checkpoint", low_memory_save=False) + + mock_get_model_config.assert_not_called() + assert mock_config_container.call_args.kwargs["model"] is model_config + save_kwargs = mock_save_checkpoint.call_args.kwargs + assert save_kwargs["state"] is mock_state + assert save_kwargs["model"] == [mock_model] + assert isinstance( + save_kwargs["checkpointing_context"]["save_strategy"], + model_load_save._CpuTorchDistSaveShardedStrategy, + ) + @patch("megatron.bridge.training.checkpointing.save_tokenizer_assets") @patch("megatron.bridge.training.checkpointing.get_checkpoint_name") @patch("megatron.bridge.training.model_load_save.build_tokenizer") @@ -1026,6 +1065,7 @@ def finalize(self) -> None: optimizer=None, opt_param_scheduler=None, num_floating_point_operations_so_far=0, + checkpointing_context=None, callback_manager=None, ) @@ -1213,6 +1253,7 @@ def finalize(self) -> None: optimizer=None, opt_param_scheduler=None, num_floating_point_operations_so_far=0, + checkpointing_context=None, callback_manager=None, ) diff --git a/tests/unit_tests/training/utils/test_config_utils.py b/tests/unit_tests/training/utils/test_config_utils.py index ba50d5ae6f..72eedba372 100644 --- a/tests/unit_tests/training/utils/test_config_utils.py +++ b/tests/unit_tests/training/utils/test_config_utils.py @@ -25,6 +25,7 @@ import pytest import torch from megatron.core.msc_utils import MultiStorageClientFeature +from transformers import PreTrainedConfig from megatron.bridge.models.common import Serializable from megatron.bridge.training.utils.config_utils import _ConfigContainerBase, create_ddp_config @@ -774,6 +775,24 @@ class ConfigWithNone(_ConfigContainerBase): assert result["optional_field"] is None assert result["required_field"] == "required" + def test_dataclass_pretrained_config_reads_raw_fields(self): + """Dataclass config serialization must not invoke guarded dynamic access.""" + + @dataclass + class HeterogeneousConfig(PreTrainedConfig): + num_key_value_heads: int = 8 + + def __getattribute__(self, name): + if name == "num_key_value_heads": + raise RuntimeError("read this value from the per-layer config") + return super().__getattribute__(name) + + config = HeterogeneousConfig() + + serialized = _ConfigContainerBase._convert_pretrained_config_to_dict(config, include_target=True) + + assert serialized["num_key_value_heads"] == 8 + def test_config_with_complex_nested_types(self): """Test ConfigContainer with complex nested types.""" diff --git a/tests/unit_tests/utils/test_slice_batch_for_cp.py b/tests/unit_tests/utils/test_slice_batch_for_cp.py index 8e7f71b0b6..cfd1097888 100644 --- a/tests/unit_tests/utils/test_slice_batch_for_cp.py +++ b/tests/unit_tests/utils/test_slice_batch_for_cp.py @@ -204,6 +204,40 @@ def mock_get_batch(batch_dict, is_hybrid_cp=False, cp_group=None, hybrid_cp_grou class TestSliceBatchForContextParallelTHD: """Tests for slice_batch_for_context_parallel with THD (packed) format.""" + def test_thd_format_slices_metadata_without_pipeline_input_embeddings(self): + """Test non-first pipeline stages use labels to construct THD CP indices.""" + labels = torch.arange(16).unsqueeze(0) + loss_mask = torch.ones(1, 16) + position_ids = torch.arange(16).unsqueeze(0) + packed_seq_params = MockPackedSeqParams( + cu_seqlens_q=torch.tensor([0, 8, 16], dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 8, 16], dtype=torch.int32), + qkv_format="thd", + ) + pg_collection = MockPGCollection(cp_size=2, cp_rank=0) + mock_indices = torch.tensor([0, 1, 2, 3, 8, 9, 10, 11]) + + with patch( + "megatron.bridge.training.utils.packed_seq_utils.get_thd_cp_partition_indices", + return_value=mock_indices, + ): + result = slice_batch_for_context_parallel( + inputs_embeds=None, + labels=labels, + loss_mask=loss_mask, + position_ids=position_ids, + attention_mask=None, + packed_seq_params=packed_seq_params, + pg_collection=pg_collection, + ) + + out_embeds, out_labels, out_loss_mask, out_position_ids, out_attention_mask = result + assert out_embeds is None + torch.testing.assert_close(out_labels, labels.index_select(1, mock_indices)) + torch.testing.assert_close(out_loss_mask, loss_mask.index_select(1, mock_indices)) + torch.testing.assert_close(out_position_ids, position_ids.index_select(1, mock_indices)) + assert out_attention_mask is None + def test_thd_format_uses_mcore_partitioned_indices(self): """Test that THD format uses the centralized MCore partition helper.""" batch_size, seq_len, hidden = 1, 16, 64 diff --git a/uv.lock b/uv.lock index e10dffcf84..64af438385 100644 --- a/uv.lock +++ b/uv.lock @@ -2192,7 +2192,7 @@ requires-dist = [ { name = "timm" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", specifier = ">=4.67.1" }, - { name = "transformers", specifier = ">=5.8,<=5.12.1" }, + { name = "transformers", specifier = "<=5.15.0,>=5.15" }, { name = "typing-extensions" }, { name = "wandb", specifier = ">=0.25.0" }, ] @@ -4713,7 +4713,7 @@ dependencies = [ [[package]] name = "transformers" -version = "5.12.1" +version = "5.15.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, @@ -4726,9 +4726,9 @@ dependencies = [ { name = "tqdm" }, { name = "typer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/aa/7c/8240f612819718100a9346dc28dea6a11370c3ca9c8c6eabadd3dea4ef29/transformers-5.12.1.tar.gz", hash = "sha256:679ee731c8225347889ad4fb3b2c926a62e9da3b7d284e9d12c791da7272466b", size = 8924054, upload-time = "2026-06-15T17:27:50.604Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a6/3f/d89353267d511e18f137dfd7769d07837350c11b88408ce1dfe2e93e56c7/transformers-5.15.0.tar.gz", hash = "sha256:bbf98f57b2ddd7c4ecbccfa2c0069017aa6fd01cc204bd50cbc0eeadcf2a13b8", size = 9377983, upload-time = "2026-08-10T10:27:23.261Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/df/56/bbd60dd8668055803bf8ba55a81f9b8a8b31497f620109a9671d26a2076d/transformers-5.12.1-py3-none-any.whl", hash = "sha256:2a5e109d2021265df7098ffbb738295acaf5ad256f12cbc586db2ea4dcbb1a8a", size = 11150587, upload-time = "2026-06-15T17:27:46.679Z" }, + { url = "https://files.pythonhosted.org/packages/d8/43/81355710a4c84e9420e11a86d41a5364deb561f2ef36dfdf254a07371bbb/transformers-5.15.0-py3-none-any.whl", hash = "sha256:d7f007736f67749ae9490c4f8cb5d30b452ae2d68c8675e50ba8d63ea7feb107", size = 11749280, upload-time = "2026-08-10T10:27:20.416Z" }, ] [[package]]