From 207b92a6f53cc0160963787bf0c51179d4803d12 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 10 Aug 2026 11:22:48 -0700 Subject: [PATCH 01/72] feat(model): add Muse Glimmer support Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 200 +++++ examples/models/muse_glimmer/README.md | 52 ++ src/megatron/bridge/models/__init__.py | 11 + .../bridge/models/muse_glimmer/__init__.py | 42 ++ .../configuration_muse_glimmer.py | 188 +++++ .../models/muse_glimmer/model_config.py | 84 +++ .../muse_glimmer/modeling_muse_glimmer.py | 686 ++++++++++++++++++ .../muse_glimmer/muse_glimmer_bridge.py | 401 ++++++++++ .../muse_glimmer/muse_glimmer_builder.py | 80 ++ .../models/muse_glimmer/test_muse_glimmer.py | 294 ++++++++ .../test_autobridge_registration_matrix.py | 4 + 11 files changed, 2042 insertions(+) create mode 100644 examples/model_verification_cards/muse-glimmer-30b/card.yaml create mode 100644 examples/models/muse_glimmer/README.md create mode 100644 src/megatron/bridge/models/muse_glimmer/__init__.py create mode 100644 src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py create mode 100644 src/megatron/bridge/models/muse_glimmer/model_config.py create mode 100644 src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py create mode 100644 src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py create mode 100644 src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py create mode 100644 tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py 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..57832b4624 --- /dev/null +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -0,0 +1,200 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: muse_glimmer_30b +summary: > + Performance disclaimer: this model has not been performance-tuned; reported + timing and throughput metrics are sanity checks, not optimized performance + results. 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 dense text decoder; it does not use a legacy model provider. + Conversion, numerical parity, inference, and training remain unverified until + the corresponding pinned public commands complete their declared gates. +verification_index: + model_level: + unverified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - megatron_to_hf_cpu + - megatron_to_hf_gpu + - manual_forward_pass + - inference + training: + H100: + unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] +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: 6d00aafb874b9d06c7126d2e7239798982ca44d6 # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Import the complete pinned checkpoint through AutoBridge.get_model(), + persist a reloadable Megatron checkpoint, and audit every language, + vision, adapter, and projection tensor for key, shape, dtype, and value + agreement. + + hf_to_megatron_gpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Import the complete pinned checkpoint on GPUs through the builder-backed + path, persist and reload it, and report an exact audit of all mapped + tensors with no missing or unexpected parameters. + + megatron_to_hf_cpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Export a complete CPU-imported Megatron checkpoint to Hugging Face, + strictly reload it with the minimum supported Transformers version, and + compare all source and exported tensors under a declared exact criterion. + + megatron_to_hf_gpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Export a complete GPU-imported Megatron checkpoint, strictly reload the + resulting Hugging Face model, and match all language, vision, adapter, + and projection keys, shapes, dtypes, and values to the pinned source. + + manual_forward_pass: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Process the same pinned image and prompt through native Hugging Face and + Megatron implementations, match the next token, achieve logit cosine + similarity of at least 0.99, and report maximum and mean absolute logit + differences as non-gating diagnostics. + + inference: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Reload the persisted complete Megatron checkpoint, process a pinned + public image and prompt, and complete deterministic bounded generation + with finite logits and a semantically correct response. + + pretrain: + H100: + status: unverified + precision: bf16 + enabled_features: {} + 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 + expected_result: > + Complete a bounded 100-step multimodal pretraining run with finite loss, + no skipped or NaN iterations, all four metrics, and a reloadable final + checkpoint. + + sft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + 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 + expected_result: > + Complete a pinned-data 100-step multimodal SFT run with finite loss, no + skipped or NaN iterations, all four metrics, and a reloadable final + checkpoint. + + sft_export_inference: + H100: + status: unverified + precision: bf16 + depends_on: sft + commands: null + last_verified: null + expected_result: > + Export the completed SFT checkpoint to a complete Hugging Face VLM, + strictly reload it, and complete deterministic image-conditioned + generation with the pinned prompt. + + sft_long_context: + H100: + status: unverified + precision: bf16 + enabled_features: {} + 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 + expected_result: > + Complete a 100-step long-context multimodal SFT run with declared + packing and context-parallel settings, finite loss, no skipped or NaN + iterations, all four metrics, and a reloadable checkpoint. + + peft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + 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 + expected_result: > + Complete a pinned-data 100-step PEFT run with an audited adapter target + set, finite loss, no skipped or NaN iterations, all four metrics, and a + reloadable adapter 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 + 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. diff --git a/examples/models/muse_glimmer/README.md b/examples/models/muse_glimmer/README.md new file mode 100644 index 0000000000..c23f21c6c6 --- /dev/null +++ b/examples/models/muse_glimmer/README.md @@ -0,0 +1,52 @@ +# 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`; no model provider is involved. + +The text decoder uses native Megatron-Core components for 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", +) +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/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/muse_glimmer/__init__.py b/src/megatron/bridge/models/muse_glimmer/__init__.py new file mode 100644 index 0000000000..c78333f11e --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/__init__.py @@ -0,0 +1,42 @@ +# 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 megatron.bridge.models.muse_glimmer.configuration_muse_glimmer import ( + MuseGlimmerConfig, + MuseGlimmerTextConfig, + MuseGlimmerVisionConfig, +) +from megatron.bridge.models.muse_glimmer.model_config import ( + MuseGlimmerModelConfig, + MuseGlimmerTransformerConfig, + MuseGlimmerVisionModelConfig, +) +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 + + +__all__ = [ + "MuseGlimmerBridge", + "MuseGlimmerConfig", + "MuseGlimmerModel", + "MuseGlimmerModelBuilder", + "MuseGlimmerModelConfig", + "MuseGlimmerTextConfig", + "MuseGlimmerTransformerConfig", + "MuseGlimmerVisionConfig", + "MuseGlimmerVisionModelConfig", +] diff --git a/src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py new file mode 100644 index 0000000000..dbdba0f470 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py @@ -0,0 +1,188 @@ +# 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. + +"""Transformers compatibility configuration for Muse Glimmer. + +Muse Glimmer landed after the Transformers version currently supported by +Megatron Bridge. Weight import only needs a parsed configuration and the lazy +safetensors state source, so keep a small configuration fallback here instead +of requiring the full Hugging Face modeling implementation at runtime. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig, PretrainedConfig + + +try: + from transformers import MuseGlimmerConfig, MuseGlimmerTextConfig, MuseGlimmerVisionConfig +except ImportError: + + class MuseGlimmerTextConfig(PretrainedConfig): + """Configuration for the Muse Glimmer text decoder.""" + + model_type = "muse_glimmer_text" + + def __init__( + self, + *, + vocab_size: int = 202_048, + hidden_size: int = 6_656, + intermediate_size: int = 19_968, + num_hidden_layers: int = 52, + num_attention_heads: int = 32, + num_key_value_heads: int = 2, + head_dim: int = 128, + hidden_activation: str = "silu", + max_position_embeddings: int = 131_072, + rms_norm_eps: float = 1e-5, + sliding_window: int = 2_048, + final_logit_softcapping: float = 20.0, + qk_scale_factor: float = 3.87, + output_multiplier: float = 0.19611613513818404, + post_norm_eps: float = 1e-8, + layer_types: list[str] | None = None, + layer_rope_theta: list[float] | None = None, + rope_parameters: dict[str, Any] | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + initializer_range: float = 0.02, + tie_word_embeddings: bool = False, + bos_token_id: int | None = 200_000, + eos_token_id: int | list[int] | None = 200_001, + pad_token_id: int | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_activation = hidden_activation + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.sliding_window = sliding_window + self.final_logit_softcapping = final_logit_softcapping + self.qk_scale_factor = qk_scale_factor + self.output_multiplier = output_multiplier + self.post_norm_eps = post_norm_eps + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.initializer_range = initializer_range + self.rope_parameters = rope_parameters or {"rope_theta": 500_000.0, "rope_type": "default"} + self.layer_types = layer_types or [ + "full_attention" if (num_hidden_layers - 1 - index) % 4 == 0 else "sliding_attention" + for index in range(num_hidden_layers) + ] + self.layer_rope_theta = layer_rope_theta or [ + 0.0 if layer_type == "full_attention" else float(self.rope_parameters["rope_theta"]) + for layer_type in self.layer_types + ] + + class MuseGlimmerVisionConfig(PretrainedConfig): + """Configuration for the Muse Glimmer vision encoder.""" + + model_type = "muse_glimmer_vision" + + def __init__( + self, + *, + 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_eps: float = 1e-5, + hidden_act: str = "gelu", + layer_types: list[str] | None = None, + rope_parameters: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.patch_size = patch_size + self.patch_temporal = patch_temporal + self.merge_size = merge_size + self.pos_emb_height = pos_emb_height + self.pos_emb_width = pos_emb_width + self.max_position_embeddings = max_position_embeddings + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.rope_parameters = rope_parameters or {"rope_theta": 10_000.0, "rope_type": "default"} + self.layer_types = layer_types or [ + "full_attention" if index % 4 == 3 or index == num_hidden_layers - 1 else "window_attention" + for index in range(num_hidden_layers) + ] + + class MuseGlimmerConfig(PretrainedConfig): + """Top-level configuration for Muse Glimmer.""" + + model_type = "muse_glimmer" + sub_configs = { + "text_config": MuseGlimmerTextConfig, + "vision_config": MuseGlimmerVisionConfig, + } + + def __init__( + self, + *, + text_config: dict[str, Any] | PretrainedConfig | None = None, + vision_config: dict[str, Any] | PretrainedConfig | None = None, + image_token_id: int = 200_092, + video_token_id: int = 200_091, + out_hidden_size: int = 6_144, + projector_hidden_size: int = 4_096, + projector_hidden_act: str = "gelu", + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.text_config = ( + text_config + if isinstance(text_config, PretrainedConfig) + else MuseGlimmerTextConfig(**(text_config or {})) + ) + self.vision_config = ( + vision_config + if isinstance(vision_config, PretrainedConfig) + else MuseGlimmerVisionConfig(**(vision_config or {})) + ) + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.out_hidden_size = out_hidden_size + self.projector_hidden_size = projector_hidden_size + self.projector_hidden_act = projector_hidden_act + + AutoConfig.register(MuseGlimmerConfig.model_type, MuseGlimmerConfig, exist_ok=True) + + +__all__ = ["MuseGlimmerConfig", "MuseGlimmerTextConfig", "MuseGlimmerVisionConfig"] diff --git a/src/megatron/bridge/models/muse_glimmer/model_config.py b/src/megatron/bridge/models/muse_glimmer/model_config.py new file mode 100644 index 0000000000..0990831b5f --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/model_config.py @@ -0,0 +1,84 @@ +# 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. + +"""Builder-backed configuration objects for Muse Glimmer.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import ClassVar + +from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig +from megatron.bridge.models.transformer_config import TransformerConfig + + +@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) + + +@dataclass(kw_only=True) +class MuseGlimmerModelConfig(BridgeGPTModelConfig): + """Complete builder configuration for Muse Glimmer.""" + + builder: ClassVar[str] = "megatron.bridge.models.muse_glimmer.MuseGlimmerModelBuilder" + + 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 + + @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} + + +__all__ = [ + "MuseGlimmerModelConfig", + "MuseGlimmerTransformerConfig", + "MuseGlimmerVisionModelConfig", +] 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..f0503b3ef6 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -0,0 +1,686 @@ +# 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.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.models.backends import LocalSpecProvider +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.module import MegatronModule +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 torch import Tensor + +from megatron.bridge.models.gemma.modules import extend_instance +from megatron.bridge.models.muse_glimmer.model_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 + + +def _rms_norm(hidden_states: Tensor, eps: float) -> Tensor: + """Apply the float32 RMS expression used by the Hugging Face reference.""" + normalized = hidden_states.float() * torch.pow( + hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + eps, + -0.5, + ) + return normalized.to(hidden_states.dtype) + + +def _centered_rms_norm(hidden_states: Tensor, eps: float) -> Tensor: + """Apply the rsqrt-based expression used by Muse's centered layer norms.""" + normalized = hidden_states.float() * torch.rsqrt(hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + eps) + return normalized.to(hidden_states.dtype) + + +class MuseGlimmerCenteredRMSNorm(nn.Module): + """Gemma2-style RMSNorm whose checkpoint weight is centered around zero.""" + + def __init__(self, config: TransformerConfig, hidden_size: int, eps: float = 1e-5) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=config.params_dtype)) + setattr(self.weight, "sequence_parallel", bool(config.sequence_parallel)) + self.eps = eps + + def forward(self, hidden_states: Tensor) -> Tensor: + return _centered_rms_norm(hidden_states, self.eps) * (1.0 + self.weight.float()).to(hidden_states.dtype) + + +class MuseGlimmerRMSNorm(nn.Module): + """Learnable RMSNorm used after the final Muse decoder layer.""" + + def __init__(self, config: TransformerConfig, hidden_size: int, eps: float = 1e-5) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=config.params_dtype)) + setattr(self.weight, "sequence_parallel", bool(config.sequence_parallel)) + self.eps = eps + + def forward(self, hidden_states: Tensor) -> Tensor: + return _rms_norm(hidden_states, self.eps) * self.weight.float().to(hidden_states.dtype) + + +class MuseGlimmerWeightlessRMSNorm(nn.Module): + """Scaleless RMSNorm used for Muse Q/K, embeddings, and perception features.""" + + def __init__(self, config: TransformerConfig, hidden_size: int, eps: float = 1e-5) -> None: + super().__init__() + del hidden_size + self.eps = eps + + def forward(self, hidden_states: Tensor) -> Tensor: + return _rms_norm(hidden_states, self.eps) + + +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) + self.post_layernorm = MuseGlimmerCenteredRMSNorm( + 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 = MuseGlimmerCenteredRMSNorm( + 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_layer_spec(config: TransformerConfig) -> ModuleSpec: + """Build the local MCore layer spec used by the Muse decoder.""" + del config + backend = LocalSpecProvider() + return ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=MuseGlimmerCenteredRMSNorm, + 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(), + q_layernorm=MuseGlimmerWeightlessRMSNorm, + k_layernorm=MuseGlimmerWeightlessRMSNorm, + ), + ), + self_attn_bda=get_bias_dropout_add, + pre_mlp_layernorm=MuseGlimmerCenteredRMSNorm, + 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 MuseGlimmerEmbeddingNormMixin(nn.Module): + """Apply the weightless Muse RMSNorm after word embedding lookup.""" + + def forward(self, *args: Any, **kwargs: Any) -> Tensor: + embeddings = super().forward(*args, **kwargs) + return _rms_norm(embeddings, self.config.layernorm_epsilon) + + +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 customize_muse_glimmer_language_model(model: nn.Module) -> None: + """Install parameter-free Muse embedding/logit behavior on an MCore GPT model.""" + if hasattr(model, "embedding"): + extend_instance(model.embedding, MuseGlimmerEmbeddingNormMixin) + if hasattr(model, "decoder") and model.decoder.final_layernorm is not None: + model.decoder.final_layernorm = MuseGlimmerRMSNorm( + model.config, + model.config.hidden_size, + eps=model.config.layernorm_epsilon, + ) + if hasattr(model, "output_layer"): + extend_instance(model.output_layer, MuseGlimmerOutputLayerMixin) + + +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) + output = F.scaled_dot_product_attention(q, k, v, scale=self.scaling) + 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) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.attn = MuseGlimmerVisionAttention(config) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + 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) -> None: + super().__init__() + self.config = config + self.patch_embedder = MuseGlimmerVisionPatchEmbedder(config) + self.rotary_emb = MuseGlimmerVisionRotaryEmbedding(config) + self.ln_pre = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.layers = nn.ModuleList([MuseGlimmerVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.ln_post = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + 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 + 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(MegatronModule): + """Combined Muse vision tower, adapter, and MCore language model.""" + + def __init__( + self, + config: MuseGlimmerModelConfig, + language_model: MegatronModule, + *, + pre_process: bool, + post_process: bool, + vp_stage: int | None, + ) -> None: + super().__init__(config=config.transformer) + self.model_config = config + self.pre_process = pre_process + self.post_process = post_process + self.vp_stage = vp_stage + self.language_model = language_model + self.pg_collection = language_model.pg_collection + if pre_process: + self.vision_tower = MuseGlimmerVisionModel(config.vision) + 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.share_embeddings_and_output_weights = config.share_embeddings_and_output_weights + self.shared_embedding_or_output_weight = language_model.shared_embedding_or_output_weight + self.freeze( + freeze_language_model=config.freeze_language_model, + freeze_vision_model=config.freeze_vision_model, + freeze_vision_projection=config.freeze_vision_projection, + ) + + @property + def decoder(self) -> nn.Module | None: + """Expose the text decoder to Megatron-Core inference utilities.""" + return getattr(self.language_model, "decoder", None) + + def set_input_tensor(self, input_tensor: Tensor | list[Tensor]) -> None: + self.language_model.set_input_tensor(input_tensor) + + 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.append(self.language_model) + 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 _rms_norm(features, 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, + extra_block_kwargs: dict[str, Any] | None = None, + *, + inference_params: BaseInferenceContext | None = None, + loss_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.language_model.embedding(input_ids=llm_input_ids, position_ids=None) + 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 = self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=inputs_embeds, + labels=labels, + loss_mask=loss_mask, + inference_context=inference_context, + runtime_gather_output=runtime_gather_output, + packed_seq_params=packed_seq_params, + extra_block_kwargs=extra_block_kwargs, + inference_params=inference_params, + ) + if loss_mask is not None: + return output, loss_mask + return output + + +__all__ = [ + "MuseGlimmerCenteredRMSNorm", + "MuseGlimmerModel", + "MuseGlimmerRMSNorm", + "MuseGlimmerVisionModel", + "MuseGlimmerWeightlessRMSNorm", + "customize_muse_glimmer_language_model", + "get_muse_glimmer_layer_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..531372c292 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -0,0 +1,401 @@ +# 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, + merge_qkv_weights, + split_qkv_weights, +) +from megatron.bridge.models.conversion.utils import remove_non_pickleables +from megatron.bridge.models.muse_glimmer.model_config import ( + MuseGlimmerModelConfig, + MuseGlimmerTransformerConfig, + MuseGlimmerVisionModelConfig, +) +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import MuseGlimmerModel + + +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 + TRANSFORMER_CONFIG_CLASS = MuseGlimmerTransformerConfig + + @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 HF config directly into the new model-config/builder path.""" + 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(text_config.tie_word_embeddings), + position_embedding_type="rope", + rotary_percent=1.0, + rotary_base=rotary_base, + scatter_embedding_sequence_parallel=False, + 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="language_model.embedding.word_embeddings.weight", + hf_param=f"{text_prefix}.embed_tokens.weight", + ), + AutoMapping( + megatron_param="language_model.output_layer.weight", + hf_param="lm_head.weight", + ), + ReplicatedMapping( + megatron_param="language_model.decoder.final_layernorm.weight", + hf_param=f"{text_prefix}.norm.weight", + ), + ReplicatedMapping( + megatron_param="language_model.decoder.layers.*.input_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.input_layernorm.weight", + ), + ReplicatedMapping( + megatron_param="language_model.decoder.layers.*.self_attention.post_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.post_attention_layernorm.weight", + ), + ReplicatedMapping( + megatron_param="language_model.decoder.layers.*.pre_mlp_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.pre_feedforward_layernorm.weight", + ), + ReplicatedMapping( + megatron_param="language_model.decoder.layers.*.mlp.post_layernorm.weight", + hf_param=f"{text_prefix}.layers.*.post_feedforward_layernorm.weight", + ), + AutoMapping( + megatron_param="language_model.decoder.layers.*.self_attention.linear_proj.weight", + hf_param=f"{text_prefix}.layers.*.self_attn.o_proj.weight", + ), + AutoMapping( + megatron_param="language_model.decoder.layers.*.mlp.linear_fc2.weight", + hf_param=f"{text_prefix}.layers.*.mlp.down_proj.weight", + ), + MuseGlimmerQKVGMapping( + "language_model.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="language_model.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..8c14c1c604 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py @@ -0,0 +1,80 @@ +# 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 dataclasses import replace + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.training.models.gpt import GPTModelBuilder + +from megatron.bridge.models.muse_glimmer.model_config import MuseGlimmerModelConfig +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( + MuseGlimmerModel, + customize_muse_glimmer_language_model, + get_muse_glimmer_layer_spec, +) + + +class MuseGlimmerModelBuilder(GPTModelBuilder): + """Construct the replicated vision stack around a builder-backed MCore decoder.""" + + 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__}.") + + language_config = replace( + model_config, + transformer_layer_spec=get_muse_glimmer_layer_spec(model_config.transformer), + ) + language_model = GPTModelBuilder(language_config).build_model( + pg_collection, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + customize_muse_glimmer_language_model(language_model) + return MuseGlimmerModel( + model_config, + language_model, + pre_process=language_model.pre_process, + post_process=language_model.post_process, + vp_stage=vp_stage, + ) + + +__all__ = ["MuseGlimmerModelBuilder"] 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..0759f99b3a --- /dev/null +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -0,0 +1,294 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for the builder-backed Muse Glimmer implementation.""" + +from __future__ import annotations + +import math +from unittest.mock import patch + +import pytest +import torch + +from megatron.bridge import AutoBridge +from megatron.bridge.models.conversion.param_mapping import merge_qkv_weights, split_qkv_weights +from megatron.bridge.models.muse_glimmer import ( + MuseGlimmerConfig, + MuseGlimmerModel, + MuseGlimmerModelBuilder, + MuseGlimmerModelConfig, + MuseGlimmerTextConfig, + MuseGlimmerVisionConfig, +) +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( + MuseGlimmerCenteredRMSNorm, + MuseGlimmerRMSNorm, + MuseGlimmerVisionModel, +) +from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import ( + MuseGlimmerBridge, + MuseGlimmerQKVGMapping, +) + + +pytestmark = pytest.mark.unit + + +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"], + 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 model_config.get_builder_cls() is MuseGlimmerModelBuilder + 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.scatter_embedding_sequence_parallel is False + assert model_config.special_token_ids == {"images": 120, "videos": 121} + + +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_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_centered_and_final_norms_use_their_native_expressions() -> None: + config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer + hidden_states = torch.tensor([[0.25, -0.5, 1.0, -2.0]], dtype=torch.float32) + centered = MuseGlimmerCenteredRMSNorm(config, hidden_size=4, eps=1e-5) + final = MuseGlimmerRMSNorm(config, hidden_size=4, eps=1e-5) + with torch.no_grad(): + centered.weight.fill_(0.25) + final.weight.fill_(1.25) + + mean_square = hidden_states.pow(2).mean(dim=-1, keepdim=True) + expected_centered = hidden_states * torch.rsqrt(mean_square + 1e-5) * 1.25 + expected_final = hidden_states * torch.pow(mean_square + 1e-5, -0.5) * 1.25 + + torch.testing.assert_close(centered(hidden_states), expected_centered) + torch.testing.assert_close(final(hidden_states), expected_final) + + +def test_mapping_registry_covers_complete_checkpoint() -> None: + registry = MuseGlimmerBridge().mapping_registry() + + qkvg = registry.megatron_to_hf_lookup("language_model.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" + + +def test_tiny_vision_model_preserves_expected_token_count() -> None: + vision_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).vision + model = MuseGlimmerVisionModel(vision_config) + pixel_values = torch.randn(16, 24) + grid_thw = torch.tensor([[1, 4, 4]]) + + output = model(pixel_values, grid_thw) + + assert output.shape == (4, 64) + + +def test_builder_constructs_combined_model_on_cpu() -> None: + auto_bridge = AutoBridge.from_hf_config(_tiny_hf_config()) + model_config = auto_bridge.get_model_config() + model_config.use_cpu_initialization = True + 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, + ) + try: + assert len(models) == 1 + model = models[0] + assert isinstance(model, MuseGlimmerModel) + names = dict(model.named_parameters()) + assert "vision_tower.patch_embedder.patch_embedding.weight" in names + assert "language_model.decoder.layers.0.self_attention.post_layernorm.weight" in names + assert "language_model.decoder.layers.0.mlp.post_layernorm.weight" in names + assert isinstance(model.language_model.decoder.final_layernorm, MuseGlimmerRMSNorm) + 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(model.language_model, "forward", return_value=expected_output) as language_forward, + ): + output = model(input_ids=input_ids, inputs_embeds=inputs_embeds) + + assert output is expected_output + torch.testing.assert_close(language_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.language_model.embedding, "forward", return_value=embedding_output + ) as embedding_forward, + patch.object(model.language_model, "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) + 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() diff --git a/tests/unit_tests/models/test_autobridge_registration_matrix.py b/tests/unit_tests/models/test_autobridge_registration_matrix.py index d2396f02b4..f69f29920d 100644 --- a/tests/unit_tests/models/test_autobridge_registration_matrix.py +++ b/tests/unit_tests/models/test_autobridge_registration_matrix.py @@ -58,6 +58,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" @@ -113,6 +116,7 @@ "MiMoV2FlashForCausalLM", "MiniMaxM2ForCausalLM", "MiniMaxM3SparseForConditionalGeneration", + "MuseGlimmerForConditionalGeneration", "NemotronHForCausalLM", "NemotronH_Nano_Omni_Reasoning_V3", "NemotronH_Super_Omni_Reasoning_V3", From 2e3b8f3667daf371ed403396d7b6554cac9da459 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 10 Aug 2026 11:31:21 -0700 Subject: [PATCH 02/72] docs(model): pin Muse Glimmer verification commit Signed-off-by: yaoyu-33 --- examples/model_verification_cards/muse-glimmer-30b/card.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index 57832b4624..0712a5530a 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -30,7 +30,7 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: 6d00aafb874b9d06c7126d2e7239798982ca44d6 # pragma: allowlist secret + bridge_commit: 2e5efbab0e2039446693f61e2fc99d0233832775 # pragma: allowlist secret items: hf_to_megatron_cpu: From 95bc2220f8d721744388a1e6dc8b32df34ec421c Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 13:53:11 -0700 Subject: [PATCH 03/72] refactor(model): migrate Muse Glimmer to HybridModel Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 3 +- examples/models/muse_glimmer/README.md | 21 ++- .../models/muse_glimmer/model_config.py | 53 +++++- .../muse_glimmer/modeling_muse_glimmer.py | 131 ++++++++------ .../muse_glimmer/muse_glimmer_bridge.py | 26 +-- .../muse_glimmer/muse_glimmer_builder.py | 41 +++-- .../models/muse_glimmer/test_muse_glimmer.py | 167 +++++++++++------- 7 files changed, 277 insertions(+), 165 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index 0712a5530a..782ce5c4a6 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -8,7 +8,8 @@ summary: > results. 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 dense text decoder; it does not use a legacy model provider. + projection, and a native Megatron-Core HybridModel text decoder; it does not + use a legacy model provider or GPTModel wrapper. Conversion, numerical parity, inference, and training remain unverified until the corresponding pinned public commands complete their declared gates. verification_index: diff --git a/examples/models/muse_glimmer/README.md b/examples/models/muse_glimmer/README.md index c23f21c6c6..332134443c 100644 --- a/examples/models/muse_glimmer/README.md +++ b/examples/models/muse_glimmer/README.md @@ -9,14 +9,17 @@ projection, and text decoder all participate in checkpoint conversion. Muse Glimmer uses the builder-backed API. `MuseGlimmerBridge` translates the Hugging Face configuration directly into `MuseGlimmerModelConfig`, and -`MuseGlimmerModelBuilder` constructs `MuseGlimmerModel`; no model provider is involved. - -The text decoder uses native Megatron-Core components for 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. +`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, @@ -30,7 +33,7 @@ from megatron.bridge import AutoBridge bridge = AutoBridge.from_hf_pretrained( "meta-models/Muse-Glimmer-30B", - revision="f84ecc3a0ea984a4c04542a84269e3d065350a6e", + revision="f84ecc3a0ea984a4c04542a84269e3d065350a6e", # pragma: allowlist secret ) model_config = bridge.get_model_config() diff --git a/src/megatron/bridge/models/muse_glimmer/model_config.py b/src/megatron/bridge/models/muse_glimmer/model_config.py index 0990831b5f..c9aafebe2b 100644 --- a/src/megatron/bridge/models/muse_glimmer/model_config.py +++ b/src/megatron/bridge/models/muse_glimmer/model_config.py @@ -17,10 +17,14 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import ClassVar +from typing import Any, ClassVar -from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig +from megatron.training.models.hybrid import HybridModelConfig + +from megatron.bridge.models.common import ModelConfig 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 @@ -53,7 +57,7 @@ class MuseGlimmerVisionModelConfig: @dataclass(kw_only=True) -class MuseGlimmerModelConfig(BridgeGPTModelConfig): +class MuseGlimmerModelConfig(HybridModelConfig): """Complete builder configuration for Muse Glimmer.""" builder: ClassVar[str] = "megatron.bridge.models.muse_glimmer.MuseGlimmerModelBuilder" @@ -76,6 +80,49 @@ 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 = ModelConfig.from_dict(restored_data) + if not isinstance(result, cls): + raise TypeError(f"Expected {cls.__name__}, got {type(result).__name__}.") + return result + __all__ = [ "MuseGlimmerModelConfig", diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index f0503b3ef6..13fc82726c 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -24,11 +24,13 @@ import torch.nn.functional as F from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.backends import LocalSpecProvider +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.module import MegatronModule 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 @@ -147,35 +149,40 @@ def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: return self.post_layernorm(output), bias -def get_muse_glimmer_layer_spec(config: TransformerConfig) -> ModuleSpec: - """Build the local MCore layer spec used by the Muse decoder.""" +def get_muse_glimmer_hybrid_stack_spec(config: TransformerConfig) -> ModuleSpec: + """Build the native Hybrid stack spec used by the Muse decoder.""" del config backend = LocalSpecProvider() return ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - input_layernorm=MuseGlimmerCenteredRMSNorm, - 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(), - q_layernorm=MuseGlimmerWeightlessRMSNorm, - k_layernorm=MuseGlimmerWeightlessRMSNorm, + module=HybridStack, + submodules=HybridStackSubmodules( + attention_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=MuseGlimmerCenteredRMSNorm, + 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(), + q_layernorm=MuseGlimmerWeightlessRMSNorm, + k_layernorm=MuseGlimmerWeightlessRMSNorm, + ), + ), + self_attn_bda=get_bias_dropout_add, + pre_mlp_layernorm=MuseGlimmerCenteredRMSNorm, + 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, ), ), - self_attn_bda=get_bias_dropout_add, - pre_mlp_layernorm=MuseGlimmerCenteredRMSNorm, - 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, ), ) @@ -200,12 +207,12 @@ def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: return output, bias -def customize_muse_glimmer_language_model(model: nn.Module) -> None: - """Install parameter-free Muse embedding/logit behavior on an MCore GPT model.""" +def customize_muse_glimmer_language_model(model: HybridModel) -> None: + """Install Muse embedding, final-norm, and logit behavior on a Hybrid model.""" if hasattr(model, "embedding"): extend_instance(model.embedding, MuseGlimmerEmbeddingNormMixin) - if hasattr(model, "decoder") and model.decoder.final_layernorm is not None: - model.decoder.final_layernorm = MuseGlimmerRMSNorm( + if hasattr(model, "decoder") and getattr(model.decoder, "final_norm", None) is not None: + model.decoder.final_norm = MuseGlimmerRMSNorm( model.config, model.config.hidden_size, eps=model.config.layernorm_epsilon, @@ -517,25 +524,41 @@ def forward(self, hidden_states: Tensor) -> Tensor: return F.gelu(self.fc2(F.gelu(self.fc1(hidden_states)))) -class MuseGlimmerModel(MegatronModule): - """Combined Muse vision tower, adapter, and MCore language model.""" +class MuseGlimmerModel(HybridModel): + """Native MCore Hybrid model with the Muse vision modules attached.""" def __init__( self, config: MuseGlimmerModelConfig, - language_model: MegatronModule, + 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) + 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 - self.pre_process = pre_process - self.post_process = post_process - self.vp_stage = vp_stage - self.language_model = language_model - self.pg_collection = language_model.pg_collection + customize_muse_glimmer_language_model(self) if pre_process: self.vision_tower = MuseGlimmerVisionModel(config.vision) self.vision_adapter = MuseGlimmerVisionAdapter(config) @@ -547,22 +570,12 @@ def __init__( 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.share_embeddings_and_output_weights = config.share_embeddings_and_output_weights - self.shared_embedding_or_output_weight = language_model.shared_embedding_or_output_weight self.freeze( freeze_language_model=config.freeze_language_model, freeze_vision_model=config.freeze_vision_model, freeze_vision_projection=config.freeze_vision_projection, ) - @property - def decoder(self) -> nn.Module | None: - """Expose the text decoder to Megatron-Core inference utilities.""" - return getattr(self.language_model, "decoder", None) - - def set_input_tensor(self, input_tensor: Tensor | list[Tensor]) -> None: - self.language_model.set_input_tensor(input_tensor) - def freeze( self, *, @@ -572,7 +585,15 @@ def freeze( ) -> None: modules: list[nn.Module] = [] if freeze_language_model: - modules.append(self.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"): @@ -610,10 +631,10 @@ def forward( inference_context: BaseInferenceContext | None = None, runtime_gather_output: bool | None = None, packed_seq_params: PackedSeqParams | None = None, - extra_block_kwargs: dict[str, Any] | 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: @@ -623,7 +644,7 @@ def forward( input_ids == self.model_config.video_token_id ) llm_input_ids = input_ids.masked_fill(multimodal_mask, 0) - inputs_embeds = self.language_model.embedding(input_ids=llm_input_ids, position_ids=None) + inputs_embeds = self.embedding(input_ids=llm_input_ids, position_ids=None) 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: @@ -657,18 +678,18 @@ def forward( 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 = self.language_model( + output = super().forward( input_ids=None, position_ids=position_ids, attention_mask=attention_mask, decoder_input=inputs_embeds, labels=labels, - loss_mask=loss_mask, inference_context=inference_context, runtime_gather_output=runtime_gather_output, packed_seq_params=packed_seq_params, - extra_block_kwargs=extra_block_kwargs, inference_params=inference_params, + loss_mask=loss_mask, + padding_mask=padding_mask, ) if loss_mask is not None: return output, loss_mask @@ -682,5 +703,5 @@ def forward( "MuseGlimmerVisionModel", "MuseGlimmerWeightlessRMSNorm", "customize_muse_glimmer_language_model", - "get_muse_glimmer_layer_spec", + "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 index 531372c292..e4b77b963f 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -177,7 +177,7 @@ def _validate_architecture(hf_config: PretrainedConfig) -> tuple[Any, Any]: return text_config, vision_config def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerModelConfig: - """Translate HF config directly into the new model-config/builder path.""" + """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) @@ -246,10 +246,10 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerM seq_length=int(text_config.max_position_embeddings), parallel_output=True, share_embeddings_and_output_weights=bool(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, - scatter_embedding_sequence_parallel=False, vision=vision, image_token_id=int(hf_config.image_token_id), video_token_id=int(hf_config.video_token_id), @@ -344,50 +344,50 @@ def mapping_registry(self) -> MegatronMappingRegistry: text_prefix = "model.language_model" mappings: list[MegatronParamMapping[Any]] = [ AutoMapping( - megatron_param="language_model.embedding.word_embeddings.weight", + megatron_param="embedding.word_embeddings.weight", hf_param=f"{text_prefix}.embed_tokens.weight", ), AutoMapping( - megatron_param="language_model.output_layer.weight", + megatron_param="output_layer.weight", hf_param="lm_head.weight", ), ReplicatedMapping( - megatron_param="language_model.decoder.final_layernorm.weight", + megatron_param="decoder.final_norm.weight", hf_param=f"{text_prefix}.norm.weight", ), ReplicatedMapping( - megatron_param="language_model.decoder.layers.*.input_layernorm.weight", + megatron_param="decoder.layers.*.input_layernorm.weight", hf_param=f"{text_prefix}.layers.*.input_layernorm.weight", ), ReplicatedMapping( - megatron_param="language_model.decoder.layers.*.self_attention.post_layernorm.weight", + megatron_param="decoder.layers.*.self_attention.post_layernorm.weight", hf_param=f"{text_prefix}.layers.*.post_attention_layernorm.weight", ), ReplicatedMapping( - megatron_param="language_model.decoder.layers.*.pre_mlp_layernorm.weight", + megatron_param="decoder.layers.*.pre_mlp_layernorm.weight", hf_param=f"{text_prefix}.layers.*.pre_feedforward_layernorm.weight", ), ReplicatedMapping( - megatron_param="language_model.decoder.layers.*.mlp.post_layernorm.weight", + megatron_param="decoder.layers.*.mlp.post_layernorm.weight", hf_param=f"{text_prefix}.layers.*.post_feedforward_layernorm.weight", ), AutoMapping( - megatron_param="language_model.decoder.layers.*.self_attention.linear_proj.weight", + megatron_param="decoder.layers.*.self_attention.linear_proj.weight", hf_param=f"{text_prefix}.layers.*.self_attn.o_proj.weight", ), AutoMapping( - megatron_param="language_model.decoder.layers.*.mlp.linear_fc2.weight", + megatron_param="decoder.layers.*.mlp.linear_fc2.weight", hf_param=f"{text_prefix}.layers.*.mlp.down_proj.weight", ), MuseGlimmerQKVGMapping( - "language_model.decoder.layers.*.self_attention.linear_qkv.weight", + "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="language_model.decoder.layers.*.mlp.linear_fc1.weight", + 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", ), diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py index 8c14c1c604..57bd374039 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py @@ -16,21 +16,20 @@ from __future__ import annotations -from dataclasses import replace - +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.gpt import GPTModelBuilder +from megatron.training.models.hybrid import HybridModelBuilder +from megatron.training.vocab_utils import calculate_padded_vocab_size from megatron.bridge.models.muse_glimmer.model_config import MuseGlimmerModelConfig from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( MuseGlimmerModel, - customize_muse_glimmer_language_model, - get_muse_glimmer_layer_spec, + get_muse_glimmer_hybrid_stack_spec, ) -class MuseGlimmerModelBuilder(GPTModelBuilder): - """Construct the replicated vision stack around a builder-backed MCore decoder.""" +class MuseGlimmerModelBuilder(HybridModelBuilder): + """Construct the complete Muse Glimmer model on MCore HybridModel.""" def __init__(self, model_config: MuseGlimmerModelConfig) -> None: super().__init__(model_config) @@ -57,24 +56,28 @@ def build_model( if not isinstance(model_config, MuseGlimmerModelConfig): raise TypeError(f"Expected MuseGlimmerModelConfig, got {type(model_config).__name__}.") - language_config = replace( + 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, - transformer_layer_spec=get_muse_glimmer_layer_spec(model_config.transformer), - ) - language_model = GPTModelBuilder(language_config).build_model( + 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, ) - customize_muse_glimmer_language_model(language_model) - return MuseGlimmerModel( - model_config, - language_model, - pre_process=language_model.pre_process, - post_process=language_model.post_process, - vp_stage=vp_stage, - ) __all__ = ["MuseGlimmerModelBuilder"] diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 0759f99b3a..6d403216d0 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -5,10 +5,14 @@ from __future__ import annotations import math +from collections.abc import Iterator from unittest.mock import patch import pytest import torch +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.training.models.hybrid import HybridModelConfig from megatron.bridge import AutoBridge from megatron.bridge.models.conversion.param_mapping import merge_qkv_weights, split_qkv_weights @@ -94,6 +98,7 @@ def test_config_conversion_uses_model_config_path_only() -> None: 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.num_layers == 4 assert model_config.hidden_size == 64 @@ -106,7 +111,7 @@ def test_config_conversion_uses_model_config_path_only() -> None: 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.scatter_embedding_sequence_parallel is False + assert model_config.hybrid_layer_pattern == "****" assert model_config.special_token_ids == {"images": 120, "videos": 121} @@ -181,7 +186,7 @@ def test_centered_and_final_norms_use_their_native_expressions() -> None: def test_mapping_registry_covers_complete_checkpoint() -> None: registry = MuseGlimmerBridge().mapping_registry() - qkvg = registry.megatron_to_hf_lookup("language_model.decoder.layers.2.self_attention.linear_qkv.weight") + 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 ( @@ -202,7 +207,8 @@ def test_tiny_vision_model_preserves_expected_token_count() -> None: assert output.shape == (4, 64) -def test_builder_constructs_combined_model_on_cpu() -> None: +@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 @@ -222,69 +228,10 @@ def test_builder_constructs_combined_model_on_cpu() -> None: wrap_with_ddp=False, mixed_precision_wrapper=None, ) + assert len(models) == 1 + model = models[0] try: - assert len(models) == 1 - model = models[0] - assert isinstance(model, MuseGlimmerModel) - names = dict(model.named_parameters()) - assert "vision_tower.patch_embedder.patch_embedding.weight" in names - assert "language_model.decoder.layers.0.self_attention.post_layernorm.weight" in names - assert "language_model.decoder.layers.0.mlp.post_layernorm.weight" in names - assert isinstance(model.language_model.decoder.final_layernorm, MuseGlimmerRMSNorm) - 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(model.language_model, "forward", return_value=expected_output) as language_forward, - ): - output = model(input_ids=input_ids, inputs_embeds=inputs_embeds) - - assert output is expected_output - torch.testing.assert_close(language_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.language_model.embedding, "forward", return_value=embedding_output - ) as embedding_forward, - patch.object(model.language_model, "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) + yield model finally: from megatron.core import parallel_state @@ -292,3 +239,93 @@ def test_builder_constructs_combined_model_on_cpu() -> None: 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.final_norm, MuseGlimmerRMSNorm) + 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_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) From 6cef8df6d2f18d28d9c43b9cecbda5589a70dbe8 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 13:56:11 -0700 Subject: [PATCH 04/72] docs(model): pin Muse Glimmer HybridModel verification commit Signed-off-by: yaoyu-33 --- examples/model_verification_cards/muse-glimmer-30b/card.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index 782ce5c4a6..a931e0ea78 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -31,7 +31,7 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: 2e5efbab0e2039446693f61e2fc99d0233832775 # pragma: allowlist secret + bridge_commit: 95bc2220f8d721744388a1e6dc8b32df34ec421c # pragma: allowlist secret items: hf_to_megatron_cpu: From 01ba98337a5f218cd7933de96b3ad8c7d31b2caf Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 14:56:11 -0700 Subject: [PATCH 05/72] fix(conversion): support builder-backed model checkpoints Signed-off-by: yaoyu-33 --- scripts/conversion/gpu_backend.py | 88 +++++++++++++++---- .../bridge/models/conversion/auto_bridge.py | 16 +++- .../bridge/training/model_load_save.py | 18 ++-- .../conversion/launcher/test_gpu_backend.py | 55 ++++++++++++ tests/unit_tests/models/test_auto_bridge.py | 40 +++++++++ .../training/test_model_load_save.py | 38 ++++++++ 6 files changed, 227 insertions(+), 28 deletions(-) diff --git a/scripts/conversion/gpu_backend.py b/scripts/conversion/gpu_backend.py index b199e96251..a012f0f8c9 100644 --- a/scripts/conversion/gpu_backend.py +++ b/scripts/conversion/gpu_backend.py @@ -167,6 +167,30 @@ def _configure_model_provider( model_provider.params_dtype = dtype +def _uses_model_builder(bridge: AutoBridge) -> bool: + """Return whether the selected bridge supports native builder construction.""" + return getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None + + +def _configure_model_config( + model_config, + *, + tp: int, + pp: int, + ep: int, + etp: int, + dtype: torch.dtype, +) -> 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 + + 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] = {} @@ -304,12 +328,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) - 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) + megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False) bridge.save_megatron_model( megatron_model, @@ -383,12 +417,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) - _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) + 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_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) + _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) model_parallel_overrides: dict[str, object] = { "tensor_model_parallel_size": tp, @@ -457,11 +499,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) - 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) + 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/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index b6f89a16af..b61102d567 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -1467,8 +1467,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.MODEL_CONFIG_CLASS is not None: + 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 = {} @@ -1843,6 +1853,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/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index a077b8a520..d24f1640cd 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/src/megatron/bridge/training/model_load_save.py @@ -571,16 +571,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]) - - # Validate that the model config is a model provider - if not isinstance(model_config, ModelProviderMixin): + # 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]) + + 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 diff --git a/tests/unit_tests/conversion/launcher/test_gpu_backend.py b/tests/unit_tests/conversion/launcher/test_gpu_backend.py index 905890e749..d1b6423b17 100644 --- a/tests/unit_tests/conversion/launcher/test_gpu_backend.py +++ b/tests/unit_tests/conversion/launcher/test_gpu_backend.py @@ -114,6 +114,8 @@ def provide_distributed_model(self, *args, **kwargs): class _FakeModelBridge: + MODEL_CONFIG_CLASS = None + def get_hf_tokenizer_kwargs(self): return {"padding_side": "left"} @@ -182,6 +184,59 @@ def fake_from_hf_pretrained(*args, **kwargs): assert from_hf_call[2]["revision"] == "0123456789abcdef" # pragma: allowlist secret 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 + + 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/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index 2bb402f54a..d4bb57c6da 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -711,6 +711,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 @@ -1858,6 +1859,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.MODEL_CONFIG_CLASS = None mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model @@ -1889,6 +1891,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.MODEL_CONFIG_CLASS = None mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model @@ -1933,6 +1936,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.MODEL_CONFIG_CLASS = None mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model mock_bridge.save_megatron_model = Mock() @@ -1961,6 +1965,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.MODEL_CONFIG_CLASS = object + 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") @@ -1972,6 +2010,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.MODEL_CONFIG_CLASS = None mock_bridge.to_megatron_model.return_value = [Mock()] mock_bridge.save_megatron_model = Mock() mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {} @@ -1995,6 +2034,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.MODEL_CONFIG_CLASS = None 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/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index 1ce8bb4eb7..9e199dcb67 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 @@ -911,6 +912,43 @@ def finalize(self) -> 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) + 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 + mock_save_checkpoint.assert_called_once_with( + state=mock_state, + model=[mock_model], + optimizer=None, + opt_param_scheduler=None, + num_floating_point_operations_so_far=0, + callback_manager=None, + ) + @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") From 841867344d6649fcf65374206f55ce5ff422c89c Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 15:07:58 -0700 Subject: [PATCH 06/72] fix(inference): load builder-backed VLM checkpoints Signed-off-by: yaoyu-33 --- .../compare_hf_and_megatron/compare.py | 57 ++++++++---- scripts/inference/vlm_generation.py | 88 ++++++++++++++----- .../unit_tests/test_compare_mask_handling.py | 44 ++++++++++ 3 files changed, 153 insertions(+), 36 deletions(-) diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 2e2efe6235..9d72542355 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -694,14 +694,25 @@ 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, "MODEL_CONFIG_CLASS", None) is not None: + model_config = bridge.get_model_config() + model_config.tensor_model_parallel_size = tp + model_config.pipeline_model_parallel_size = pp + model_config.expert_model_parallel_size = ep + model_config.expert_tensor_parallel_size = etp + model_config.pipeline_dtype = torch.bfloat16 + model_config.params_dtype = torch.bfloat16 + 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.finalize() + model_provider.initialize_model_parallel(seed=0) megatron_model = bridge.load_megatron_model( args.megatron_model_path, mp_overrides={ @@ -722,14 +733,28 @@ 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, "MODEL_CONFIG_CLASS", None) is not None: + model_config = bridge.get_model_config() + model_config.tensor_model_parallel_size = tp + model_config.pipeline_model_parallel_size = pp + model_config.expert_model_parallel_size = ep + model_config.expert_tensor_parallel_size = etp + model_config.pipeline_dtype = torch.bfloat16 + model_config.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/scripts/inference/vlm_generation.py b/scripts/inference/vlm_generation.py index 1a48e91d22..51d8d9ace5 100644 --- a/scripts/inference/vlm_generation.py +++ b/scripts/inference/vlm_generation.py @@ -148,6 +148,34 @@ 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, "MODEL_CONFIG_CLASS", None) is not None + + +def _configure_builder_model( + bridge: AutoBridge, + *, + tp: int, + pp: int, + ep: int, + etp: int, + pp_layout: str | None = None, +): + """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 + 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.""" maybe_initialize_distributed() @@ -187,17 +215,29 @@ 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, + ) + 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.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, @@ -215,15 +255,23 @@ 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) + 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() + 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 diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index 7a34025aaa..1711de9302 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -210,6 +210,50 @@ 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(MODEL_CONFIG_CLASS=object) + 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 model_config.tensor_model_parallel_size == 8 + assert model_config.params_dtype is torch.bfloat16 + 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_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) From 37df4b0bec9f5fff6240818baa55c5a7c0aab1e9 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 15:13:14 -0700 Subject: [PATCH 07/72] fix(checkpoint): identify Muse TP norm replicas Signed-off-by: yaoyu-33 --- .../muse_glimmer/modeling_muse_glimmer.py | 18 ++++++++++++ .../models/muse_glimmer/test_muse_glimmer.py | 28 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 13fc82726c..1079ac895c 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -22,6 +22,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from megatron.core import parallel_state from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.backends import LocalSpecProvider from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules @@ -34,6 +35,7 @@ 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, make_sharded_tensors_for_checkpoint from torch import Tensor from megatron.bridge.models.gemma.modules import extend_instance @@ -79,6 +81,22 @@ def __init__(self, config: TransformerConfig, hidden_size: int, eps: float = 1e- def forward(self, hidden_states: Tensor) -> Tensor: return _centered_rms_norm(hidden_states, self.eps) * (1.0 + self.weight.float()).to(hidden_states.dtype) + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: tuple[tuple[int, int, int], ...] = (), + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Represent the replicated weight with a distinct writer on each TP rank.""" + metadata = ensure_metadata_has_dp_cp_group(metadata) + return make_sharded_tensors_for_checkpoint( + self.state_dict(keep_vars=True), + prefix, + sharded_offsets=sharded_offsets, + tp_group=parallel_state.get_tensor_model_parallel_group(), + dp_cp_group=metadata["dp_cp_group"], + ) + class MuseGlimmerRMSNorm(nn.Module): """Learnable RMSNorm used after the final Muse decoder layer.""" diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 6d403216d0..0203ee3cf1 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -183,6 +183,34 @@ def test_centered_and_final_norms_use_their_native_expressions() -> None: torch.testing.assert_close(final(hidden_states), expected_final) +def test_centered_norm_sharded_state_dict_identifies_tensor_parallel_replicas() -> None: + config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer + norm = MuseGlimmerCenteredRMSNorm(config, hidden_size=4) + tp_group = object() + dp_cp_group = object() + + with ( + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.parallel_state.get_tensor_model_parallel_group", + return_value=tp_group, + ), + patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.make_sharded_tensors_for_checkpoint", + return_value={"norm.weight": object()}, + ) as make_sharded, + ): + result = norm.sharded_state_dict(prefix="norm.", metadata={"dp_cp_group": dp_cp_group}) + + assert set(result) == {"norm.weight"} + make_sharded.assert_called_once_with( + norm.state_dict(keep_vars=True), + "norm.", + sharded_offsets=(), + tp_group=tp_group, + dp_cp_group=dp_cp_group, + ) + + def test_mapping_registry_covers_complete_checkpoint() -> None: registry = MuseGlimmerBridge().mapping_registry() From 0bbd420fada019562647a17fa420ec981c6f81e3 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 15:29:36 -0700 Subject: [PATCH 08/72] fix(checkpoint): restore persisted builder values Signed-off-by: yaoyu-33 --- src/megatron/bridge/models/common/__init__.py | 2 + src/megatron/bridge/models/common/base.py | 79 +++++++++++-------- .../models/muse_glimmer/model_config.py | 4 +- .../bridge/training/model_load_save.py | 12 ++- .../models/muse_glimmer/test_muse_glimmer.py | 19 +++++ 5 files changed, 80 insertions(+), 36 deletions(-) diff --git a/src/megatron/bridge/models/common/__init__.py b/src/megatron/bridge/models/common/__init__.py index 03e7396120..63525c25d2 100644 --- a/src/megatron/bridge/models/common/__init__.py +++ b/src/megatron/bridge/models/common/__init__.py @@ -19,6 +19,7 @@ ModelT, Serializable, compose_hooks, + deserialize_model_config, ) from megatron.bridge.models.common.unimodal import build_virtual_pipeline_stages, unimodal_build_distributed_models @@ -30,6 +31,7 @@ "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..19e3c42454 100644 --- a/src/megatron/bridge/models/common/base.py +++ b/src/megatron/bridge/models/common/base.py @@ -27,7 +27,50 @@ ModelConfig as _MegatronModelConfig, ) -from megatron.bridge.utils.instantiate_utils import _resolve_target, _validate_target_prefix +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 ModelConfig(_MegatronModelConfig): @@ -43,35 +86,7 @@ def get_builder_cls(self) -> type: @classmethod def from_dict(cls, data: dict[str, Any]) -> "ModelConfig": """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 - + result = deserialize_model_config(data) + if not isinstance(result, cls): + raise ValueError(f"Cannot deserialize: expected {cls.__name__}, got {type(result).__name__}") return result diff --git a/src/megatron/bridge/models/muse_glimmer/model_config.py b/src/megatron/bridge/models/muse_glimmer/model_config.py index c9aafebe2b..18f5554e6c 100644 --- a/src/megatron/bridge/models/muse_glimmer/model_config.py +++ b/src/megatron/bridge/models/muse_glimmer/model_config.py @@ -21,7 +21,7 @@ from megatron.training.models.hybrid import HybridModelConfig -from megatron.bridge.models.common import ModelConfig +from megatron.bridge.models.common import 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 @@ -118,7 +118,7 @@ def from_dict(cls, data: dict[str, Any]) -> "MuseGlimmerModelConfig": restored_transformer["activation_func"] = str_to_callable(activation_name) restored_data["transformer"] = restored_transformer - result = ModelConfig.from_dict(restored_data) + result = deserialize_model_config(restored_data) if not isinstance(result, cls): raise TypeError(f"Expected {cls.__name__}, got {type(result).__name__}.") return result diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index d24f1640cd..ffa58e034f 100644 --- a/src/megatron/bridge/training/model_load_save.py +++ b/src/megatron/bridge/training/model_load_save.py @@ -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: diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 0203ee3cf1..75fd55de41 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -130,6 +130,25 @@ def test_autobridge_selects_string_registration_and_serializes_config() -> None: 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()) From a34fb67d9f893ff9bcb75f9fa5ca1af7a88bbc11 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 15:43:49 -0700 Subject: [PATCH 09/72] fix(checkpoint): skip disabled RNG capture Signed-off-by: yaoyu-33 --- src/megatron/bridge/training/checkpointing.py | 14 ++++++++------ tests/unit_tests/training/test_checkpointing.py | 8 ++++++++ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 46439a9a46..6598e3f125 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -1166,12 +1166,14 @@ def save_checkpoint( # Collect rng state across data parallel ranks. if pg_collection is None: pg_collection = get_pg_collection(model) - rng_state = get_rng_state( - data_parallel_random_init=cfg.rng.data_parallel_random_init, - ckpt_format=ckpt_cfg.ckpt_format, - pg_collection=pg_collection, - module_name=module_name, - ) + rng_state = None + if ckpt_cfg.save_rng: + rng_state = get_rng_state( + data_parallel_random_init=cfg.rng.data_parallel_random_init, + ckpt_format=ckpt_cfg.ckpt_format, + pg_collection=pg_collection, + module_name=module_name, + ) # Collect rerun state across all ranks rerun_state_machine = get_rerun_state_machine() diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index a0e2290e55..aee3315ca6 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -617,6 +617,7 @@ def save_checkpoint_fixtures(): class TestSaveCheckpoint: """Test checkpoint saving functionality.""" + @pytest.mark.parametrize("save_rng", [True, False]) @patch("megatron.bridge.training.checkpointing.wandb_utils") @patch("megatron.bridge.training.checkpointing.is_last_rank") @patch("builtins.open", new_callable=mock_open) @@ -664,6 +665,7 @@ def test_save_checkpoint_global( mock_is_last_rank, mock_wandb, save_checkpoint_fixtures, + save_rng, ): """Test saving a global checkpoint.""" # Setup mocks @@ -695,6 +697,7 @@ def test_save_checkpoint_global( # Add wandb logger to state save_checkpoint_fixtures["mock_state"].wandb_logger = Mock() save_checkpoint_fixtures["mock_state"].cfg.checkpoint.most_recent_k = -1 + save_checkpoint_fixtures["mock_state"].cfg.checkpoint.save_rng = save_rng # Call save_checkpoint save_checkpoint( @@ -711,6 +714,11 @@ def test_save_checkpoint_global( mock_ft.on_checkpointing_start.assert_called_once() mock_gen_state.assert_called_once() mock_dist_ckpt.save.assert_called_once() + if save_rng: + 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 = [ From b031e3400afc933720bfad3dbd6fca90b90372fa Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 15:53:48 -0700 Subject: [PATCH 10/72] fix(conversion): save CPU checkpoints without CUDA Signed-off-by: yaoyu-33 --- src/megatron/bridge/training/checkpointing.py | 25 ++++++++++++++ .../bridge/training/model_load_save.py | 15 ++++++++- .../unit_tests/training/test_checkpointing.py | 33 +++++++++++++++++++ .../training/test_model_load_save.py | 17 ++++++---- 4 files changed, 82 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 6598e3f125..4566073053 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -16,6 +16,7 @@ import contextlib import gc +import inspect import os import random import shutil @@ -24,6 +25,7 @@ from abc import ABC 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 @@ -159,6 +161,29 @@ DATALOADER_STATE_SUBDIR = "energon" +class _CpuTorchDistSaveShardedStrategy(TorchDistSaveShardedStrategy): + """Run MCore's synchronous torch-dist writer without a CUDA staging barrier.""" + + 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.execute_sync() + + # ============================================================================ # Checkpoint version and utilities # ============================================================================ diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py index ffa58e034f..3cd3e3de68 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 @@ -615,6 +615,17 @@ def save_megatron_model( dist=None, ) + 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 @@ -777,6 +788,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: @@ -787,6 +799,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/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index aee3315ca6..bd196ae0a3 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -16,12 +16,15 @@ import os import pickle import tempfile +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 ( @@ -34,6 +37,7 @@ 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, @@ -736,6 +740,35 @@ 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 = [] + + 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=[], + preload_fn=partial(preload, ["cpu-tensor"], True), + ) + strategy = _CpuTorchDistSaveShardedStrategy() + + 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"] + 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 9e199dcb67..8b5fe31374 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/tests/unit_tests/training/test_model_load_save.py @@ -909,6 +909,7 @@ def finalize(self) -> None: optimizer=None, opt_param_scheduler=None, num_floating_point_operations_so_far=0, + checkpointing_context=None, callback_manager=None, ) @@ -932,6 +933,7 @@ def test_save_megatron_model_accepts_builder_config( """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 @@ -940,13 +942,12 @@ def test_save_megatron_model_accepts_builder_config( mock_get_model_config.assert_not_called() assert mock_config_container.call_args.kwargs["model"] is model_config - mock_save_checkpoint.assert_called_once_with( - state=mock_state, - model=[mock_model], - optimizer=None, - opt_param_scheduler=None, - num_floating_point_operations_so_far=0, - callback_manager=None, + 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") @@ -1027,6 +1028,7 @@ def finalize(self) -> None: optimizer=None, opt_param_scheduler=None, num_floating_point_operations_so_far=0, + checkpointing_context=None, callback_manager=None, ) @@ -1097,6 +1099,7 @@ def finalize(self) -> None: optimizer=None, opt_param_scheduler=None, num_floating_point_operations_so_far=0, + checkpointing_context=None, callback_manager=None, ) From e322510641fd910a48ef6f6dc143eac9573366e5 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 16:00:22 -0700 Subject: [PATCH 11/72] fix(checkpoint): finalize CPU saves on Gloo Signed-off-by: yaoyu-33 --- src/megatron/bridge/training/checkpointing.py | 13 +++++++++++++ tests/unit_tests/training/test_checkpointing.py | 6 +++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 4566073053..d2bf088a8a 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -164,6 +164,16 @@ class _CpuTorchDistSaveShardedStrategy(TorchDistSaveShardedStrategy): """Run MCore's synchronous torch-dist writer without a CUDA staging barrier.""" + @staticmethod + def _run_cpu_finalize(finalize_fn: Callable[[], None]) -> None: + """Make MCore's failure-status tensor use the CPU for a Gloo save.""" + 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") @@ -181,6 +191,9 @@ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> No 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_cpu_finalize, finalize_fn) for finalize_fn in async_request.finalize_fns], + ) async_request.execute_sync() diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index bd196ae0a3..307b7999e6 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -744,6 +744,7 @@ 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) @@ -755,10 +756,11 @@ def write(_rank, write_buckets, _results_queue): request = AsyncRequest( async_fn=write, async_fn_args=(0, None, None), - finalize_fns=[], + 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), @@ -768,6 +770,8 @@ def write(_rank, write_buckets, _results_queue): 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_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.""" From 8b096c99c1a4556a40951d54138c22119270197c Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 16:32:22 -0700 Subject: [PATCH 12/72] docs(model): verify Muse Glimmer model workflows Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 149 +++++++++++++----- 1 file changed, 106 insertions(+), 43 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index a931e0ea78..716dbac334 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -10,11 +10,12 @@ summary: > 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. - Conversion, numerical parity, inference, and training remain unverified until - the corresponding pinned public commands complete their declared gates. + CPU and TP8 GPU conversion, exact round-trip tensor equality, numerical + forward parity, and deterministic image-conditioned inference are verified. + Training remains unverified. verification_index: model_level: - unverified: + verified: - hf_to_megatron_cpu - hf_to_megatron_gpu - megatron_to_hf_cpu @@ -31,70 +32,132 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: 95bc2220f8d721744388a1e6dc8b32df34ec421c # pragma: allowlist secret + bridge_commit: e322510641fd910a48ef6f6dc143eac9573366e5 # pragma: allowlist secret items: hf_to_megatron_cpu: - status: unverified + status: verified precision: bf16 - command: null - last_verified: null + 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-11 expected_result: > - Import the complete pinned checkpoint through AutoBridge.get_model(), - persist a reloadable Megatron checkpoint, and audit every language, - vision, adapter, and projection tensor for key, shape, dtype, and value - agreement. + 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: unverified + status: verified precision: bf16 - command: null - last_verified: null + bridge_commit: 37df4b0bec9f5fff6240818baa55c5a7c0aab1e9 # pragma: allowlist secret + 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-11 expected_result: > - Import the complete pinned checkpoint on GPUs through the builder-backed - path, persist and reload it, and report an exact audit of all mapped - tensors with no missing or unexpected parameters. + 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: unverified + status: verified precision: bf16 - command: null - last_verified: null + 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-11 expected_result: > - Export a complete CPU-imported Megatron checkpoint to Hugging Face, - strictly reload it with the minimum supported Transformers version, and - compare all source and exported tensors under a declared exact criterion. + 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: unverified + status: verified precision: bf16 - command: null - last_verified: null + bridge_commit: 0bbd420fada019562647a17fa420ec981c6f81e3 # pragma: allowlist secret + 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-11 expected_result: > - Export a complete GPU-imported Megatron checkpoint, strictly reload the - resulting Hugging Face model, and match all language, vision, adapter, - and projection keys, shapes, dtypes, and values to the pinned source. + 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: unverified + status: verified precision: bf16 - command: null - last_verified: null + 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-11 expected_result: > - Process the same pinned image and prompt through native Hugging Face and - Megatron implementations, match the next token, achieve logit cosine - similarity of at least 0.99, and report maximum and mean absolute logit - differences as non-gating diagnostics. + 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.999929, above the 0.99 gate; maximum and + mean absolute logit differences are 0.687500 and 0.080432, respectively. inference: - status: unverified + status: verified precision: bf16 - command: null - last_verified: null - expected_result: > - Reload the persisted complete Megatron checkpoint, process a pinned - public image and prompt, and complete deterministic bounded generation - with finite logits and a semantically correct response. + 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 --tp 8 --pp 1 --ep 1 --etp 1 + last_verified: 2026-08-11 + 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: From 220621a2032006008ef6d2912f7999cdb703045f Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 21:59:27 -0700 Subject: [PATCH 13/72] feat(recipe): add Muse Glimmer training verification configs Signed-off-by: yaoyu-33 --- scripts/training/recipe_metadata.py | 1 + .../models/muse_glimmer/model_config.py | 15 ++ src/megatron/bridge/recipes/__init__.py | 2 + .../bridge/recipes/muse_glimmer/__init__.py | 18 ++ .../recipes/muse_glimmer/h100/__init__.py | 30 ++++ .../recipes/muse_glimmer/h100/muse_glimmer.py | 170 ++++++++++++++++++ .../models/muse_glimmer/test_muse_glimmer.py | 14 ++ .../recipes/test_muse_glimmer_recipes.py | 86 +++++++++ .../scripts/training/test_recipe_metadata.py | 1 + 9 files changed, 337 insertions(+) create mode 100644 src/megatron/bridge/recipes/muse_glimmer/__init__.py create mode 100644 src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py create mode 100644 src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py create mode 100644 tests/unit_tests/recipes/test_muse_glimmer_recipes.py diff --git a/scripts/training/recipe_metadata.py b/scripts/training/recipe_metadata.py index bc73425c27..dfec57ed8b 100644 --- a/scripts/training/recipe_metadata.py +++ b/scripts/training/recipe_metadata.py @@ -94,6 +94,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/src/megatron/bridge/models/muse_glimmer/model_config.py b/src/megatron/bridge/models/muse_glimmer/model_config.py index 18f5554e6c..348a7b0ae6 100644 --- a/src/megatron/bridge/models/muse_glimmer/model_config.py +++ b/src/megatron/bridge/models/muse_glimmer/model_config.py @@ -55,6 +55,16 @@ class MuseGlimmerVisionModelConfig: 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(HybridModelConfig): @@ -75,6 +85,11 @@ class MuseGlimmerModelConfig(HybridModelConfig): freeze_vision_model: bool = False freeze_vision_projection: 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.""" 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..402701b235 --- /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_128gpu_h100_bf16_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_128gpu_h100_bf16_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..873922d70d --- /dev/null +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -0,0 +1,170 @@ +# 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 + +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 _apply_model_and_data(cfg: ConfigContainer, *, seq_length: int, context_parallel_size: int) -> None: + """Apply the exact model, multimodal data, and common execution settings.""" + cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.model.seq_length = seq_length + cfg.model.tensor_model_parallel_size = 4 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_dtype = None + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = context_parallel_size + 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.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 = seq_length + cfg.dataset.hf_processor_path = _MODEL_ID + cfg.dataset.hf_processor_kwargs = {"revision": _MODEL_REVISION} + 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.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.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + + +def _set_optimizer(cfg: ConfigContainer, *, max_lr: float, warmup_iters: int) -> None: + """Set the bounded cosine schedule shared by verification workloads.""" + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=warmup_iters, + max_lr=max_lr, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + + +def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: + """Return a 100-step random-init multimodal pretraining config on 128 H100 GPUs.""" + cfg = _sft_common_vlm() + _apply_model_and_data(cfg, seq_length=4096, context_parallel_size=1) + cfg.dataset.pad_to_max_length = True + cfg.dataset.enable_in_batch_packing = False + cfg.rng.seed = 1234 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 + _set_optimizer(cfg, max_lr=3e-4, warmup_iters=40) + cfg.optimizer.use_precision_aware_optimizer = True + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + 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() + _apply_model_and_data(cfg, seq_length=4096, context_parallel_size=1) + cfg.dataset.pad_to_max_length = True + cfg.dataset.enable_in_batch_packing = False + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 8 + _set_optimizer(cfg, max_lr=5e-6, warmup_iters=10) + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.save_optim = False + cfg.checkpoint.save_rng = False + return cfg + + +def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContainer: + """Return a 100-step 8K full SFT config with CP4 on 32 H100 GPUs.""" + cfg = _sft_common_vlm() + _apply_model_and_data(cfg, seq_length=8192, context_parallel_size=4) + cfg.dataset.pad_to_max_length = True + cfg.dataset.enable_in_batch_packing = False + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 8 + _set_optimizer(cfg, max_lr=5e-6, warmup_iters=10) + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.save_optim = False + cfg.checkpoint.save_rng = False + 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() + _apply_model_and_data(cfg, seq_length=8192, context_parallel_size=1) + 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.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 8 + _set_optimizer(cfg, max_lr=1e-4, warmup_iters=10) + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + return cfg + + +__all__ = [ + "muse_glimmer_30b_peft_8gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_128gpu_h100_bf16_config", + "muse_glimmer_30b_sft_32gpu_h100_bf16_config", + "muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config", +] diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 75fd55de41..b10dfcec0b 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -6,6 +6,7 @@ import math from collections.abc import Iterator +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -33,6 +34,7 @@ MuseGlimmerBridge, MuseGlimmerQKVGMapping, ) +from megatron.bridge.training.utils.flop_utils import vit_flops_from_grid_thw pytestmark = pytest.mark.unit @@ -113,6 +115,18 @@ def test_config_conversion_uses_model_config_path_only() -> None: 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_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_autobridge_selects_string_registration_and_serializes_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..0944585e06 --- /dev/null +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Focused tests for Muse Glimmer verification recipes.""" + +from __future__ import annotations + +import pytest + +from megatron.bridge.peft.lora import LoRA +from megatron.bridge.recipes.muse_glimmer.h100 import ( + muse_glimmer_30b_peft_8gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_128gpu_h100_bf16_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 +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", "context_parallel_size", "global_batch_size"), + [ + (muse_glimmer_30b_pretrain_128gpu_h100_bf16_config, 4096, 1, 1024), + (muse_glimmer_30b_sft_32gpu_h100_bf16_config, 4096, 1, 8), + (muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, 8192, 4, 8), + (muse_glimmer_30b_peft_8gpu_h100_bf16_config, 8192, 1, 8), + ], +) +def test_muse_glimmer_recipe_contracts(recipe, seq_length, context_parallel_size, global_batch_size) -> None: + cfg = recipe() + + assert isinstance(cfg, ConfigContainer) + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.context_parallel_size == context_parallel_size + 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.hf_processor_kwargs == { + "revision": "f84ecc3a0ea984a4c04542a84269e3d065350a6e" # pragma: allowlist secret + } + assert cfg.dataset.source.load_kwargs == { + "revision": "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret + } + assert cfg.dataset.pad_to_max_length is True + assert cfg.dataset.enable_in_batch_packing is False + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == global_batch_size + assert cfg.train.micro_batch_size == 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_128gpu_h100_bf16_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.checkpoint.save_interval == 50 + assert cfg.checkpoint.load is None + + +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) 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"), From ec67e110e7b9ba8a4e80b5a6137b63709caa0bf2 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 22:06:27 -0700 Subject: [PATCH 14/72] fix(recipe): keep Muse training on HybridModel config Signed-off-by: yaoyu-33 --- .../recipes/muse_glimmer/h100/muse_glimmer.py | 4 +-- .../recipes/test_muse_glimmer_recipes.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 873922d70d..c0c0a5bc72 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -33,9 +33,7 @@ def _apply_model_and_data(cfg: ConfigContainer, *, seq_length: int, context_parallel_size: int) -> None: """Apply the exact model, multimodal data, and common execution settings.""" - cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).to_megatron_provider( - load_weights=False - ) + cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).get_model_config() cfg.model.seq_length = seq_length cfg.model.tensor_model_parallel_size = 4 cfg.model.pipeline_model_parallel_size = 1 diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 0944585e06..561cecccdc 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -4,9 +4,12 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest from megatron.bridge.peft.lora import LoRA +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_128gpu_h100_bf16_config, @@ -84,3 +87,25 @@ def test_muse_glimmer_lora_targets_native_attention_projections() -> None: 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_128gpu_h100_bf16_config() + + assert cfg.model is not None From 22ddbf0bdcc44ced258cad6b1a478f1b2f6d90b1 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 22:26:15 -0700 Subject: [PATCH 15/72] feat(model): add Muse Glimmer training collator Signed-off-by: yaoyu-33 --- .../bridge/data/collators/registry.py | 3 + .../models/muse_glimmer/data/__init__.py | 20 +++ .../models/muse_glimmer/data/collate_fn.py | 114 ++++++++++++++++++ .../models/muse_glimmer/test_collate.py | 65 ++++++++++ 4 files changed, 202 insertions(+) create mode 100644 src/megatron/bridge/models/muse_glimmer/data/__init__.py create mode 100644 src/megatron/bridge/models/muse_glimmer/data/collate_fn.py create mode 100644 tests/unit_tests/models/muse_glimmer/test_collate.py diff --git a/src/megatron/bridge/data/collators/registry.py b/src/megatron/bridge/data/collators/registry.py index 6fd485b555..3b07d53f46 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" + ), } 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..ec5d558f22 --- /dev/null +++ b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py @@ -0,0 +1,114 @@ +# 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, + shared_chat_template_kwargs_from_examples, +) +from megatron.bridge.data.datasets.utils import IGNORE_INDEX +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.""" + if enable_in_batch_packing: + raise ValueError("Muse Glimmer direct-HF training does not support in-batch packing.") + del in_batch_packing_pad_to_multiple_of + + 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, + ) + 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/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..6854bf15af --- /dev/null +++ b/tests/unit_tests/models/muse_glimmer/test_collate.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from types import SimpleNamespace + +import pytest +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"] is True + assert self.tokenizer.padding_side == "right" + batch_size = len(conversations) + return { + "input_ids": torch.tensor([[11, 12, 13, 14]] * batch_size), + "attention_mask": torch.ones(batch_size, 4, 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_rejects_in_batch_packing(): + with pytest.raises(ValueError, match="does not support in-batch packing"): + muse_glimmer_collate_fn([], SimpleNamespace(), enable_in_batch_packing=True) From 18b42fc683bfdefa3a9478ec97f6c120ec7f9dd3 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 11 Aug 2026 22:39:30 -0700 Subject: [PATCH 16/72] fix(recipe): fit Muse Glimmer training on H100 Signed-off-by: yaoyu-33 --- .../recipes/muse_glimmer/h100/muse_glimmer.py | 53 +++++++++++++++---- .../recipes/test_muse_glimmer_recipes.py | 35 +++++++++--- 2 files changed, 71 insertions(+), 17 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index c0c0a5bc72..9b45d66e86 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -16,6 +16,8 @@ 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 @@ -31,14 +33,23 @@ _CORD_V2_REVISION = "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret -def _apply_model_and_data(cfg: ConfigContainer, *, seq_length: int, context_parallel_size: int) -> None: +def _apply_model_and_data( + cfg: ConfigContainer, + *, + seq_length: int, + tensor_parallel_size: int, + pipeline_parallel_size: int, + context_parallel_size: int, +) -> None: """Apply the exact model, multimodal data, and common execution settings.""" cfg.model = AutoBridge.from_hf_pretrained(_MODEL_ID, revision=_MODEL_REVISION).get_model_config() cfg.model.seq_length = seq_length - cfg.model.tensor_model_parallel_size = 4 - cfg.model.pipeline_model_parallel_size = 1 - cfg.model.pipeline_dtype = None + cfg.model.tensor_model_parallel_size = tensor_parallel_size + cfg.model.pipeline_model_parallel_size = pipeline_parallel_size + cfg.model.pipeline_dtype = torch.bfloat16 if pipeline_parallel_size > 1 else None cfg.model.virtual_pipeline_model_parallel_size = None + if pipeline_parallel_size == 2: + cfg.model.hybrid_layer_pattern = f"{'*' * 23}|{'*' * 29}" cfg.model.context_parallel_size = context_parallel_size cfg.model.sequence_parallel = True cfg.model.freeze_language_model = False @@ -89,7 +100,13 @@ def _set_optimizer(cfg: ConfigContainer, *, max_lr: float, warmup_iters: int) -> def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: """Return a 100-step random-init multimodal pretraining config on 128 H100 GPUs.""" cfg = _sft_common_vlm() - _apply_model_and_data(cfg, seq_length=4096, context_parallel_size=1) + _apply_model_and_data( + cfg, + seq_length=4096, + tensor_parallel_size=8, + pipeline_parallel_size=2, + context_parallel_size=1, + ) cfg.dataset.pad_to_max_length = True cfg.dataset.enable_in_batch_packing = False cfg.rng.seed = 1234 @@ -105,7 +122,13 @@ def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: 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() - _apply_model_and_data(cfg, seq_length=4096, context_parallel_size=1) + _apply_model_and_data( + cfg, + seq_length=4096, + tensor_parallel_size=8, + pipeline_parallel_size=2, + context_parallel_size=1, + ) cfg.dataset.pad_to_max_length = True cfg.dataset.enable_in_batch_packing = False cfg.rng.seed = 5678 @@ -120,9 +143,15 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_config() -> ConfigContainer: def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContainer: - """Return a 100-step 8K full SFT config with CP4 on 32 H100 GPUs.""" + """Return a 100-step 8K full SFT config with CP2 on 32 H100 GPUs.""" cfg = _sft_common_vlm() - _apply_model_and_data(cfg, seq_length=8192, context_parallel_size=4) + _apply_model_and_data( + cfg, + seq_length=8192, + tensor_parallel_size=8, + pipeline_parallel_size=2, + context_parallel_size=2, + ) cfg.dataset.pad_to_max_length = True cfg.dataset.enable_in_batch_packing = False cfg.rng.seed = 5678 @@ -139,7 +168,13 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContaine 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() - _apply_model_and_data(cfg, seq_length=8192, context_parallel_size=1) + _apply_model_and_data( + cfg, + seq_length=8192, + tensor_parallel_size=8, + pipeline_parallel_size=1, + context_parallel_size=1, + ) cfg.dataset.pad_to_max_length = True cfg.dataset.enable_in_batch_packing = False diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 561cecccdc..c774b2e87b 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -7,6 +7,7 @@ from types import SimpleNamespace import pytest +import torch from megatron.bridge.peft.lora import LoRA from megatron.bridge.recipes.muse_glimmer.h100 import muse_glimmer as muse_glimmer_recipes @@ -29,20 +30,38 @@ def _offline_recipe_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.parametrize( - ("recipe", "seq_length", "context_parallel_size", "global_batch_size"), + ( + "recipe", + "seq_length", + "tensor_parallel_size", + "pipeline_parallel_size", + "context_parallel_size", + "global_batch_size", + ), [ - (muse_glimmer_30b_pretrain_128gpu_h100_bf16_config, 4096, 1, 1024), - (muse_glimmer_30b_sft_32gpu_h100_bf16_config, 4096, 1, 8), - (muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, 8192, 4, 8), - (muse_glimmer_30b_peft_8gpu_h100_bf16_config, 8192, 1, 8), + (muse_glimmer_30b_pretrain_128gpu_h100_bf16_config, 4096, 8, 2, 1, 1024), + (muse_glimmer_30b_sft_32gpu_h100_bf16_config, 4096, 8, 2, 1, 8), + (muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, 8192, 8, 2, 2, 8), + (muse_glimmer_30b_peft_8gpu_h100_bf16_config, 8192, 8, 1, 1, 8), ], ) -def test_muse_glimmer_recipe_contracts(recipe, seq_length, context_parallel_size, global_batch_size) -> None: +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 == 4 - assert cfg.model.pipeline_model_parallel_size == 1 + 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 = f"{'*' * 23}|{'*' * 29}" if pipeline_parallel_size == 2 else "*" * 52 + assert getattr(cfg.model, "hybrid_layer_pattern", "*" * 52) == expected_pattern + assert getattr(cfg.model, "pipeline_model_parallel_layout", None) is None assert cfg.model.context_parallel_size == context_parallel_size assert cfg.model.sequence_parallel is True assert cfg.model.recompute_granularity == "selective" From a30dcb4765c13d0e3cf52b8c0b9c9abe67b0e3c1 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 00:10:18 -0700 Subject: [PATCH 17/72] fix(recipe): increase Muse pipeline memory headroom Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 2 +- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 9b45d66e86..86951c618d 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -49,7 +49,7 @@ def _apply_model_and_data( cfg.model.pipeline_dtype = torch.bfloat16 if pipeline_parallel_size > 1 else None cfg.model.virtual_pipeline_model_parallel_size = None if pipeline_parallel_size == 2: - cfg.model.hybrid_layer_pattern = f"{'*' * 23}|{'*' * 29}" + cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" cfg.model.context_parallel_size = context_parallel_size cfg.model.sequence_parallel = True cfg.model.freeze_language_model = False diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index c774b2e87b..3207c6597c 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -59,7 +59,7 @@ def test_muse_glimmer_recipe_contracts( 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 = f"{'*' * 23}|{'*' * 29}" if pipeline_parallel_size == 2 else "*" * 52 + expected_pattern = f"{'*' * 20}|{'*' * 32}" if pipeline_parallel_size == 2 else "*" * 52 assert getattr(cfg.model, "hybrid_layer_pattern", "*" * 52) == expected_pattern assert getattr(cfg.model, "pipeline_model_parallel_layout", None) is None assert cfg.model.context_parallel_size == context_parallel_size From a1ce18bf6500c2d3f2cf01298f8a26e64106d1a6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 00:28:31 -0700 Subject: [PATCH 18/72] fix(model): honor Muse transformer backend Signed-off-by: yaoyu-33 --- .../compare_hf_and_megatron/compare.py | 28 ++++++----- .../muse_glimmer/modeling_muse_glimmer.py | 5 +- .../models/muse_glimmer/__init__.py | 1 + .../models/muse_glimmer/test_muse_glimmer.py | 26 ++++++++++ .../unit_tests/test_compare_mask_handling.py | 50 ++++++++++++++++++- 5 files changed, 92 insertions(+), 18 deletions(-) create mode 100644 tests/unit_tests/models/muse_glimmer/__init__.py diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 9d72542355..0d07dccff1 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -696,14 +696,15 @@ def _load_megatron_model(args): ) if getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None: model_config = bridge.get_model_config() - model_config.tensor_model_parallel_size = tp - model_config.pipeline_model_parallel_size = pp - model_config.expert_model_parallel_size = ep - model_config.expert_tensor_parallel_size = etp - model_config.pipeline_dtype = torch.bfloat16 - model_config.params_dtype = torch.bfloat16 + 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(model_config.transformer) + bridge._get_or_initialize_pg_collection(transformer) else: model_provider = bridge.to_megatron_provider(load_weights=False) model_provider.tensor_model_parallel_size = tp @@ -735,12 +736,13 @@ def _load_megatron_model(args): ) if getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None: model_config = bridge.get_model_config() - model_config.tensor_model_parallel_size = tp - model_config.pipeline_model_parallel_size = pp - model_config.expert_model_parallel_size = ep - model_config.expert_tensor_parallel_size = etp - model_config.pipeline_dtype = torch.bfloat16 - model_config.params_dtype = torch.bfloat16 + 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, diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 1079ac895c..d11f03323f 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -24,7 +24,7 @@ import torch.nn.functional as F from megatron.core import parallel_state from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add -from megatron.core.models.backends import LocalSpecProvider +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 @@ -169,8 +169,7 @@ def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: def get_muse_glimmer_hybrid_stack_spec(config: TransformerConfig) -> ModuleSpec: """Build the native Hybrid stack spec used by the Muse decoder.""" - del config - backend = LocalSpecProvider() + backend = get_backend(config.transformer_impl) return ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( 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_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index b10dfcec0b..205b91f2a6 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -11,8 +11,11 @@ import pytest import torch +from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TEDotProductAttention 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 megatron.bridge import AutoBridge @@ -29,6 +32,7 @@ MuseGlimmerCenteredRMSNorm, MuseGlimmerRMSNorm, MuseGlimmerVisionModel, + get_muse_glimmer_hybrid_stack_spec, ) from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import ( MuseGlimmerBridge, @@ -129,6 +133,28 @@ def test_vision_config_contributes_to_runtime_flops() -> None: assert vision_flops > 0 +@pytest.mark.parametrize( + ("transformer_impl", "expected_qkv", "expected_attention"), + [ + ("local", ColumnParallelLinear, DotProductAttention), + ("transformer_engine", TEColumnParallelLinear, TEDotProductAttention), + ], +) +def test_decoder_stack_honors_transformer_backend( + transformer_impl: str, + expected_qkv: type, + expected_attention: 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 + + 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() diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index 1711de9302..9f87d5571e 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -238,8 +238,13 @@ def test_checkpoint_load_uses_builder_config_when_available(self): ): compare._load_megatron_model(args) - assert model_config.tensor_model_parallel_size == 8 - assert model_config.params_dtype is torch.bfloat16 + 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() @@ -254,6 +259,47 @@ def test_checkpoint_load_uses_builder_config_when_available(self): 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(MODEL_CONFIG_CLASS=object) + 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_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) From 927e1b165adc7521a306e398abe6e98dec03e5af Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 00:49:27 -0700 Subject: [PATCH 19/72] fix(recipe): add Muse pipeline workspace headroom Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 2 +- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 86951c618d..cfe4e340f9 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -49,7 +49,7 @@ def _apply_model_and_data( cfg.model.pipeline_dtype = torch.bfloat16 if pipeline_parallel_size > 1 else None cfg.model.virtual_pipeline_model_parallel_size = None if pipeline_parallel_size == 2: - cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" + cfg.model.hybrid_layer_pattern = f"{'*' * 19}|{'*' * 33}" cfg.model.context_parallel_size = context_parallel_size cfg.model.sequence_parallel = True cfg.model.freeze_language_model = False diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 3207c6597c..77406eb8ea 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -59,7 +59,7 @@ def test_muse_glimmer_recipe_contracts( 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 = f"{'*' * 20}|{'*' * 32}" if pipeline_parallel_size == 2 else "*" * 52 + expected_pattern = f"{'*' * 19}|{'*' * 33}" if pipeline_parallel_size == 2 else "*" * 52 assert getattr(cfg.model, "hybrid_layer_pattern", "*" * 52) == expected_pattern assert getattr(cfg.model, "pipeline_model_parallel_layout", None) is None assert cfg.model.context_parallel_size == context_parallel_size From 8b0d06b5d7fffe1cb7a84e5ff79f0a811ffc52c9 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 00:57:00 -0700 Subject: [PATCH 20/72] fix(model): recompute Muse vision layers during training Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/model_config.py | 1 + .../muse_glimmer/modeling_muse_glimmer.py | 20 ++++++++++++++++--- .../recipes/muse_glimmer/h100/muse_glimmer.py | 3 ++- .../models/muse_glimmer/test_muse_glimmer.py | 18 +++++++++++++++++ .../recipes/test_muse_glimmer_recipes.py | 3 ++- 5 files changed, 40 insertions(+), 5 deletions(-) diff --git a/src/megatron/bridge/models/muse_glimmer/model_config.py b/src/megatron/bridge/models/muse_glimmer/model_config.py index 348a7b0ae6..4dea9f753c 100644 --- a/src/megatron/bridge/models/muse_glimmer/model_config.py +++ b/src/megatron/bridge/models/muse_glimmer/model_config.py @@ -84,6 +84,7 @@ class MuseGlimmerModelConfig(HybridModelConfig): 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: diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index d11f03323f..a0f3589a8c 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -37,6 +37,7 @@ from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules from megatron.core.transformer.utils import ensure_metadata_has_dp_cp_group, make_sharded_tensors_for_checkpoint from torch import Tensor +from torch.utils.checkpoint import checkpoint as torch_checkpoint from megatron.bridge.models.gemma.modules import extend_instance from megatron.bridge.models.muse_glimmer.model_config import ( @@ -476,9 +477,10 @@ def forward( class MuseGlimmerVisionModel(nn.Module): """Muse vision tower with 3:1 window/full attention and 2x2 pixel shuffle.""" - def __init__(self, config: MuseGlimmerVisionModelConfig) -> None: + def __init__(self, config: MuseGlimmerVisionModelConfig, *, 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) self.ln_pre = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) @@ -524,7 +526,16 @@ def forward(self, pixel_values: Tensor, grid_thw: Tensor) -> Tensor: 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 - hidden_states = layer(hidden_states, cu_seqlens, position_embeddings) + 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) @@ -577,7 +588,10 @@ def __init__( self.model_config = config customize_muse_glimmer_language_model(self) if pre_process: - self.vision_tower = MuseGlimmerVisionModel(config.vision) + self.vision_tower = MuseGlimmerVisionModel( + config.vision, + recompute_layers=config.recompute_vision_layers, + ) self.vision_adapter = MuseGlimmerVisionAdapter(config) self.vision_projection = nn.Linear( config.projector_hidden_size, diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index cfe4e340f9..910983a528 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -49,12 +49,13 @@ def _apply_model_and_data( cfg.model.pipeline_dtype = torch.bfloat16 if pipeline_parallel_size > 1 else None cfg.model.virtual_pipeline_model_parallel_size = None if pipeline_parallel_size == 2: - cfg.model.hybrid_layer_pattern = f"{'*' * 19}|{'*' * 33}" + cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" cfg.model.context_parallel_size = context_parallel_size 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 diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 205b91f2a6..a48b56c216 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -294,6 +294,24 @@ def test_tiny_vision_model_preserves_expected_token_count() -> None: assert output.shape == (4, 64) +def test_vision_layer_recomputation_is_training_only() -> None: + vision_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).vision + model = MuseGlimmerVisionModel(vision_config, recompute_layers=True) + pixel_values = torch.randn(16, 24) + grid_thw = torch.tensor([[1, 4, 4]]) + + with patch( + "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.torch_checkpoint", + side_effect=lambda function, *args, **kwargs: function(*args), + ) as checkpoint: + 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()) diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 77406eb8ea..ebd5b595dd 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -59,8 +59,9 @@ def test_muse_glimmer_recipe_contracts( 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 = f"{'*' * 19}|{'*' * 33}" if pipeline_parallel_size == 2 else "*" * 52 + expected_pattern = f"{'*' * 20}|{'*' * 32}" if pipeline_parallel_size == 2 else "*" * 52 assert getattr(cfg.model, "hybrid_layer_pattern", "*" * 52) == expected_pattern + assert cfg.model.recompute_vision_layers is True assert getattr(cfg.model, "pipeline_model_parallel_layout", None) is None assert cfg.model.context_parallel_size == context_parallel_size assert cfg.model.sequence_parallel is True From 5df3d61ddf84c44a6544a67c59d32f8b38d47692 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 01:10:54 -0700 Subject: [PATCH 21/72] fix(recipe): bound Muse training image tokens Signed-off-by: yaoyu-33 --- .../bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 5 ++++- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 7 ++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 910983a528..c5265622de 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -65,7 +65,10 @@ def _apply_model_and_data( cfg.dataset.seq_length = seq_length cfg.dataset.hf_processor_path = _MODEL_ID - cfg.dataset.hf_processor_kwargs = {"revision": _MODEL_REVISION} + cfg.dataset.hf_processor_kwargs = { + "revision": _MODEL_REVISION, + "max_image_tokens": 3_072, + } cfg.dataset.source.split = "train" cfg.dataset.source.load_kwargs = {"revision": _CORD_V2_REVISION} cfg.dataset.do_validation = False diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index ebd5b595dd..aee4d051c8 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -62,6 +62,10 @@ def test_muse_glimmer_recipe_contracts( expected_pattern = f"{'*' * 20}|{'*' * 32}" if pipeline_parallel_size == 2 else "*" * 52 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": 3_072, + } assert getattr(cfg.model, "pipeline_model_parallel_layout", None) is None assert cfg.model.context_parallel_size == context_parallel_size assert cfg.model.sequence_parallel is True @@ -69,9 +73,6 @@ def test_muse_glimmer_recipe_contracts( assert cfg.model.recompute_modules == ["core_attn"] assert cfg.model.seq_length == seq_length assert cfg.dataset.seq_length == seq_length - assert cfg.dataset.hf_processor_kwargs == { - "revision": "f84ecc3a0ea984a4c04542a84269e3d065350a6e" # pragma: allowlist secret - } assert cfg.dataset.source.load_kwargs == { "revision": "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret } From a0e1b3ef0d61752890f5c723f09d0a4812a76923 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 01:26:26 -0700 Subject: [PATCH 22/72] perf(recipe): bound Muse verification images Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 4 ++-- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index c5265622de..8c801ef955 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -55,7 +55,7 @@ def _apply_model_and_data( 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.recompute_vision_layers = False cfg.model.transformer_impl = "transformer_engine" cfg.model.attention_backend = "auto" cfg.model.cross_entropy_loss_fusion = True @@ -67,7 +67,7 @@ def _apply_model_and_data( cfg.dataset.hf_processor_path = _MODEL_ID cfg.dataset.hf_processor_kwargs = { "revision": _MODEL_REVISION, - "max_image_tokens": 3_072, + "max_image_tokens": 256, } cfg.dataset.source.split = "train" cfg.dataset.source.load_kwargs = {"revision": _CORD_V2_REVISION} diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index aee4d051c8..de8c20a937 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -61,10 +61,10 @@ def test_muse_glimmer_recipe_contracts( assert cfg.model.pipeline_dtype is (torch.bfloat16 if pipeline_parallel_size > 1 else None) expected_pattern = f"{'*' * 20}|{'*' * 32}" if pipeline_parallel_size == 2 else "*" * 52 assert getattr(cfg.model, "hybrid_layer_pattern", "*" * 52) == expected_pattern - assert cfg.model.recompute_vision_layers is True + assert cfg.model.recompute_vision_layers is False assert cfg.dataset.hf_processor_kwargs == { "revision": muse_glimmer_recipes._MODEL_REVISION, - "max_image_tokens": 3_072, + "max_image_tokens": 256, } assert getattr(cfg.model, "pipeline_model_parallel_layout", None) is None assert cfg.model.context_parallel_size == context_parallel_size From f34ec264096ed770be3caceea5a21d898e800989 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 01:43:07 -0700 Subject: [PATCH 23/72] fix(recipe): retain Muse vision recompute Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 2 +- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 8c801ef955..8686e18aa2 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -55,7 +55,7 @@ def _apply_model_and_data( cfg.model.freeze_language_model = False cfg.model.freeze_vision_model = False cfg.model.freeze_vision_projection = False - cfg.model.recompute_vision_layers = 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 diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index de8c20a937..757a924c7a 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -61,7 +61,7 @@ def test_muse_glimmer_recipe_contracts( assert cfg.model.pipeline_dtype is (torch.bfloat16 if pipeline_parallel_size > 1 else None) expected_pattern = f"{'*' * 20}|{'*' * 32}" if pipeline_parallel_size == 2 else "*" * 52 assert getattr(cfg.model, "hybrid_layer_pattern", "*" * 52) == expected_pattern - assert cfg.model.recompute_vision_layers is False + assert cfg.model.recompute_vision_layers is True assert cfg.dataset.hf_processor_kwargs == { "revision": muse_glimmer_recipes._MODEL_REVISION, "max_image_tokens": 256, From 44c76b9d3fb21df16546a736e2942a80b1249a53 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 01:48:39 -0700 Subject: [PATCH 24/72] perf(recipe): pack Muse pretraining sequences Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/data/collate_fn.py | 8 ++--- .../recipes/muse_glimmer/h100/muse_glimmer.py | 4 ++- .../models/muse_glimmer/test_collate.py | 36 ++++++++++++++----- .../recipes/test_muse_glimmer_recipes.py | 6 ++-- 4 files changed, 37 insertions(+), 17 deletions(-) diff --git a/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py index ec5d558f22..dc0e6c838e 100644 --- a/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py +++ b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py @@ -46,10 +46,6 @@ def muse_glimmer_collate_fn( in_batch_packing_pad_to_multiple_of: int = 1, ) -> dict[str, Any]: """Collate Muse Glimmer conversations and their THW-grid media tensors.""" - if enable_in_batch_packing: - raise ValueError("Muse Glimmer direct-HF training does not support in-batch packing.") - del in_batch_packing_pad_to_multiple_of - skipped_tokens = extract_skipped_token_ids(processor) boundary_config = assistant_mask_boundary_config_from_markers( processor, @@ -104,8 +100,8 @@ def muse_glimmer_collate_fn( 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, + enable_in_batch_packing=enable_in_batch_packing, + in_batch_packing_pad_to_multiple_of=in_batch_packing_pad_to_multiple_of, ignore_index=IGNORE_INDEX, ) return batch diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 8686e18aa2..3512e06072 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -112,10 +112,12 @@ def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: context_parallel_size=1, ) cfg.dataset.pad_to_max_length = True - cfg.dataset.enable_in_batch_packing = False + cfg.dataset.enable_in_batch_packing = True + cfg.dataset.in_batch_packing_pad_to_multiple_of = 8 cfg.rng.seed = 1234 cfg.train.train_iters = 100 cfg.train.global_batch_size = 1024 + cfg.train.micro_batch_size = 16 _set_optimizer(cfg, max_lr=3e-4, warmup_iters=40) cfg.optimizer.use_precision_aware_optimizer = True cfg.checkpoint.save_interval = 50 diff --git a/tests/unit_tests/models/muse_glimmer/test_collate.py b/tests/unit_tests/models/muse_glimmer/test_collate.py index 6854bf15af..0e73283d34 100644 --- a/tests/unit_tests/models/muse_glimmer/test_collate.py +++ b/tests/unit_tests/models/muse_glimmer/test_collate.py @@ -1,8 +1,5 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -from types import SimpleNamespace - -import pytest import torch from megatron.bridge.data.collators.registry import resolve_model_collate @@ -31,8 +28,8 @@ def apply_chat_template(self, conversations, **kwargs): return { "input_ids": torch.tensor([[11, 12, 13, 14]] * batch_size), "attention_mask": torch.ones(batch_size, 4, dtype=torch.long), - "pixel_values": torch.arange(24, dtype=torch.float32).reshape(2, 12), - "image_grid_thw": torch.tensor([[1, 2, 4]]), + "pixel_values": torch.arange(batch_size * 24, dtype=torch.float32).reshape(batch_size * 2, 12), + "image_grid_thw": torch.tensor([[1, 2, 4]] * batch_size), } @@ -60,6 +57,29 @@ def test_muse_glimmer_collator_builds_shifted_labels_and_visual_inputs(monkeypat assert "pixel_values" not in batch -def test_muse_glimmer_collator_rejects_in_batch_packing(): - with pytest.raises(ValueError, match="does not support in-batch packing"): - muse_glimmer_collate_fn([], SimpleNamespace(), enable_in_batch_packing=True) +def test_muse_glimmer_collator_packs_text_and_preserves_visual_order(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() + examples = [ + {"conversation": [{"role": "user", "content": "image"}, {"role": "assistant", "content": "ok"}]}, + {"conversation": [{"role": "user", "content": "image"}, {"role": "assistant", "content": "ok"}]}, + ] + + batch = muse_glimmer_collate_fn( + examples, + processor, + sequence_length=8, + enable_in_batch_packing=True, + in_batch_packing_pad_to_multiple_of=4, + ) + + assert batch["input_ids"].shape == (1, 8) + assert batch["cu_seqlens_q"].tolist() == [0, 4, 8] + assert batch["cu_seqlens_kv"].tolist() == [0, 4, 8] + assert batch["max_seqlen_q"].item() == 4 + assert batch["total_tokens"] == 8 + assert torch.equal(batch["visual_inputs"].pixel_values, torch.arange(48, dtype=torch.float32).reshape(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/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 757a924c7a..9656bf6842 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -77,10 +77,12 @@ def test_muse_glimmer_recipe_contracts( "revision": "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret } assert cfg.dataset.pad_to_max_length is True - assert cfg.dataset.enable_in_batch_packing is False + is_pretrain = recipe is muse_glimmer_30b_pretrain_128gpu_h100_bf16_config + assert cfg.dataset.enable_in_batch_packing is is_pretrain + assert cfg.dataset.in_batch_packing_pad_to_multiple_of == (8 if is_pretrain else 1) assert cfg.train.train_iters == 100 assert cfg.train.global_batch_size == global_batch_size - assert cfg.train.micro_batch_size == 1 + assert cfg.train.micro_batch_size == (16 if is_pretrain else 1) assert cfg.validation.eval_iters == 0 assert cfg.validation.eval_interval == 0 assert cfg.logger.log_throughput is True From cb029f132d34eea96a21859f9d1dcf3fd4e31f6a Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 01:59:23 -0700 Subject: [PATCH 25/72] perf(recipe): add Muse decoder benchmark Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/data/collate_fn.py | 8 +++-- .../recipes/muse_glimmer/h100/__init__.py | 2 ++ .../recipes/muse_glimmer/h100/muse_glimmer.py | 36 ++++++++++++++++--- .../models/muse_glimmer/test_collate.py | 36 +++++-------------- .../recipes/test_muse_glimmer_recipes.py | 31 +++++++++++++--- 5 files changed, 74 insertions(+), 39 deletions(-) diff --git a/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py index dc0e6c838e..ec5d558f22 100644 --- a/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py +++ b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py @@ -46,6 +46,10 @@ def muse_glimmer_collate_fn( in_batch_packing_pad_to_multiple_of: int = 1, ) -> dict[str, Any]: """Collate Muse Glimmer conversations and their THW-grid media tensors.""" + if enable_in_batch_packing: + raise ValueError("Muse Glimmer direct-HF training does not support in-batch packing.") + del in_batch_packing_pad_to_multiple_of + skipped_tokens = extract_skipped_token_ids(processor) boundary_config = assistant_mask_boundary_config_from_markers( processor, @@ -100,8 +104,8 @@ def muse_glimmer_collate_fn( sequence_length=sequence_length, pad_to_max_length=pad_to_max_length, pad_to_multiple_of=pad_to_multiple_of, - enable_in_batch_packing=enable_in_batch_packing, - in_batch_packing_pad_to_multiple_of=in_batch_packing_pad_to_multiple_of, + enable_in_batch_packing=False, + in_batch_packing_pad_to_multiple_of=1, ignore_index=IGNORE_INDEX, ) return batch diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py index 402701b235..9c631e859f 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py @@ -17,6 +17,7 @@ from megatron.bridge.recipes.muse_glimmer.h100.muse_glimmer import ( muse_glimmer_30b_peft_8gpu_h100_bf16_config, muse_glimmer_30b_pretrain_128gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config, muse_glimmer_30b_sft_32gpu_h100_bf16_config, muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, ) @@ -25,6 +26,7 @@ __all__ = [ "muse_glimmer_30b_peft_8gpu_h100_bf16_config", "muse_glimmer_30b_pretrain_128gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_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 index 3512e06072..8512fd2a41 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -24,7 +24,7 @@ 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.config import ConfigContainer, MockGPTDatasetConfig from megatron.bridge.training.mixed_precision import bf16_mixed @@ -112,12 +112,10 @@ def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: context_parallel_size=1, ) cfg.dataset.pad_to_max_length = True - cfg.dataset.enable_in_batch_packing = True - cfg.dataset.in_batch_packing_pad_to_multiple_of = 8 + cfg.dataset.enable_in_batch_packing = False cfg.rng.seed = 1234 cfg.train.train_iters = 100 cfg.train.global_batch_size = 1024 - cfg.train.micro_batch_size = 16 _set_optimizer(cfg, max_lr=3e-4, warmup_iters=40) cfg.optimizer.use_precision_aware_optimizer = True cfg.checkpoint.save_interval = 50 @@ -125,6 +123,35 @@ def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: return cfg +def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigContainer: + """Return a 50-step mock-data benchmark for the dense decoder on 32 H100 GPUs.""" + cfg = muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() + cfg.model.hybrid_layer_pattern = f"{'*' * 26}|{'*' * 26}" + 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.train_iters = 50 + cfg.train.global_batch_size = 256 + cfg.scheduler.lr_warmup_iters = 5 + cfg.scheduler.lr_decay_iters = 50 + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.checkpoint.save_interval = 0 + 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() @@ -204,6 +231,7 @@ def muse_glimmer_30b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora" __all__ = [ "muse_glimmer_30b_peft_8gpu_h100_bf16_config", "muse_glimmer_30b_pretrain_128gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config", "muse_glimmer_30b_sft_32gpu_h100_bf16_config", "muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config", ] diff --git a/tests/unit_tests/models/muse_glimmer/test_collate.py b/tests/unit_tests/models/muse_glimmer/test_collate.py index 0e73283d34..6854bf15af 100644 --- a/tests/unit_tests/models/muse_glimmer/test_collate.py +++ b/tests/unit_tests/models/muse_glimmer/test_collate.py @@ -1,5 +1,8 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +from types import SimpleNamespace + +import pytest import torch from megatron.bridge.data.collators.registry import resolve_model_collate @@ -28,8 +31,8 @@ def apply_chat_template(self, conversations, **kwargs): return { "input_ids": torch.tensor([[11, 12, 13, 14]] * batch_size), "attention_mask": torch.ones(batch_size, 4, dtype=torch.long), - "pixel_values": torch.arange(batch_size * 24, dtype=torch.float32).reshape(batch_size * 2, 12), - "image_grid_thw": torch.tensor([[1, 2, 4]] * batch_size), + "pixel_values": torch.arange(24, dtype=torch.float32).reshape(2, 12), + "image_grid_thw": torch.tensor([[1, 2, 4]]), } @@ -57,29 +60,6 @@ def test_muse_glimmer_collator_builds_shifted_labels_and_visual_inputs(monkeypat assert "pixel_values" not in batch -def test_muse_glimmer_collator_packs_text_and_preserves_visual_order(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() - examples = [ - {"conversation": [{"role": "user", "content": "image"}, {"role": "assistant", "content": "ok"}]}, - {"conversation": [{"role": "user", "content": "image"}, {"role": "assistant", "content": "ok"}]}, - ] - - batch = muse_glimmer_collate_fn( - examples, - processor, - sequence_length=8, - enable_in_batch_packing=True, - in_batch_packing_pad_to_multiple_of=4, - ) - - assert batch["input_ids"].shape == (1, 8) - assert batch["cu_seqlens_q"].tolist() == [0, 4, 8] - assert batch["cu_seqlens_kv"].tolist() == [0, 4, 8] - assert batch["max_seqlen_q"].item() == 4 - assert batch["total_tokens"] == 8 - assert torch.equal(batch["visual_inputs"].pixel_values, torch.arange(48, dtype=torch.float32).reshape(4, 12)) - assert torch.equal(batch["visual_inputs"].image_grid_thw, torch.tensor([[1, 2, 4], [1, 2, 4]])) +def test_muse_glimmer_collator_rejects_in_batch_packing(): + with pytest.raises(ValueError, match="does not support in-batch packing"): + muse_glimmer_collate_fn([], SimpleNamespace(), enable_in_batch_packing=True) diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 9656bf6842..63171541f4 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -14,10 +14,11 @@ from megatron.bridge.recipes.muse_glimmer.h100 import ( muse_glimmer_30b_peft_8gpu_h100_bf16_config, muse_glimmer_30b_pretrain_128gpu_h100_bf16_config, + muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_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 +from megatron.bridge.training.config import ConfigContainer, MockGPTDatasetConfig from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies @@ -77,12 +78,10 @@ def test_muse_glimmer_recipe_contracts( "revision": "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret } assert cfg.dataset.pad_to_max_length is True - is_pretrain = recipe is muse_glimmer_30b_pretrain_128gpu_h100_bf16_config - assert cfg.dataset.enable_in_batch_packing is is_pretrain - assert cfg.dataset.in_batch_packing_pad_to_multiple_of == (8 if is_pretrain else 1) + assert cfg.dataset.enable_in_batch_packing is False assert cfg.train.train_iters == 100 assert cfg.train.global_batch_size == global_batch_size - assert cfg.train.micro_batch_size == (16 if is_pretrain else 1) + assert cfg.train.micro_batch_size == 1 assert cfg.validation.eval_iters == 0 assert cfg.validation.eval_interval == 0 assert cfg.logger.log_throughput is True @@ -100,6 +99,28 @@ def test_muse_glimmer_pretrain_owns_resume_checkpoint_contract() -> None: assert cfg.checkpoint.load is None +def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: + cfg = muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() + + assert isinstance(cfg.dataset, MockGPTDatasetConfig) + assert cfg.dataset.seq_length == 4096 + assert cfg.model.hybrid_layer_pattern == f"{'*' * 26}|{'*' * 26}" + 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 == 8 + assert cfg.model.pipeline_model_parallel_size == 2 + assert cfg.train.train_iters == 50 + assert cfg.train.global_batch_size == 256 + assert cfg.train.micro_batch_size == 1 + assert cfg.scheduler.lr_warmup_iters == 5 + 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.checkpoint.save_interval == 0 + + def test_muse_glimmer_lora_targets_native_attention_projections() -> None: cfg = muse_glimmer_30b_peft_8gpu_h100_bf16_config() From 58f9aa3ac00e902f1701a19419a44234c16ec1fa Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 02:06:35 -0700 Subject: [PATCH 26/72] perf(recipe): fold Muse decoder benchmark Signed-off-by: yaoyu-33 --- .../bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 5 ++++- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 8 ++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 8512fd2a41..57485111fe 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -126,7 +126,9 @@ def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigContainer: """Return a 50-step mock-data benchmark for the dense decoder on 32 H100 GPUs.""" cfg = muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() - cfg.model.hybrid_layer_pattern = f"{'*' * 26}|{'*' * 26}" + cfg.model.tensor_model_parallel_size = 4 + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.hybrid_layer_pattern = "|".join(["*" * 13] * 4) cfg.model.freeze_vision_model = True cfg.model.freeze_vision_projection = True cfg.model.recompute_vision_layers = False @@ -144,6 +146,7 @@ def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigCont ) cfg.train.train_iters = 50 cfg.train.global_batch_size = 256 + cfg.train.micro_batch_size = 4 cfg.scheduler.lr_warmup_iters = 5 cfg.scheduler.lr_decay_iters = 50 cfg.ddp.overlap_grad_reduce = True diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 63171541f4..d61d723d67 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -104,16 +104,16 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: assert isinstance(cfg.dataset, MockGPTDatasetConfig) assert cfg.dataset.seq_length == 4096 - assert cfg.model.hybrid_layer_pattern == f"{'*' * 26}|{'*' * 26}" + assert cfg.model.hybrid_layer_pattern == "|".join(["*" * 13] * 4) 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 == 8 - assert cfg.model.pipeline_model_parallel_size == 2 + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 4 assert cfg.train.train_iters == 50 assert cfg.train.global_batch_size == 256 - assert cfg.train.micro_batch_size == 1 + assert cfg.train.micro_batch_size == 4 assert cfg.scheduler.lr_warmup_iters == 5 assert cfg.scheduler.lr_decay_iters == 50 assert cfg.ddp.overlap_grad_reduce is True From 768e758ad2a6c2cf3d67f8bca3c4383479bda07e Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 02:10:14 -0700 Subject: [PATCH 27/72] perf(recipe): fit Muse folded microbatch Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 4 ++-- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 57485111fe..7d43b89942 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -145,8 +145,8 @@ def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigCont skip_getting_attention_mask_from_dataset=True, ) cfg.train.train_iters = 50 - cfg.train.global_batch_size = 256 - cfg.train.micro_batch_size = 4 + cfg.train.global_batch_size = 192 + cfg.train.micro_batch_size = 3 cfg.scheduler.lr_warmup_iters = 5 cfg.scheduler.lr_decay_iters = 50 cfg.ddp.overlap_grad_reduce = True diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index d61d723d67..420b9c39fd 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -112,8 +112,8 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: assert cfg.model.tensor_model_parallel_size == 4 assert cfg.model.pipeline_model_parallel_size == 4 assert cfg.train.train_iters == 50 - assert cfg.train.global_batch_size == 256 - assert cfg.train.micro_batch_size == 4 + assert cfg.train.global_batch_size == 192 + assert cfg.train.micro_batch_size == 3 assert cfg.scheduler.lr_warmup_iters == 5 assert cfg.scheduler.lr_decay_iters == 50 assert cfg.ddp.overlap_grad_reduce is True From 5857c4dfd6971e777ad94d26752b57d411b17466 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 02:27:04 -0700 Subject: [PATCH 28/72] fix(model): count full Muse transformer layers Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/model_config.py | 1 + .../bridge/training/utils/flop_utils.py | 8 ++++++ .../models/muse_glimmer/test_muse_glimmer.py | 27 ++++++++++++++++++- 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/muse_glimmer/model_config.py b/src/megatron/bridge/models/muse_glimmer/model_config.py index 4dea9f753c..99ac1ec4ad 100644 --- a/src/megatron/bridge/models/muse_glimmer/model_config.py +++ b/src/megatron/bridge/models/muse_glimmer/model_config.py @@ -71,6 +71,7 @@ class MuseGlimmerModelConfig(HybridModelConfig): """Complete builder configuration for Muse Glimmer.""" builder: ClassVar[str] = "megatron.bridge.models.muse_glimmer.MuseGlimmerModelBuilder" + hybrid_attention_layers_include_mlp: ClassVar[bool] = True vision: MuseGlimmerVisionModelConfig = field(default_factory=MuseGlimmerVisionModelConfig) image_token_id: int = 200_092 diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index 10ffe3a69f..cabb96044e 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -1572,6 +1572,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/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index a48b56c216..5fce95d425 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -38,7 +38,7 @@ MuseGlimmerBridge, MuseGlimmerQKVGMapping, ) -from megatron.bridge.training.utils.flop_utils import vit_flops_from_grid_thw +from megatron.bridge.training.utils.flop_utils import num_floating_point_operations, vit_flops_from_grid_thw pytestmark = pytest.mark.unit @@ -133,6 +133,31 @@ def test_vision_config_contributes_to_runtime_flops() -> None: 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"), [ From b5b9290153f531539c2fafa8aad666a781f8d4b2 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 02:47:50 -0700 Subject: [PATCH 29/72] fix(recipe): use supported Muse context parallelism Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 2 ++ tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 1 + 2 files changed, 3 insertions(+) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 7d43b89942..8f33154334 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -51,6 +51,8 @@ def _apply_model_and_data( if pipeline_parallel_size == 2: cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" cfg.model.context_parallel_size = context_parallel_size + if context_parallel_size > 1: + cfg.model.cp_comm_type = "all_gather" cfg.model.sequence_parallel = True cfg.model.freeze_language_model = False cfg.model.freeze_vision_model = False diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 420b9c39fd..2256c64324 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -69,6 +69,7 @@ def test_muse_glimmer_recipe_contracts( } 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") == ("all_gather" 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"] From fc302adea4b2eae1211e3e0090a6a3f6ba78fa2e Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 03:29:26 -0700 Subject: [PATCH 30/72] perf(recipe): use Muse context parallelism Signed-off-by: yaoyu-33 --- .../bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 6 ++++-- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 8f33154334..c5a8cf0318 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -130,6 +130,8 @@ def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigCont cfg = muse_glimmer_30b_pretrain_128gpu_h100_bf16_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(["*" * 13] * 4) cfg.model.freeze_vision_model = True cfg.model.freeze_vision_projection = True @@ -147,8 +149,8 @@ def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigCont skip_getting_attention_mask_from_dataset=True, ) cfg.train.train_iters = 50 - cfg.train.global_batch_size = 192 - cfg.train.micro_batch_size = 3 + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 4 cfg.scheduler.lr_warmup_iters = 5 cfg.scheduler.lr_decay_iters = 50 cfg.ddp.overlap_grad_reduce = True diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 2256c64324..6607a7da92 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -112,9 +112,11 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: 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 == 192 - assert cfg.train.micro_batch_size == 3 + assert cfg.train.global_batch_size == 128 + assert cfg.train.micro_batch_size == 4 assert cfg.scheduler.lr_warmup_iters == 5 assert cfg.scheduler.lr_decay_iters == 50 assert cfg.ddp.overlap_grad_reduce is True From 196b8b26adbd387e7d44a77d9b37b12745e06508 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 03:41:54 -0700 Subject: [PATCH 31/72] perf(recipe): overlap Muse tensor parallelism Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 2 ++ tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index c5a8cf0318..582bb55da6 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -24,6 +24,7 @@ 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.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer, MockGPTDatasetConfig from megatron.bridge.training.mixed_precision import bf16_mixed @@ -155,6 +156,7 @@ def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigCont cfg.scheduler.lr_decay_iters = 50 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) cfg.checkpoint.save_interval = 0 return cfg diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 6607a7da92..6b65a4baf5 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -121,6 +121,8 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: 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_interval == 0 From e773fd0bfa2ba8a5c295f11f50e7758fa09dab6c Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 03:51:01 -0700 Subject: [PATCH 32/72] fix(checkpoint): read HybridModel parallelism Signed-off-by: yaoyu-33 --- src/megatron/bridge/training/checkpointing.py | 21 ++++++++++++------- .../unit_tests/training/test_checkpointing.py | 12 +++++++++++ 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index d2bf088a8a..e061c84eca 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -23,6 +23,7 @@ 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 @@ -161,6 +162,16 @@ 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.""" @@ -2816,10 +2827,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, @@ -2936,10 +2944,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/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 307b7999e6..749451e11a 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -41,6 +41,7 @@ _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, @@ -98,6 +99,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", [ From 891147e96ccf163dd6fdd6b71276445dc764c342 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 04:14:45 -0700 Subject: [PATCH 33/72] docs(model): verify Muse Glimmer LoRA Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 40 +++++++++++++------ 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index 716dbac334..325396695a 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -12,7 +12,7 @@ summary: > use a legacy model provider or GPTModel wrapper. CPU and TP8 GPU conversion, exact round-trip tensor equality, numerical forward parity, and deterministic image-conditioned inference are verified. - Training remains unverified. + TP8 LoRA PEFT is also verified; the remaining training modes are unverified. verification_index: model_level: verified: @@ -24,7 +24,8 @@ verification_index: - inference training: H100: - unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] + verified: [peft] + unverified: [pretrain, sft, sft_export_inference, sft_long_context, checkpoint_resume] model: hf_id: meta-models/Muse-Glimmer-30B hf_revision: f84ecc3a0ea984a4c04542a84269e3d065350a6e # pragma: allowlist secret @@ -224,20 +225,35 @@ items: peft: H100: - status: unverified + status: verified precision: bf16 enabled_features: {} - command: null - last_verified: null + bridge_commit: e773fd0bfa2ba8a5c295f11f50e7758fa09dab6c # 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 + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-08-12 metrics: - initial_loss: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + initial_loss: 0.7247973 + final_loss: 0.04332113 + last_10_steps_step_time_ms_avg: 5291.7 + last_10_steps_model_tflops_per_gpu_avg: 258.34 expected_result: > - Complete a pinned-data 100-step PEFT run with an audited adapter target - set, finite loss, no skipped or NaN iterations, all four metrics, and a - reloadable adapter checkpoint. + 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.7247973 to 0.04332113, all four metrics are recorded, and the complete + eight-shard iter_0000100 adapter checkpoint is saved with its run config, + metadata, and train state. checkpoint_resume: H100: From 7d5037e6763bd6b88f71260270bcc42e46281320 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 05:13:12 -0700 Subject: [PATCH 34/72] perf(recipe): canonicalize Muse Glimmer benchmark Signed-off-by: yaoyu-33 --- scripts/training/recipe_metadata.py | 1 + .../perf_recipes/muse_glimmer/__init__.py | 22 ++++++ .../muse_glimmer/h100/__init__.py | 22 ++++++ .../muse_glimmer/h100/muse_glimmer.py | 68 +++++++++++++++++++ .../recipes/muse_glimmer/h100/__init__.py | 2 - .../recipes/muse_glimmer/h100/muse_glimmer.py | 39 +---------- .../recipes/test_muse_glimmer_recipes.py | 23 +++++-- 7 files changed, 131 insertions(+), 46 deletions(-) create mode 100644 src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py create mode 100644 src/megatron/bridge/perf_recipes/muse_glimmer/h100/__init__.py create mode 100644 src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py diff --git a/scripts/training/recipe_metadata.py b/scripts/training/recipe_metadata.py index dfec57ed8b..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"), 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..b76dfca4f5 --- /dev/null +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py @@ -0,0 +1,22 @@ +# 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, +) + + +__all__ = ["muse_glimmer_30b_pretrain_32gpu_h100_bf16_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..d918d1d089 --- /dev/null +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/__init__.py @@ -0,0 +1,22 @@ +# 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, +) + + +__all__ = ["muse_glimmer_30b_pretrain_32gpu_h100_bf16_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..f9bd579099 --- /dev/null +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py @@ -0,0 +1,68 @@ +# 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 +from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS +from megatron.bridge.recipes.muse_glimmer.h100 import muse_glimmer_30b_pretrain_128gpu_h100_bf16_config +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ConfigContainer, MockGPTDatasetConfig + + +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_128gpu_h100_bf16_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(["*" * 13] * 4) + 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") + 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, + } + return cfg + + +__all__ = ["muse_glimmer_30b_pretrain_32gpu_h100_bf16_config"] diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py index 9c631e859f..402701b235 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py @@ -17,7 +17,6 @@ from megatron.bridge.recipes.muse_glimmer.h100.muse_glimmer import ( muse_glimmer_30b_peft_8gpu_h100_bf16_config, muse_glimmer_30b_pretrain_128gpu_h100_bf16_config, - muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config, muse_glimmer_30b_sft_32gpu_h100_bf16_config, muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, ) @@ -26,7 +25,6 @@ __all__ = [ "muse_glimmer_30b_peft_8gpu_h100_bf16_config", "muse_glimmer_30b_pretrain_128gpu_h100_bf16_config", - "muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_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 index 582bb55da6..e1fd6161f1 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -24,8 +24,7 @@ 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.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ConfigContainer, MockGPTDatasetConfig +from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.mixed_precision import bf16_mixed @@ -126,41 +125,6 @@ def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: return cfg -def muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() -> ConfigContainer: - """Return a 50-step mock-data benchmark for the dense decoder on 32 H100 GPUs.""" - cfg = muse_glimmer_30b_pretrain_128gpu_h100_bf16_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(["*" * 13] * 4) - 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.train_iters = 50 - cfg.train.global_batch_size = 128 - cfg.train.micro_batch_size = 4 - cfg.scheduler.lr_warmup_iters = 5 - cfg.scheduler.lr_decay_iters = 50 - 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) - cfg.checkpoint.save_interval = 0 - 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() @@ -240,7 +204,6 @@ def muse_glimmer_30b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora" __all__ = [ "muse_glimmer_30b_peft_8gpu_h100_bf16_config", "muse_glimmer_30b_pretrain_128gpu_h100_bf16_config", - "muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config", "muse_glimmer_30b_sft_32gpu_h100_bf16_config", "muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config", ] diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 6b65a4baf5..ab53edc26e 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -10,11 +10,13 @@ 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, +) 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_128gpu_h100_bf16_config, - muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config, muse_glimmer_30b_sft_32gpu_h100_bf16_config, muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, ) @@ -101,7 +103,7 @@ def test_muse_glimmer_pretrain_owns_resume_checkpoint_contract() -> None: def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: - cfg = muse_glimmer_30b_pretrain_performance_32gpu_h100_bf16_config() + cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_config() assert isinstance(cfg.dataset, MockGPTDatasetConfig) assert cfg.dataset.seq_length == 4096 @@ -115,15 +117,24 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: 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 == 128 - assert cfg.train.micro_batch_size == 4 - assert cfg.scheduler.lr_warmup_iters == 5 + 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_interval == 0 + 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, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + "TORCH_NCCL_HIGH_PRIORITY": 1, + } def test_muse_glimmer_lora_targets_native_attention_projections() -> None: From a9ec3a059357f72782306ebe07964151414e3c89 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 05:17:24 -0700 Subject: [PATCH 35/72] feat(recipe): add reproducible Muse pretrain Signed-off-by: yaoyu-33 --- .../bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py | 6 ++++-- src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py | 4 ++-- .../bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 8 ++++---- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 8 ++++---- 4 files changed, 14 insertions(+), 12 deletions(-) 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 index f9bd579099..fbe2b0ec50 100644 --- a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py @@ -16,14 +16,16 @@ from megatron.bridge.perf_recipes._common import _benchmark_common from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS -from megatron.bridge.recipes.muse_glimmer.h100 import muse_glimmer_30b_pretrain_128gpu_h100_bf16_config +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_bf16_config() -> ConfigContainer: """Muse Glimmer dense decoder pretrain: 32× H100, BF16, TP=4 PP=4 CP=2.""" - cfg = muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() + 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 diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py index 402701b235..ff7c457c45 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/__init__.py @@ -16,7 +16,7 @@ from megatron.bridge.recipes.muse_glimmer.h100.muse_glimmer import ( muse_glimmer_30b_peft_8gpu_h100_bf16_config, - muse_glimmer_30b_pretrain_128gpu_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, ) @@ -24,7 +24,7 @@ __all__ = [ "muse_glimmer_30b_peft_8gpu_h100_bf16_config", - "muse_glimmer_30b_pretrain_128gpu_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 index e1fd6161f1..bb53d0c99c 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -103,8 +103,8 @@ def _set_optimizer(cfg: ConfigContainer, *, max_lr: float, warmup_iters: int) -> cfg.scheduler.lr_decay_iters = 100 -def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: - """Return a 100-step random-init multimodal pretraining config on 128 H100 GPUs.""" +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() _apply_model_and_data( cfg, @@ -117,7 +117,7 @@ def muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer: cfg.dataset.enable_in_batch_packing = False cfg.rng.seed = 1234 cfg.train.train_iters = 100 - cfg.train.global_batch_size = 1024 + cfg.train.global_batch_size = 256 _set_optimizer(cfg, max_lr=3e-4, warmup_iters=40) cfg.optimizer.use_precision_aware_optimizer = True cfg.checkpoint.save_interval = 50 @@ -203,7 +203,7 @@ def muse_glimmer_30b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora" __all__ = [ "muse_glimmer_30b_peft_8gpu_h100_bf16_config", - "muse_glimmer_30b_pretrain_128gpu_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/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index ab53edc26e..7396a9584b 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -16,7 +16,7 @@ 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_128gpu_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, ) @@ -42,7 +42,7 @@ def _offline_recipe_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: "global_batch_size", ), [ - (muse_glimmer_30b_pretrain_128gpu_h100_bf16_config, 4096, 8, 2, 1, 1024), + (muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config, 4096, 8, 2, 1, 256), (muse_glimmer_30b_sft_32gpu_h100_bf16_config, 4096, 8, 2, 1, 8), (muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, 8192, 8, 2, 2, 8), (muse_glimmer_30b_peft_8gpu_h100_bf16_config, 8192, 8, 1, 1, 8), @@ -91,7 +91,7 @@ def test_muse_glimmer_recipe_contracts( def test_muse_glimmer_pretrain_owns_resume_checkpoint_contract() -> None: - cfg = muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() + cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() assert cfg.rng.seed == 1234 assert cfg.scheduler.lr_warmup_iters == 40 @@ -166,6 +166,6 @@ def from_hf_pretrained(*args, **kwargs): monkeypatch.setattr(muse_glimmer_recipes, "AutoBridge", _AutoBridge) - cfg = muse_glimmer_30b_pretrain_128gpu_h100_bf16_config() + cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() assert cfg.model is not None From 1d20fafe1d529f8eb94233efe6ea3c5f1f18dca5 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 05:24:21 -0700 Subject: [PATCH 36/72] perf(recipe): add Muse Glimmer FP8 benchmark Signed-off-by: yaoyu-33 --- .../perf_recipes/muse_glimmer/__init__.py | 6 ++++- .../muse_glimmer/h100/__init__.py | 6 ++++- .../muse_glimmer/h100/muse_glimmer.py | 25 ++++++++++++++++--- .../recipes/test_muse_glimmer_recipes.py | 13 ++++++++++ 4 files changed, 44 insertions(+), 6 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py b/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py index b76dfca4f5..cd5cea1076 100644 --- a/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/__init__.py @@ -16,7 +16,11 @@ 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"] +__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 index d918d1d089..b3410cde73 100644 --- a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/__init__.py +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/__init__.py @@ -16,7 +16,11 @@ 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"] +__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 index fbe2b0ec50..51029c4e64 100644 --- a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py @@ -14,7 +14,7 @@ """H100 performance recipes for Muse Glimmer.""" -from megatron.bridge.perf_recipes._common import _benchmark_common +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, @@ -23,8 +23,8 @@ from megatron.bridge.training.config import ConfigContainer, MockGPTDatasetConfig -def muse_glimmer_30b_pretrain_32gpu_h100_bf16_config() -> ConfigContainer: - """Muse Glimmer dense decoder pretrain: 32× H100, BF16, TP=4 PP=4 CP=2.""" +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 @@ -67,4 +67,21 @@ def muse_glimmer_30b_pretrain_32gpu_h100_bf16_config() -> ConfigContainer: return cfg -__all__ = ["muse_glimmer_30b_pretrain_32gpu_h100_bf16_config"] +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") + 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") + return cfg + + +__all__ = [ + "muse_glimmer_30b_pretrain_32gpu_h100_bf16_config", + "muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config", +] diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 7396a9584b..3a9fa6218c 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -12,6 +12,7 @@ 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 ( @@ -137,6 +138,18 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: } +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 == 96 + assert cfg.train.micro_batch_size == 3 + + def test_muse_glimmer_lora_targets_native_attention_projections() -> None: cfg = muse_glimmer_30b_peft_8gpu_h100_bf16_config() From fc4cbf99d310a5c542b168b9d30ab906941a6ea6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 05:31:53 -0700 Subject: [PATCH 37/72] perf(recipe): raise Muse FP8 microbatch Signed-off-by: yaoyu-33 --- .../bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py | 2 ++ tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) 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 index 51029c4e64..2c3c6fd7bb 100644 --- a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py @@ -78,6 +78,8 @@ 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 = 128 + cfg.train.micro_batch_size = 4 return cfg diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 3a9fa6218c..278f4947d3 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -146,8 +146,8 @@ def test_muse_glimmer_fp8_performance_recipe_uses_current_scaling() -> None: 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 == 96 - assert cfg.train.micro_batch_size == 3 + assert cfg.train.global_batch_size == 128 + assert cfg.train.micro_batch_size == 4 def test_muse_glimmer_lora_targets_native_attention_projections() -> None: From 01c9249155490d2a3b028ced2d4719bcff09641d Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 06:05:51 -0700 Subject: [PATCH 38/72] perf(model): use native Muse decoder norms Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/modeling_muse_glimmer.py | 16 ++++++++++++---- .../models/muse_glimmer/test_muse_glimmer.py | 11 +++++++---- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index a0f3589a8c..11a880f42b 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -124,12 +124,19 @@ def forward(self, hidden_states: Tensor) -> Tensor: return _rms_norm(hidden_states, self.eps) +def _centered_rms_norm_cls(config: TransformerConfig) -> type[nn.Module]: + """Select the backend-native centered RMSNorm when it supports Muse semantics.""" + if config.transformer_impl == "local": + return MuseGlimmerCenteredRMSNorm + return get_backend(config.transformer_impl).layer_norm(rms_norm=True) + + 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) - self.post_layernorm = MuseGlimmerCenteredRMSNorm( + self.post_layernorm = _centered_rms_norm_cls(config)( config, config.hidden_size, eps=float(getattr(config, "post_norm_epsilon", 1e-8)), @@ -157,7 +164,7 @@ def __init__( ffn_hidden_size=ffn_hidden_size, **kwargs, ) - self.post_layernorm = MuseGlimmerCenteredRMSNorm( + self.post_layernorm = _centered_rms_norm_cls(config)( config, config.hidden_size, eps=float(getattr(config, "post_norm_epsilon", 1e-8)), @@ -171,13 +178,14 @@ def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: 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) + centered_rms_norm = _centered_rms_norm_cls(config) return ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( attention_layer=ModuleSpec( module=TransformerLayer, submodules=TransformerLayerSubmodules( - input_layernorm=MuseGlimmerCenteredRMSNorm, + input_layernorm=centered_rms_norm, self_attention=ModuleSpec( module=MuseGlimmerSelfAttention, params={"attn_mask_type": AttnMaskType.causal}, @@ -190,7 +198,7 @@ def get_muse_glimmer_hybrid_stack_spec(config: TransformerConfig) -> ModuleSpec: ), ), self_attn_bda=get_bias_dropout_add, - pre_mlp_layernorm=MuseGlimmerCenteredRMSNorm, + pre_mlp_layernorm=centered_rms_norm, mlp=functools.partial( MuseGlimmerMLP.as_mlp_submodule, submodules=MLPSubmodules( diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 5fce95d425..703f8c9be6 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -11,7 +11,7 @@ import pytest import torch -from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TEDotProductAttention +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 @@ -159,16 +159,17 @@ def test_hybrid_attention_layers_use_full_transformer_flops() -> None: @pytest.mark.parametrize( - ("transformer_impl", "expected_qkv", "expected_attention"), + ("transformer_impl", "expected_qkv", "expected_attention", "expected_norm"), [ - ("local", ColumnParallelLinear, DotProductAttention), - ("transformer_engine", TEColumnParallelLinear, TEDotProductAttention), + ("local", ColumnParallelLinear, DotProductAttention, MuseGlimmerCenteredRMSNorm), + ("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 @@ -178,6 +179,8 @@ def test_decoder_stack_honors_transformer_backend( 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: From 1fbf9630639f83f0fbfbc840a2f77878bc9a6a97 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 06:28:12 -0700 Subject: [PATCH 39/72] perf(recipe): tune Muse Glimmer FP8 layout Signed-off-by: yaoyu-33 --- .../bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py | 6 +++--- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) 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 index 2c3c6fd7bb..985823bad7 100644 --- a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py @@ -30,7 +30,7 @@ def _muse_glimmer_30b_pretrain_32gpu_h100_config() -> ConfigContainer: 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(["*" * 13] * 4) + 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 @@ -78,8 +78,8 @@ 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 = 128 - cfg.train.micro_batch_size = 4 + cfg.train.global_batch_size = 192 + cfg.train.micro_batch_size = 6 return cfg diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 278f4947d3..904f05b596 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -108,7 +108,7 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: assert isinstance(cfg.dataset, MockGPTDatasetConfig) assert cfg.dataset.seq_length == 4096 - assert cfg.model.hybrid_layer_pattern == "|".join(["*" * 13] * 4) + 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 @@ -146,8 +146,8 @@ def test_muse_glimmer_fp8_performance_recipe_uses_current_scaling() -> None: 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 == 128 - assert cfg.train.micro_batch_size == 4 + assert cfg.train.global_batch_size == 192 + assert cfg.train.micro_batch_size == 6 def test_muse_glimmer_lora_targets_native_attention_projections() -> None: From e9da92f6255ed560ea85b6a78ff7ff17fb4a289a Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 06:59:41 -0700 Subject: [PATCH 40/72] docs(model): record Muse performance evidence Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 43 +++++++++++++++---- examples/models/muse_glimmer/README.md | 2 +- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index 325396695a..fe03d1dd3f 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -3,13 +3,13 @@ title: muse_glimmer_30b summary: > - Performance disclaimer: this model has not been performance-tuned; reported - timing and throughput metrics are sanity checks, not optimized performance - results. 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. + 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. CPU and TP8 GPU conversion, exact round-trip tensor equality, numerical forward parity, and deterministic image-conditioned inference are verified. TP8 LoRA PEFT is also verified; the remaining training modes are unverified. @@ -26,6 +26,8 @@ verification_index: H100: verified: [peft] unverified: [pretrain, sft, sft_export_inference, sft_long_context, checkpoint_resume] + performance: + H100: verified model: hf_id: meta-models/Muse-Glimmer-30B hf_revision: f84ecc3a0ea984a4c04542a84269e3d065350a6e # pragma: allowlist secret @@ -33,7 +35,7 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: e322510641fd910a48ef6f6dc143eac9573366e5 # pragma: allowlist secret + bridge_commit: 1fbf9630639f83f0fbfbc840a2f77878bc9a6a97 # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -278,3 +280,28 @@ items: 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 + 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 + last_verified: 2026-08-12 + metrics: + initial_loss: 12.26874 + final_loss: 0.5229914 + last_10_steps_step_time_ms_avg: 9353.63 + last_10_steps_model_tflops_per_gpu_avg: 429.18 + 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.26874 to 0.5229914 with zero skipped or NaN + iterations. Steps 41-50 average 9,353.63 ms and 429.18 model + TFLOP/s/GPU, and the process exits successfully without checkpoint + output. diff --git a/examples/models/muse_glimmer/README.md b/examples/models/muse_glimmer/README.md index 332134443c..d3420443d1 100644 --- a/examples/models/muse_glimmer/README.md +++ b/examples/models/muse_glimmer/README.md @@ -38,7 +38,7 @@ bridge = AutoBridge.from_hf_pretrained( model_config = bridge.get_model_config() # Apply deployment-specific parallelism before construction. -model_config.tensor_model_parallel_size = 8 +model_config.transformer.tensor_model_parallel_size = 8 models = bridge.get_model(model_config) ``` From a1122c6091f267618e92a636b4e440d3baf2ba1a Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 07:59:45 -0700 Subject: [PATCH 41/72] fix(model): stabilize Muse checkpoint and inference verification Signed-off-by: yaoyu-33 --- scripts/inference/vlm_generation.py | 48 +++++++++- .../muse_glimmer/modeling_muse_glimmer.py | 47 +++++++++ .../models/muse_glimmer/test_muse_glimmer.py | 45 +++++++++ .../scripts/inference/test_vlm_generation.py | 96 +++++++++++++++++++ 4 files changed, 235 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/scripts/inference/test_vlm_generation.py diff --git a/scripts/inference/vlm_generation.py b/scripts/inference/vlm_generation.py index 51d8d9ace5..b461a74f42 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 @@ -153,6 +154,31 @@ def _uses_model_builder(bridge: AutoBridge) -> bool: return getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None +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, *, @@ -161,6 +187,7 @@ def _configure_builder_model( 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() @@ -171,6 +198,7 @@ def _configure_builder_model( 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 @@ -178,6 +206,8 @@ def _configure_builder_model( 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 @@ -223,6 +253,7 @@ def main(args) -> None: ep=ep, etp=etp, pp_layout=args.pp_layout, + deterministic=args.deterministic, ) model_config.finalize() bridge._get_or_initialize_pg_collection(model_config.transformer) @@ -233,6 +264,7 @@ def main(args) -> None: 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 @@ -245,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 @@ -256,7 +289,14 @@ def main(args) -> None: else: print_rank_0(f"Loading HuggingFace model from: {args.hf_model_path}") if _uses_model_builder(bridge): - model_config = _configure_builder_model(bridge, tp=tp, pp=pp, ep=ep, etp=etp) + 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, @@ -269,6 +309,7 @@ def main(args) -> None: 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) @@ -475,6 +516,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/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 11a880f42b..993ab03924 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -36,6 +36,7 @@ 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, make_sharded_tensors_for_checkpoint +from megatron.core.utils import get_pg_rank from torch import Tensor from torch.utils.checkpoint import checkpoint as torch_checkpoint @@ -563,6 +564,19 @@ def forward(self, hidden_states: Tensor) -> Tensor: class MuseGlimmerModel(HybridModel): """Native MCore Hybrid model with the Muse vision modules attached.""" + _CENTERED_NORM_EXTRA_STATE_SUFFIXES = ( + "input_layernorm._extra_state", + "self_attention.post_layernorm._extra_state", + "pre_mlp_layernorm._extra_state", + "mlp.post_layernorm._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, @@ -615,6 +629,39 @@ def __init__( 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 RMSNorm adds an empty ``_extra_state`` object that + the local centered RMSNorm does not expose. Muse conversion uses the + local backend so it can run on CPU, while training may use Transformer + Engine. Omitting only these empty norm objects keeps those checkpoints + 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._CENTERED_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 centered RMSNorm extra state must be empty: {key}.") + return sharded_state_dict + def freeze( self, *, diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 703f8c9be6..87943fcf59 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -442,6 +442,51 @@ def test_builder_constructs_native_hybrid_model_on_cpu(tiny_hybrid_model: MuseGl 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.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.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="centered RMSNorm 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() 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..15023c5441 --- /dev/null +++ b/tests/unit_tests/scripts/inference/test_vlm_generation.py @@ -0,0 +1,96 @@ +# 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_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 From 05767887281d305ea6d7c4dd145e959df79df73b Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 08:32:34 -0700 Subject: [PATCH 42/72] docs(model): verify Muse modeling workflows Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 39 +++++++++---------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index fe03d1dd3f..023c1e1e3b 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -35,7 +35,7 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: 1fbf9630639f83f0fbfbc840a2f77878bc9a6a97 # pragma: allowlist secret + bridge_commit: a1122c6091f267618e92a636b4e440d3baf2ba1a # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -47,7 +47,7 @@ items: --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e --megatron-path work/model-verification/muse-glimmer-30b/cpu-megatron --torch-dtype bfloat16 --overwrite - last_verified: 2026-08-11 + last_verified: 2026-08-12 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 @@ -59,14 +59,13 @@ items: hf_to_megatron_gpu: status: verified precision: bf16 - bridge_commit: 37df4b0bec9f5fff6240818baa55c5a7c0aab1e9 # pragma: allowlist secret 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-11 + last_verified: 2026-08-12 expected_result: > The eight-GPU command builds the complete native HybridModel, completes all 1,228 mappings, validates the TP8 distributed state, and persists @@ -84,7 +83,7 @@ items: --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-11 + last_verified: 2026-08-12 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, @@ -95,7 +94,6 @@ items: megatron_to_hf_gpu: status: verified precision: bf16 - bridge_commit: 0bbd420fada019562647a17fa420ec981c6f81e3 # pragma: allowlist secret command: > ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 1 --gpus-per-node 8 --hf-model meta-models/Muse-Glimmer-30B @@ -104,7 +102,7 @@ items: --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-11 + last_verified: 2026-08-12 expected_result: > Strict distributed export exits successfully. Every language, vision, adapter, and projection tensor matches the pinned source bitwise: 1,436 @@ -116,8 +114,9 @@ items: 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 + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path meta-models/Muse-Glimmer-30B --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e --model_class MuseGlimmerForConditionalGeneration --megatron_model_path @@ -125,15 +124,15 @@ items: --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-11 + last_verified: 2026-08-12 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.999929, above the 0.99 gate; maximum and - mean absolute logit differences are 0.687500 and 0.080432, respectively. + Logit cosine similarity is 0.999852, above the 0.99 gate; maximum and + mean absolute logit differences are 0.953125 and 0.281018, respectively. inference: status: verified @@ -146,8 +145,8 @@ items: 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 --tp 8 --pp 1 --ep 1 --etp 1 - last_verified: 2026-08-11 + --max_new_tokens 64 --deterministic --tp 8 --pp 1 --ep 1 --etp 1 + last_verified: 2026-08-12 expected_result: |- Two independent greedy runs reload the complete TP8 HybridModel checkpoint and use the same repository image (SHA-256 @@ -230,7 +229,7 @@ items: status: verified precision: bf16 enabled_features: {} - bridge_commit: e773fd0bfa2ba8a5c295f11f50e7758fa09dab6c # pragma: allowlist secret + bridge_commit: 13f8af365d264353213033b63544b7303b66ab96 # 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 @@ -243,17 +242,17 @@ items: logger.tensorboard_dir=null last_verified: 2026-08-12 metrics: - initial_loss: 0.7247973 - final_loss: 0.04332113 - last_10_steps_step_time_ms_avg: 5291.7 - last_10_steps_model_tflops_per_gpu_avg: 258.34 + initial_loss: 0.7233534 + final_loss: 0.03220246 + last_10_steps_step_time_ms_avg: 4709.57 + last_10_steps_model_tflops_per_gpu_avg: 290.42 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.7247973 to 0.04332113, all four metrics are recorded, and the complete + 0.7233534 to 0.03220246, all four metrics are recorded, and the complete eight-shard iter_0000100 adapter checkpoint is saved with its run config, metadata, and train state. From edb31e8fb86a4a29835c556a651a07114b75c4f1 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 08:39:51 -0700 Subject: [PATCH 43/72] fix(model): honor Muse embedding sharing contract Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/muse_glimmer_bridge.py | 2 +- .../models/muse_glimmer/test_muse_glimmer.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index e4b77b963f..c46959c826 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -245,7 +245,7 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerM should_pad_vocab=False, seq_length=int(text_config.max_position_embeddings), parallel_output=True, - share_embeddings_and_output_weights=bool(text_config.tie_word_embeddings), + share_embeddings_and_output_weights=bool(hf_config.tie_word_embeddings), hybrid_layer_pattern="*" * int(text_config.num_hidden_layers), position_embedding_type="rope", rotary_percent=1.0, diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 87943fcf59..dcca25b925 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -85,6 +85,7 @@ def _tiny_hf_config() -> MuseGlimmerConfig: projector_hidden_size=32, projector_hidden_act="gelu", architectures=["MuseGlimmerForConditionalGeneration"], + tie_word_embeddings=False, torch_dtype=torch.bfloat16, ) @@ -124,6 +125,17 @@ def test_config_conversion_uses_model_config_path_only() -> None: assert model_config.vision_config.spatial_merge_size == 2 +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_vision_config_contributes_to_runtime_flops() -> None: model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) cfg = SimpleNamespace(model=model_config) From 3c8a0e5acb32036a43aa16a17d0a87705bca1cc6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 09:04:07 -0700 Subject: [PATCH 44/72] feat(model): verify packed Muse long-context SFT Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/data/collate_fn.py | 70 +++++++++++++++++-- .../recipes/muse_glimmer/h100/muse_glimmer.py | 7 +- .../models/muse_glimmer/test_collate.py | 43 +++++++++--- .../recipes/test_muse_glimmer_recipes.py | 7 +- 4 files changed, 108 insertions(+), 19 deletions(-) diff --git a/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py index ec5d558f22..cdc3ad5a89 100644 --- a/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py +++ b/src/megatron/bridge/models/muse_glimmer/data/collate_fn.py @@ -24,9 +24,11 @@ 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 @@ -46,16 +48,76 @@ def muse_glimmer_collate_fn( in_batch_packing_pad_to_multiple_of: int = 1, ) -> dict[str, Any]: """Collate Muse Glimmer conversations and their THW-grid media tensors.""" - if enable_in_batch_packing: - raise ValueError("Muse Glimmer direct-HF training does not support in-batch packing.") - del in_batch_packing_pad_to_multiple_of - 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( diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index bb53d0c99c..b73d442ca0 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -149,7 +149,7 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_config() -> ConfigContainer: def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContainer: - """Return a 100-step 8K full SFT config with CP2 on 32 H100 GPUs.""" + """Return a 100-step packed 8K full SFT config with CP2 on 32 H100 GPUs.""" cfg = _sft_common_vlm() _apply_model_and_data( cfg, @@ -158,11 +158,12 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContaine pipeline_parallel_size=2, context_parallel_size=2, ) - cfg.dataset.pad_to_max_length = True - cfg.dataset.enable_in_batch_packing = False + cfg.dataset.pad_to_max_length = False + cfg.dataset.enable_in_batch_packing = True cfg.rng.seed = 5678 cfg.train.train_iters = 100 cfg.train.global_batch_size = 8 + cfg.train.micro_batch_size = 2 _set_optimizer(cfg, max_lr=5e-6, warmup_iters=10) cfg.checkpoint.save_interval = 100 cfg.checkpoint.load = None diff --git a/tests/unit_tests/models/muse_glimmer/test_collate.py b/tests/unit_tests/models/muse_glimmer/test_collate.py index 6854bf15af..7de7391f64 100644 --- a/tests/unit_tests/models/muse_glimmer/test_collate.py +++ b/tests/unit_tests/models/muse_glimmer/test_collate.py @@ -1,8 +1,5 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -from types import SimpleNamespace - -import pytest import torch from megatron.bridge.data.collators.registry import resolve_model_collate @@ -25,12 +22,13 @@ def __init__(self) -> None: def apply_chat_template(self, conversations, **kwargs): assert kwargs["tokenize"] is True - assert kwargs["padding"] 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.tensor([[11, 12, 13, 14]] * batch_size), - "attention_mask": torch.ones(batch_size, 4, dtype=torch.long), + "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]]), } @@ -60,6 +58,33 @@ def test_muse_glimmer_collator_builds_shifted_labels_and_visual_inputs(monkeypat assert "pixel_values" not in batch -def test_muse_glimmer_collator_rejects_in_batch_packing(): - with pytest.raises(ValueError, match="does not support in-batch packing"): - muse_glimmer_collate_fn([], SimpleNamespace(), enable_in_batch_packing=True) +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/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 904f05b596..630053822b 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -81,11 +81,12 @@ def test_muse_glimmer_recipe_contracts( assert cfg.dataset.source.load_kwargs == { "revision": "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret } - assert cfg.dataset.pad_to_max_length is True - assert cfg.dataset.enable_in_batch_packing is False + 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 == 1 + 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 From da607cec8b03dd4fe4dee130523df2888ab6652e Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 09:11:22 -0700 Subject: [PATCH 45/72] fix(model): support released Muse tying config Signed-off-by: yaoyu-33 --- .../bridge/models/muse_glimmer/muse_glimmer_bridge.py | 4 +++- .../models/muse_glimmer/test_muse_glimmer.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index c46959c826..eefe629757 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -245,7 +245,9 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerM should_pad_vocab=False, seq_length=int(text_config.max_position_embeddings), parallel_output=True, - share_embeddings_and_output_weights=bool(hf_config.tie_word_embeddings), + 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, diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index dcca25b925..7dfc7845dd 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -136,6 +136,17 @@ def test_config_conversion_uses_top_level_embedding_sharing_contract() -> None: 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) From a0a3583acdbd3bb1abde5dbc65f2641fef9c76ae Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 09:27:41 -0700 Subject: [PATCH 46/72] fix(model): select packed Muse CP backend Signed-off-by: yaoyu-33 --- .../recipes/muse_glimmer/h100/muse_glimmer.py | 10 ++++---- .../recipes/test_muse_glimmer_recipes.py | 23 +++++++++++++++---- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index b73d442ca0..993b4a7adf 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -50,9 +50,11 @@ def _apply_model_and_data( cfg.model.virtual_pipeline_model_parallel_size = None if pipeline_parallel_size == 2: cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" + elif pipeline_parallel_size == 4: + cfg.model.hybrid_layer_pattern = "|".join("*" * layers for layers in (9, 15, 15, 13)) cfg.model.context_parallel_size = context_parallel_size if context_parallel_size > 1: - cfg.model.cp_comm_type = "all_gather" + cfg.model.cp_comm_type = "a2a" cfg.model.sequence_parallel = True cfg.model.freeze_language_model = False cfg.model.freeze_vision_model = False @@ -154,8 +156,8 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContaine _apply_model_and_data( cfg, seq_length=8192, - tensor_parallel_size=8, - pipeline_parallel_size=2, + tensor_parallel_size=1, + pipeline_parallel_size=4, context_parallel_size=2, ) cfg.dataset.pad_to_max_length = False @@ -163,7 +165,7 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContaine cfg.rng.seed = 5678 cfg.train.train_iters = 100 cfg.train.global_batch_size = 8 - cfg.train.micro_batch_size = 2 + cfg.train.micro_batch_size = 1 _set_optimizer(cfg, max_lr=5e-6, warmup_iters=10) cfg.checkpoint.save_interval = 100 cfg.checkpoint.load = None diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 630053822b..246ef95e2b 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -45,7 +45,7 @@ def _offline_recipe_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: [ (muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config, 4096, 8, 2, 1, 256), (muse_glimmer_30b_sft_32gpu_h100_bf16_config, 4096, 8, 2, 1, 8), - (muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config, 8192, 8, 2, 2, 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), ], ) @@ -63,7 +63,11 @@ def test_muse_glimmer_recipe_contracts( 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 = f"{'*' * 20}|{'*' * 32}" if pipeline_parallel_size == 2 else "*" * 52 + 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 == { @@ -72,7 +76,7 @@ def test_muse_glimmer_recipe_contracts( } 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") == ("all_gather" if context_parallel_size > 1 else "p2p") + 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"] @@ -86,7 +90,7 @@ def test_muse_glimmer_recipe_contracts( 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.train.micro_batch_size == 1 assert cfg.validation.eval_iters == 0 assert cfg.validation.eval_interval == 0 assert cfg.logger.log_throughput is True @@ -104,6 +108,17 @@ def test_muse_glimmer_pretrain_owns_resume_checkpoint_contract() -> None: 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.num_query_groups == 2 + assert cfg.model.hybrid_layer_pattern == "|".join("*" * layers for layers in (9, 15, 15, 13)) + + def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: cfg = muse_glimmer_30b_pretrain_32gpu_h100_bf16_config() From 99dac20669ac00d9d9e1dfa17ecdb6cad333a3e0 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 09:32:01 -0700 Subject: [PATCH 47/72] fix(model): satisfy Muse packing batch contract Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 2 +- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 993b4a7adf..af1a20e69d 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -89,7 +89,7 @@ def _apply_model_and_data( cfg.validation.eval_iters = 0 cfg.logger.log_interval = 1 cfg.logger.log_throughput = True - cfg.train.micro_batch_size = 1 + cfg.train.micro_batch_size = 2 cfg.train.manual_gc = True cfg.train.manual_gc_interval = 10 cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 246ef95e2b..b9b83515cd 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -90,7 +90,7 @@ def test_muse_glimmer_recipe_contracts( 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 == 1 + 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 @@ -117,6 +117,7 @@ def test_muse_glimmer_long_context_uses_te_supported_packed_cp_layout() -> None: assert cfg.model.cp_comm_type == "a2a" assert cfg.model.num_query_groups == 2 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: From 9149d3f56fe21c5dea8699ffda57db6d657af34a Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 09:35:12 -0700 Subject: [PATCH 48/72] fix(model): scope Muse packing batch override Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index af1a20e69d..07531c682f 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -89,7 +89,7 @@ def _apply_model_and_data( cfg.validation.eval_iters = 0 cfg.logger.log_interval = 1 cfg.logger.log_throughput = True - cfg.train.micro_batch_size = 2 + cfg.train.micro_batch_size = 1 cfg.train.manual_gc = True cfg.train.manual_gc_interval = 10 cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} @@ -165,7 +165,7 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContaine cfg.rng.seed = 5678 cfg.train.train_iters = 100 cfg.train.global_batch_size = 8 - cfg.train.micro_batch_size = 1 + cfg.train.micro_batch_size = 2 _set_optimizer(cfg, max_lr=5e-6, warmup_iters=10) cfg.checkpoint.save_interval = 100 cfg.checkpoint.load = None From 1cf21ff0b6174024931878018a674a13945eb940 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 09:38:05 -0700 Subject: [PATCH 49/72] test(model): keep Muse recipe fixture focused Signed-off-by: yaoyu-33 --- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index b9b83515cd..30295f2f5e 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -115,7 +115,6 @@ def test_muse_glimmer_long_context_uses_te_supported_packed_cp_layout() -> None: 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.num_query_groups == 2 assert cfg.model.hybrid_layer_pattern == "|".join("*" * layers for layers in (9, 15, 15, 13)) assert cfg.train.micro_batch_size == 2 From 00792cafec87300eb737259631f86ef319a3d16e Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 09:45:26 -0700 Subject: [PATCH 50/72] fix(training): slice packed CP metadata on pipeline stages Signed-off-by: yaoyu-33 --- src/megatron/bridge/utils/common_utils.py | 16 ++++++--- .../utils/test_slice_batch_for_cp.py | 34 +++++++++++++++++++ 2 files changed, 45 insertions(+), 5 deletions(-) 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/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 From c7026ab2cbaac76408c23c87d9ab9a27123603d6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 10:33:12 -0700 Subject: [PATCH 51/72] fix(model): stabilize Muse vision attention backward Signed-off-by: yaoyu-33 --- .../muse_glimmer/modeling_muse_glimmer.py | 4 +- .../models/muse_glimmer/test_muse_glimmer.py | 45 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 993ab03924..53c5445958 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -446,7 +446,9 @@ def forward( 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) - output = F.scaled_dot_product_attention(q, k, v, scale=self.scaling) + 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)) diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 7dfc7845dd..911662247e 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -31,6 +31,7 @@ from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( MuseGlimmerCenteredRMSNorm, MuseGlimmerRMSNorm, + MuseGlimmerVisionAttention, MuseGlimmerVisionModel, get_muse_glimmer_hybrid_stack_spec, ) @@ -345,6 +346,50 @@ def test_tiny_vision_model_preserves_expected_token_count() -> None: 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() + ) + + def test_vision_layer_recomputation_is_training_only() -> None: vision_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).vision model = MuseGlimmerVisionModel(vision_config, recompute_layers=True) From 2bf9bbb49444954f6ce4361d7435d12021c842c2 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 10:46:32 -0700 Subject: [PATCH 52/72] docs(model): finalize Muse Glimmer modeling evidence Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index 023c1e1e3b..edb9e04cdf 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -10,9 +10,10 @@ summary: > 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. - CPU and TP8 GPU conversion, exact round-trip tensor equality, numerical - forward parity, and deterministic image-conditioned inference are verified. - TP8 LoRA PEFT is also verified; the remaining training modes are unverified. + All six modeling gates are verified: CPU and TP8 GPU conversion, exact + round-trip tensor equality, numerical forward parity, and deterministic + image-conditioned inference. TP8 LoRA PEFT is also verified; the remaining + training modes are unverified. verification_index: model_level: verified: @@ -35,7 +36,7 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: a1122c6091f267618e92a636b4e440d3baf2ba1a # pragma: allowlist secret + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -131,8 +132,8 @@ items: 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.999852, above the 0.99 gate; maximum and - mean absolute logit differences are 0.953125 and 0.281018, respectively. + Logit cosine similarity is 0.999896, above the 0.99 gate; maximum and + mean absolute logit differences are 0.781250 and 0.139148, respectively. inference: status: verified @@ -148,7 +149,8 @@ items: --max_new_tokens 64 --deterministic --tp 8 --pp 1 --ep 1 --etp 1 last_verified: 2026-08-12 expected_result: |- - Two independent greedy runs reload the complete TP8 HybridModel + At bridge commit c7026ab2cbaac76408c23c87d9ab9a27123603d6, 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 From c232848851b2b3d53f26c5544e93458d273ff460 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 12:04:44 -0700 Subject: [PATCH 53/72] fix(recipe): bound Muse multimodal pretraining batch Signed-off-by: yaoyu-33 --- .../bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 5 ++++- tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 07531c682f..a528f6320b 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -119,7 +119,10 @@ def muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() -> ConfigConta cfg.dataset.enable_in_batch_packing = False cfg.rng.seed = 1234 cfg.train.train_iters = 100 - cfg.train.global_batch_size = 256 + # 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 _set_optimizer(cfg, max_lr=3e-4, warmup_iters=40) cfg.optimizer.use_precision_aware_optimizer = True cfg.checkpoint.save_interval = 50 diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 30295f2f5e..068e555612 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -43,7 +43,7 @@ def _offline_recipe_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: "global_batch_size", ), [ - (muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config, 4096, 8, 2, 1, 256), + (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), From 32d0ead9c4a554da85a38c824e454fd6a7b68107 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 12:13:21 -0700 Subject: [PATCH 54/72] fix(recipe): make Muse pretraining resume reproducible Signed-off-by: yaoyu-33 --- src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py | 3 +++ tests/unit_tests/recipes/test_muse_glimmer_recipes.py | 1 + 2 files changed, 4 insertions(+) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index a528f6320b..613dcfb758 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -117,6 +117,9 @@ def muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() -> ConfigConta ) 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.rng.seed = 1234 cfg.train.train_iters = 100 # TP8/PP2 leaves two data-parallel replicas on 32 GPUs. One sample per diff --git a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index 068e555612..bcdfe87a61 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -104,6 +104,7 @@ def test_muse_glimmer_pretrain_owns_resume_checkpoint_contract() -> None: 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 From 8be93c49130954e3257158db774d59075d2beaa6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 13:07:33 -0700 Subject: [PATCH 55/72] docs(model): verify Muse Glimmer training workflows Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 164 +++++++++++++----- 1 file changed, 121 insertions(+), 43 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index edb9e04cdf..eff388c07e 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -12,8 +12,9 @@ summary: > decoder; it does not use a legacy model provider or GPTModel wrapper. All six modeling gates are verified: CPU and TP8 GPU conversion, exact round-trip tensor equality, numerical forward parity, and deterministic - image-conditioned inference. TP8 LoRA PEFT is also verified; the remaining - training modes are unverified. + 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. verification_index: model_level: verified: @@ -25,8 +26,8 @@ verification_index: - inference training: H100: - verified: [peft] - unverified: [pretrain, sft, sft_export_inference, sft_long_context, checkpoint_resume] + verified: [pretrain, sft, sft_export_inference, sft_long_context, peft] + unverified: [checkpoint_resume] performance: H100: verified model: @@ -36,7 +37,7 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret + bridge_commit: 2bf9bbb49444954f6ce4361d7435d12021c842c2 # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -165,66 +166,143 @@ items: pretrain: H100: - status: unverified + status: verified precision: bf16 enabled_features: {} - command: null - last_verified: null + 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: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + initial_loss: 12.23545 + final_loss: 1.526005 + last_10_steps_step_time_ms_avg: 1092.84 + last_10_steps_model_tflops_per_gpu_avg: 40.59 expected_result: > - Complete a bounded 100-step multimodal pretraining run with finite loss, - no skipped or NaN iterations, all four metrics, and a reloadable final - checkpoint. + 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.23545 to 1.526005; steps 91-100 + average 1,092.84 ms and 40.59 model TFLOP/s/GPU. Complete step-50 and + step-100 model, optimizer, scheduler, RNG, metadata, train-state, and + resolved-config artifacts are saved and reloadable. sft: H100: - status: unverified + status: verified precision: bf16 enabled_features: {} - command: null - last_verified: null + 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: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + 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 expected_result: > - Complete a pinned-data 100-step multimodal SFT run with finite loss, no - skipped or NaN iterations, all four metrics, and a reloadable final - checkpoint. + 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 and 94.03 model TFLOP/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: unverified + status: verified precision: bf16 depends_on: sft - commands: null - last_verified: null - expected_result: > - Export the completed SFT checkpoint to a complete Hugging Face VLM, - strictly reload it, and complete deterministic image-conditioned - generation with the pinned prompt. + 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: unverified + status: verified precision: bf16 - enabled_features: {} - command: null - last_verified: null + 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: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + 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 expected_result: > - Complete a 100-step long-context multimodal SFT run with declared - packing and context-parallel settings, finite loss, no skipped or NaN - iterations, all four metrics, and a reloadable checkpoint. + 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, sequence parallelism, + 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 and 16.07 model TFLOP/s/GPU. The + complete step-100 full-model checkpoint, metadata, train state, and + resolved config are saved and reloadable. peft: H100: From 511124c2f6bb1d7229b00a7c94938167455717de Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 15:31:37 -0700 Subject: [PATCH 56/72] docs(model): pin Muse verification provenance Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index eff388c07e..c884c95857 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -139,6 +139,7 @@ items: inference: status: verified precision: bf16 + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret command: > ./scripts/inference/infer.sh --nodes 1 --gpus-per-node 8 --task vlm-generation --hf_model_path meta-models/Muse-Glimmer-30B @@ -150,8 +151,7 @@ items: --max_new_tokens 64 --deterministic --tp 8 --pp 1 --ep 1 --etp 1 last_verified: 2026-08-12 expected_result: |- - At bridge commit c7026ab2cbaac76408c23c87d9ab9a27123603d6, two - independent greedy runs reload the complete TP8 HybridModel + 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 @@ -168,6 +168,7 @@ items: H100: status: verified precision: bf16 + bridge_commit: 8be93c49130954e3257158db774d59075d2beaa6 # pragma: allowlist secret enabled_features: {} command: > ./scripts/training/train.sh --nodes 4 --gpus-per-node 8 @@ -202,6 +203,7 @@ items: H100: status: verified precision: bf16 + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret enabled_features: {} command: > ./scripts/training/train.sh --nodes 4 --gpus-per-node 8 @@ -237,6 +239,7 @@ items: H100: status: verified precision: bf16 + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret depends_on: sft commands: - > @@ -270,6 +273,7 @@ items: H100: status: verified precision: bf16 + bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret enabled_features: sequence_packing: in_batch context_parallel_size: 2 @@ -309,7 +313,7 @@ items: status: verified precision: bf16 enabled_features: {} - bridge_commit: 13f8af365d264353213033b63544b7303b66ab96 # pragma: allowlist secret + 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 @@ -322,17 +326,17 @@ items: logger.tensorboard_dir=null last_verified: 2026-08-12 metrics: - initial_loss: 0.7233534 - final_loss: 0.03220246 - last_10_steps_step_time_ms_avg: 4709.57 - last_10_steps_model_tflops_per_gpu_avg: 290.42 + initial_loss: 0.7254401 + final_loss: 0.03141141 + last_10_steps_step_time_ms_avg: 4812.79 + last_10_steps_model_tflops_per_gpu_avg: 283.95 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.7233534 to 0.03220246, all four metrics are recorded, and the complete + 0.7254401 to 0.03141141, all four metrics are recorded, and the complete eight-shard iter_0000100 adapter checkpoint is saved with its run config, metadata, and train state. @@ -364,6 +368,7 @@ items: H100: status: verified precision: bf16 + bridge_commit: 1fbf9630639f83f0fbfbc840a2f77878bc9a6a97 # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 4 --gpus-per-node 8 --recipe muse_glimmer_30b_pretrain_32gpu_h100_fp8cs_config From 525b82ab8e0b93da551b7955306597cf8e62a8d7 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 12 Aug 2026 18:07:04 -0700 Subject: [PATCH 57/72] docs(models): refresh Muse Glimmer verification evidence Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 74 +++++++++++-------- 1 file changed, 45 insertions(+), 29 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index c884c95857..fc29c46da7 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -14,7 +14,8 @@ summary: > 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. + canonical dense-decoder performance are also verified. Checkpoint resume + remains unverified. verification_index: model_level: verified: @@ -116,8 +117,8 @@ items: status: verified precision: bf16 command: > - uv run python -m torch.distributed.run --standalone --nproc_per_node=8 - examples/conversion/compare_hf_and_megatron/compare.py + ./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 @@ -185,19 +186,22 @@ items: logger.save_config_filepath=work/model-verification/muse-glimmer-30b/pretrain-reference-resolved.yaml last_verified: 2026-08-12 metrics: - initial_loss: 12.23545 - final_loss: 1.526005 - last_10_steps_step_time_ms_avg: 1092.84 - last_10_steps_model_tflops_per_gpu_avg: 40.59 + 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 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.23545 to 1.526005; steps 91-100 - average 1,092.84 ms and 40.59 model TFLOP/s/GPU. Complete step-50 and - step-100 model, optimizer, scheduler, RNG, metadata, train-state, and - resolved-config artifacts are saved and reloadable. + iterations. Loss decreases from 12.23977 to 1.410300; steps 91-100 + average 1,035.34 ms and 42.49 model TFLOP/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: @@ -300,8 +304,8 @@ items: 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, sequence parallelism, - A2A context-parallel communication, and selective core-attention + 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 and 16.07 model TFLOP/s/GPU. The @@ -322,23 +326,30 @@ items: --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.7254401 - final_loss: 0.03141141 - last_10_steps_step_time_ms_avg: 4812.79 - last_10_steps_model_tflops_per_gpu_avg: 283.95 + 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 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.7254401 to 0.03141141, all four metrics are recorded, and the complete - eight-shard iter_0000100 adapter checkpoint is saved with its run config, - metadata, and train state. + 0.7263421 to 0.03298394; steps 91-100 average 4,823.07 ms and 283.36 + model TFLOP/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: @@ -368,24 +379,29 @@ items: H100: status: verified precision: bf16 - bridge_commit: 1fbf9630639f83f0fbfbc840a2f77878bc9a6a97 # pragma: allowlist secret + 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.26874 - final_loss: 0.5229914 - last_10_steps_step_time_ms_avg: 9353.63 - last_10_steps_model_tflops_per_gpu_avg: 429.18 + 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 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.26874 to 0.5229914 with zero skipped or NaN - iterations. Steps 41-50 average 9,353.63 ms and 429.18 model - TFLOP/s/GPU, and the process exits successfully without checkpoint - output. + Loss remains finite from 12.26882 to 2.442728 with zero skipped or NaN + iterations. Steps 41-50 average 9,373.06 ms and 428.24 model + TFLOP/s/GPU. The persisted post-setup config records the exact runtime, + and the process exits successfully without checkpoint output. From eb4dd2f102b11ebf43ccaba1ec92954b26d05926 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Thu, 13 Aug 2026 17:24:18 -0700 Subject: [PATCH 58/72] refactor(muse-glimmer): use native Transformers config Signed-off-by: yaoyu-33 --- pyproject.toml | 2 +- .../bridge/models/muse_glimmer/__init__.py | 11 +- .../configuration_muse_glimmer.py | 188 ------------------ .../muse_glimmer/modeling_muse_glimmer.py | 2 +- .../muse_glimmer/muse_glimmer_bridge.py | 4 +- .../muse_glimmer/muse_glimmer_builder.py | 2 +- ...model_config.py => muse_glimmer_config.py} | 2 +- uv.lock | 8 +- 8 files changed, 16 insertions(+), 203 deletions(-) delete mode 100644 src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py rename src/megatron/bridge/models/muse_glimmer/{model_config.py => muse_glimmer_config.py} (98%) diff --git a/pyproject.toml b/pyproject.toml index 0966d4a2d9..300e0e6f67 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/src/megatron/bridge/models/muse_glimmer/__init__.py b/src/megatron/bridge/models/muse_glimmer/__init__.py index c78333f11e..8335b8838e 100644 --- a/src/megatron/bridge/models/muse_glimmer/__init__.py +++ b/src/megatron/bridge/models/muse_glimmer/__init__.py @@ -14,19 +14,20 @@ """Muse Glimmer model, builder, configuration, and bridge exports.""" -from megatron.bridge.models.muse_glimmer.configuration_muse_glimmer import ( +from transformers import ( MuseGlimmerConfig, MuseGlimmerTextConfig, MuseGlimmerVisionConfig, ) -from megatron.bridge.models.muse_glimmer.model_config import ( + +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, ) -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 __all__ = [ diff --git a/src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py deleted file mode 100644 index dbdba0f470..0000000000 --- a/src/megatron/bridge/models/muse_glimmer/configuration_muse_glimmer.py +++ /dev/null @@ -1,188 +0,0 @@ -# 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. - -"""Transformers compatibility configuration for Muse Glimmer. - -Muse Glimmer landed after the Transformers version currently supported by -Megatron Bridge. Weight import only needs a parsed configuration and the lazy -safetensors state source, so keep a small configuration fallback here instead -of requiring the full Hugging Face modeling implementation at runtime. -""" - -from __future__ import annotations - -from typing import Any - -from transformers import AutoConfig, PretrainedConfig - - -try: - from transformers import MuseGlimmerConfig, MuseGlimmerTextConfig, MuseGlimmerVisionConfig -except ImportError: - - class MuseGlimmerTextConfig(PretrainedConfig): - """Configuration for the Muse Glimmer text decoder.""" - - model_type = "muse_glimmer_text" - - def __init__( - self, - *, - vocab_size: int = 202_048, - hidden_size: int = 6_656, - intermediate_size: int = 19_968, - num_hidden_layers: int = 52, - num_attention_heads: int = 32, - num_key_value_heads: int = 2, - head_dim: int = 128, - hidden_activation: str = "silu", - max_position_embeddings: int = 131_072, - rms_norm_eps: float = 1e-5, - sliding_window: int = 2_048, - final_logit_softcapping: float = 20.0, - qk_scale_factor: float = 3.87, - output_multiplier: float = 0.19611613513818404, - post_norm_eps: float = 1e-8, - layer_types: list[str] | None = None, - layer_rope_theta: list[float] | None = None, - rope_parameters: dict[str, Any] | None = None, - attention_bias: bool = False, - attention_dropout: float = 0.0, - initializer_range: float = 0.02, - tie_word_embeddings: bool = False, - bos_token_id: int | None = 200_000, - eos_token_id: int | list[int] | None = 200_001, - pad_token_id: int | None = None, - **kwargs: Any, - ) -> None: - super().__init__( - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - pad_token_id=pad_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.hidden_activation = hidden_activation - self.max_position_embeddings = max_position_embeddings - self.rms_norm_eps = rms_norm_eps - self.sliding_window = sliding_window - self.final_logit_softcapping = final_logit_softcapping - self.qk_scale_factor = qk_scale_factor - self.output_multiplier = output_multiplier - self.post_norm_eps = post_norm_eps - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.initializer_range = initializer_range - self.rope_parameters = rope_parameters or {"rope_theta": 500_000.0, "rope_type": "default"} - self.layer_types = layer_types or [ - "full_attention" if (num_hidden_layers - 1 - index) % 4 == 0 else "sliding_attention" - for index in range(num_hidden_layers) - ] - self.layer_rope_theta = layer_rope_theta or [ - 0.0 if layer_type == "full_attention" else float(self.rope_parameters["rope_theta"]) - for layer_type in self.layer_types - ] - - class MuseGlimmerVisionConfig(PretrainedConfig): - """Configuration for the Muse Glimmer vision encoder.""" - - model_type = "muse_glimmer_vision" - - def __init__( - self, - *, - 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_eps: float = 1e-5, - hidden_act: str = "gelu", - layer_types: list[str] | None = None, - rope_parameters: dict[str, Any] | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.patch_size = patch_size - self.patch_temporal = patch_temporal - self.merge_size = merge_size - self.pos_emb_height = pos_emb_height - self.pos_emb_width = pos_emb_width - self.max_position_embeddings = max_position_embeddings - self.layer_norm_eps = layer_norm_eps - self.hidden_act = hidden_act - self.rope_parameters = rope_parameters or {"rope_theta": 10_000.0, "rope_type": "default"} - self.layer_types = layer_types or [ - "full_attention" if index % 4 == 3 or index == num_hidden_layers - 1 else "window_attention" - for index in range(num_hidden_layers) - ] - - class MuseGlimmerConfig(PretrainedConfig): - """Top-level configuration for Muse Glimmer.""" - - model_type = "muse_glimmer" - sub_configs = { - "text_config": MuseGlimmerTextConfig, - "vision_config": MuseGlimmerVisionConfig, - } - - def __init__( - self, - *, - text_config: dict[str, Any] | PretrainedConfig | None = None, - vision_config: dict[str, Any] | PretrainedConfig | None = None, - image_token_id: int = 200_092, - video_token_id: int = 200_091, - out_hidden_size: int = 6_144, - projector_hidden_size: int = 4_096, - projector_hidden_act: str = "gelu", - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.text_config = ( - text_config - if isinstance(text_config, PretrainedConfig) - else MuseGlimmerTextConfig(**(text_config or {})) - ) - self.vision_config = ( - vision_config - if isinstance(vision_config, PretrainedConfig) - else MuseGlimmerVisionConfig(**(vision_config or {})) - ) - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.out_hidden_size = out_hidden_size - self.projector_hidden_size = projector_hidden_size - self.projector_hidden_act = projector_hidden_act - - AutoConfig.register(MuseGlimmerConfig.model_type, MuseGlimmerConfig, exist_ok=True) - - -__all__ = ["MuseGlimmerConfig", "MuseGlimmerTextConfig", "MuseGlimmerVisionConfig"] diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 53c5445958..7fd9ea1e14 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -41,7 +41,7 @@ from torch.utils.checkpoint import checkpoint as torch_checkpoint from megatron.bridge.models.gemma.modules import extend_instance -from megatron.bridge.models.muse_glimmer.model_config import ( +from megatron.bridge.models.muse_glimmer.muse_glimmer_config import ( MuseGlimmerModelConfig, MuseGlimmerVisionModelConfig, ) diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index eefe629757..7648ab100c 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -35,12 +35,12 @@ split_qkv_weights, ) from megatron.bridge.models.conversion.utils import remove_non_pickleables -from megatron.bridge.models.muse_glimmer.model_config import ( +from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import MuseGlimmerModel +from megatron.bridge.models.muse_glimmer.muse_glimmer_config import ( MuseGlimmerModelConfig, MuseGlimmerTransformerConfig, MuseGlimmerVisionModelConfig, ) -from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import MuseGlimmerModel class MuseGlimmerQKVGMapping(MegatronParamMapping[dict[str, torch.Tensor]]): diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py index 57bd374039..787d6bea97 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_builder.py @@ -21,11 +21,11 @@ from megatron.training.models.hybrid import HybridModelBuilder from megatron.training.vocab_utils import calculate_padded_vocab_size -from megatron.bridge.models.muse_glimmer.model_config import MuseGlimmerModelConfig 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): diff --git a/src/megatron/bridge/models/muse_glimmer/model_config.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py similarity index 98% rename from src/megatron/bridge/models/muse_glimmer/model_config.py rename to src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py index 99ac1ec4ad..bbbdfd05de 100644 --- a/src/megatron/bridge/models/muse_glimmer/model_config.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Builder-backed configuration objects for Muse Glimmer.""" +"""Native HybridModel configuration objects for Muse Glimmer.""" from __future__ import annotations diff --git a/uv.lock b/uv.lock index a78ebe3de8..ffa0115200 100644 --- a/uv.lock +++ b/uv.lock @@ -2100,7 +2100,7 @@ requires-dist = [ { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", specifier = ">=4.67.1" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'" }, - { 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" }, ] @@ -4638,7 +4638,7 @@ sdist = { url = "https://pypi.nvidia.com/transformer-engine-torch/transformer_en [[package]] name = "transformers" -version = "5.12.1" +version = "5.15.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, @@ -4651,9 +4651,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]] From 743d6ad417c319d005d24d0a0b2bb8a7a1a9013c Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 14 Aug 2026 13:30:03 -0700 Subject: [PATCH 59/72] refactor(models): unify builder config overrides Signed-off-by: yaoyu-33 --- examples/models/muse_glimmer/README.md | 2 +- src/megatron/bridge/models/common/__init__.py | 2 + src/megatron/bridge/models/common/base.py | 39 ++++++++++++++++++- .../bridge/models/conversion/model_bridge.py | 25 +++++++++--- .../bridge/models/gpt/model_config.py | 31 +++------------ .../muse_glimmer/muse_glimmer_bridge.py | 1 - .../muse_glimmer/muse_glimmer_config.py | 5 ++- .../models/muse_glimmer/test_muse_glimmer.py | 19 +++++++++ .../models/test_gpt_model_config.py | 37 ++++++++++++++++++ 9 files changed, 124 insertions(+), 37 deletions(-) diff --git a/examples/models/muse_glimmer/README.md b/examples/models/muse_glimmer/README.md index d3420443d1..332134443c 100644 --- a/examples/models/muse_glimmer/README.md +++ b/examples/models/muse_glimmer/README.md @@ -38,7 +38,7 @@ bridge = AutoBridge.from_hf_pretrained( model_config = bridge.get_model_config() # Apply deployment-specific parallelism before construction. -model_config.transformer.tensor_model_parallel_size = 8 +model_config.tensor_model_parallel_size = 8 models = bridge.get_model(model_config) ``` diff --git a/src/megatron/bridge/models/common/__init__.py b/src/megatron/bridge/models/common/__init__.py index 63525c25d2..0a72385c24 100644 --- a/src/megatron/bridge/models/common/__init__.py +++ b/src/megatron/bridge/models/common/__init__.py @@ -16,6 +16,7 @@ BuildConfigT, ModelBuilder, ModelConfig, + ModelConfigOverrideMixin, ModelT, Serializable, compose_hooks, @@ -28,6 +29,7 @@ "BuildConfigT", "ModelBuilder", "ModelConfig", + "ModelConfigOverrideMixin", "ModelT", "Serializable", "compose_hooks", diff --git a/src/megatron/bridge/models/common/base.py b/src/megatron/bridge/models/common/base.py index 19e3c42454..153facb33f 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,6 +27,7 @@ ModelConfig as _MegatronModelConfig, ) +from megatron.bridge.models.transformer_config import TransformerConfig from megatron.bridge.utils.instantiate_utils import _resolve_target, _validate_target_prefix, instantiate @@ -73,6 +74,42 @@ def _from_dict(subdata: dict[str, Any], full_key: str) -> Any: 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): """Bridge compatibility wrapper for Megatron-LM model configs.""" diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index 181896b883..c45c44cc11 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -37,6 +37,7 @@ Type, TypeVar, Union, + cast, ) import torch @@ -51,6 +52,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, @@ -433,11 +435,11 @@ 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 # Leave unset unless HF export must copy nonstandard files in addition to the usual artifacts, # for example ``["*reasoning_parser.py"]``. @@ -723,12 +725,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/gpt/model_config.py b/src/megatron/bridge/models/gpt/model_config.py index ff2d679af0..191671a04f 100644 --- a/src/megatron/bridge/models/gpt/model_config.py +++ b/src/megatron/bridge/models/gpt/model_config.py @@ -15,17 +15,18 @@ """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.training.models.gpt import GPTModelConfig -from megatron.bridge.models.common.base import ModelConfig +from megatron.bridge.models.common.base import ModelConfig, ModelConfigOverrideMixin +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): """GPT model config with strict overrides and callable serialization. Outer GPT build fields and nested transformer fields keep one declared @@ -38,29 +39,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/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index 7648ab100c..64ced8aeec 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -138,7 +138,6 @@ class MuseGlimmerBridge(MegatronModelBridge): """Builder-backed bridge for the complete Muse Glimmer multimodal model.""" MODEL_CONFIG_CLASS = MuseGlimmerModelConfig - TRANSFORMER_CONFIG_CLASS = MuseGlimmerTransformerConfig @staticmethod def _validate_architecture(hf_config: PretrainedConfig) -> tuple[Any, Any]: diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py index bbbdfd05de..09006ca760 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_config.py @@ -21,7 +21,7 @@ from megatron.training.models.hybrid import HybridModelConfig -from megatron.bridge.models.common import deserialize_model_config +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 @@ -67,10 +67,11 @@ def spatial_merge_size(self) -> int: @dataclass(kw_only=True) -class MuseGlimmerModelConfig(HybridModelConfig): +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) diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 911662247e..51a5f41695 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -26,6 +26,7 @@ MuseGlimmerModelBuilder, MuseGlimmerModelConfig, MuseGlimmerTextConfig, + MuseGlimmerTransformerConfig, MuseGlimmerVisionConfig, ) from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( @@ -126,6 +127,24 @@ def test_config_conversion_uses_model_config_path_only() -> None: 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__ + + 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() diff --git a/tests/unit_tests/models/test_gpt_model_config.py b/tests/unit_tests/models/test_gpt_model_config.py index 9e25ead462..eb4a4d1ff2 100644 --- a/tests/unit_tests/models/test_gpt_model_config.py +++ b/tests/unit_tests/models/test_gpt_model_config.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import dataclass + import pytest import torch.nn.functional as F @@ -96,10 +98,45 @@ 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 BridgeGPTModelConfig.transformer_config_class is TransformerConfig + assert not hasattr(MegatronModelBridge, "TRANSFORMER_CONFIG_CLASS") def test_bridge_can_explicitly_disable_model_config() -> None: From c64a34dcb92b7ea6ceb393d7db601cf01ac75e04 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 14 Aug 2026 15:51:01 -0700 Subject: [PATCH 60/72] perf(muse-glimmer): use TE affine norms Signed-off-by: yaoyu-33 --- skills/adding-model-support/SKILL.md | 18 ++++ skills/adding-model-support/vlm-patterns.md | 8 ++ skills/parity-testing/SKILL.md | 9 ++ .../muse_glimmer/modeling_muse_glimmer.py | 89 +++++++++++++++---- .../muse_glimmer/muse_glimmer_bridge.py | 2 +- .../models/muse_glimmer/test_muse_glimmer.py | 66 ++++++++++++-- 6 files changed, 167 insertions(+), 25 deletions(-) diff --git a/skills/adding-model-support/SKILL.md b/skills/adding-model-support/SKILL.md index dc1a8c1f8e..e1fe5ed997 100644 --- a/skills/adding-model-support/SKILL.md +++ b/skills/adding-model-support/SKILL.md @@ -178,6 +178,24 @@ 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 +Transformer Engine version pinned by Megatron-Core. Use the backend-provided `TENorm` (or the +corresponding TE module) for affine LayerNorm and RMSNorm whenever TE implements the required +semantics. Preserve `eps`, parameter dtype/device initialization, sequence-parallel metadata, and +`zero_centered_gamma`; standard-gamma and zero-centered-gamma norms are not interchangeable. + +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..40dc457784 100644 --- a/skills/adding-model-support/vlm-patterns.md +++ b/skills/adding-model-support/vlm-patterns.md @@ -118,6 +118,14 @@ 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 the same Transformer Engine backend and initialization device as the +rest of the model. Do not instantiate `torch.nn.LayerNorm` unconditionally in a production TE model. +When the HF architecture specifies a weightless norm, first verify whether the pinned TE API supports +that exact parameter-free contract; 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/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 7fd9ea1e14..7e5c2c3eec 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -17,6 +17,7 @@ from __future__ import annotations import functools +from copy import copy from typing import TYPE_CHECKING, Any import torch @@ -57,7 +58,12 @@ def _rms_norm(hidden_states: Tensor, eps: float) -> Tensor: - """Apply the float32 RMS expression used by the Hugging Face reference.""" + """Apply the parameter-free RMS expression used by the HF reference. + + Transformer Engine 2.17 RMSNorm always owns an affine weight, so Muse's + scaleless Q/K, embedding, and perception norms cannot use it without + changing the architecture and checkpoint schema. + """ normalized = hidden_states.float() * torch.pow( hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + eps, -0.5, @@ -132,6 +138,39 @@ def _centered_rms_norm_cls(config: TransformerConfig) -> type[nn.Module]: return get_backend(config.transformer_impl).layer_norm(rms_norm=True) +def _standard_rms_norm(config: TransformerConfig, hidden_size: int, eps: float) -> nn.Module: + """Build a standard-gamma RMSNorm with the configured model backend.""" + if config.transformer_impl == "local": + return MuseGlimmerRMSNorm(config, hidden_size, eps=eps) + + norm_config = copy(config) + norm_config.layernorm_zero_centered_gamma = False + return get_backend(config.transformer_impl).layer_norm(rms_norm=True)( + norm_config, + hidden_size, + eps=eps, + ) + + +def _vision_layer_norm( + config: MuseGlimmerVisionModelConfig, + transformer_config: TransformerConfig, +) -> nn.Module: + """Build the vision LayerNorm with TE unless the local backend is explicit.""" + if transformer_config.transformer_impl == "local": + return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + norm_config = copy(transformer_config) + norm_config.normalization = "LayerNorm" + norm_config.layernorm_zero_centered_gamma = False + norm_config.sequence_parallel = False + return get_backend(transformer_config.transformer_impl).layer_norm()( + norm_config, + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + + class MuseGlimmerSelfAttention(SelfAttention): """MCore self-attention with Muse's residual-branch post norm.""" @@ -239,10 +278,10 @@ def customize_muse_glimmer_language_model(model: HybridModel) -> None: if hasattr(model, "embedding"): extend_instance(model.embedding, MuseGlimmerEmbeddingNormMixin) if hasattr(model, "decoder") and getattr(model.decoder, "final_norm", None) is not None: - model.decoder.final_norm = MuseGlimmerRMSNorm( + model.decoder.final_norm = _standard_rms_norm( model.config, model.config.hidden_size, - eps=model.config.layernorm_epsilon, + model.config.layernorm_epsilon, ) if hasattr(model, "output_layer"): extend_instance(model.output_layer, MuseGlimmerOutputLayerMixin) @@ -468,11 +507,11 @@ def forward(self, hidden_states: Tensor) -> Tensor: class MuseGlimmerVisionEncoderLayer(nn.Module): """Pre-normalized Muse vision transformer layer.""" - def __init__(self, config: MuseGlimmerVisionModelConfig) -> None: + def __init__(self, config: MuseGlimmerVisionModelConfig, transformer_config: TransformerConfig) -> None: super().__init__() - self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.norm1 = _vision_layer_norm(config, transformer_config) self.attn = MuseGlimmerVisionAttention(config) - self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.norm2 = _vision_layer_norm(config, transformer_config) self.mlp = MuseGlimmerVisionMLP(config) def forward( @@ -488,15 +527,23 @@ def forward( class MuseGlimmerVisionModel(nn.Module): """Muse vision tower with 3:1 window/full attention and 2x2 pixel shuffle.""" - def __init__(self, config: MuseGlimmerVisionModelConfig, *, recompute_layers: bool = False) -> None: + 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) - self.ln_pre = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) - self.layers = nn.ModuleList([MuseGlimmerVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]) - self.ln_post = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.ln_pre = _vision_layer_norm(config, transformer_config) + self.layers = nn.ModuleList( + [MuseGlimmerVisionEncoderLayer(config, transformer_config) for _ in range(config.num_hidden_layers)] + ) + self.ln_post = _vision_layer_norm(config, transformer_config) def _pixel_shuffle(self, hidden_states: Tensor, grid_thw: Tensor) -> Tensor: outputs = [] @@ -566,11 +613,16 @@ def forward(self, hidden_states: Tensor) -> Tensor: class MuseGlimmerModel(HybridModel): """Native MCore Hybrid model with the Muse vision modules attached.""" - _CENTERED_NORM_EXTRA_STATE_SUFFIXES = ( + _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", @@ -614,6 +666,7 @@ def __init__( if pre_process: self.vision_tower = MuseGlimmerVisionModel( config.vision, + config.transformer, recompute_layers=config.recompute_vision_layers, ) self.vision_adapter = MuseGlimmerVisionAdapter(config) @@ -639,12 +692,10 @@ def sharded_state_dict( ) -> dict[str, Any]: """Return a backend-stable Muse checkpoint schema. - Transformer Engine RMSNorm adds an empty ``_extra_state`` object that - the local centered RMSNorm does not expose. Muse conversion uses the - local backend so it can run on CPU, while training may use Transformer - Engine. Omitting only these empty norm objects keeps those checkpoints - interchangeable without relaxing strict loading for parameters or - stateful Transformer Engine modules. + 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) @@ -652,7 +703,7 @@ def sharded_state_dict( 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._CENTERED_NORM_EXTRA_STATE_SUFFIXES): + 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 = ( @@ -661,7 +712,7 @@ def sharded_state_dict( else extra_state_data is not None and bool(extra_state_data) ) if has_extra_state: - raise ValueError(f"Muse centered RMSNorm extra state must be empty: {key}.") + raise ValueError(f"Muse Transformer Engine norm extra state must be empty: {key}.") return sharded_state_dict def freeze( diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index 7648ab100c..b95db1aad4 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -217,7 +217,7 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerM params_dtype=params_dtype, fp16=params_dtype == torch.float16, bf16=params_dtype == torch.bfloat16, - transformer_impl="local", + transformer_impl="transformer_engine", output_multiplier=float(text_config.output_multiplier), final_logit_softcapping=float(text_config.final_logit_softcapping), ) diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 911662247e..f0ff29497d 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -17,6 +17,8 @@ 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 merge_qkv_weights, split_qkv_weights @@ -33,6 +35,8 @@ MuseGlimmerRMSNorm, MuseGlimmerVisionAttention, MuseGlimmerVisionModel, + MuseGlimmerWeightlessRMSNorm, + _standard_rms_norm, get_muse_glimmer_hybrid_stack_spec, ) from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import ( @@ -108,6 +112,7 @@ def test_config_conversion_uses_model_config_path_only() -> None: assert isinstance(model_config, MuseGlimmerModelConfig) assert isinstance(model_config, HybridModelConfig) assert model_config.get_builder_cls() is MuseGlimmerModelBuilder + assert model_config.transformer_impl == "transformer_engine" assert model_config.num_layers == 4 assert model_config.hidden_size == 64 assert model_config.ffn_hidden_size == 128 @@ -279,6 +284,7 @@ def test_full_head_gate_qkv_layout_round_trips() -> None: def test_centered_and_final_norms_use_their_native_expressions() -> None: config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer + config.transformer_impl = "local" hidden_states = torch.tensor([[0.25, -0.5, 1.0, -2.0]], dtype=torch.float32) centered = MuseGlimmerCenteredRMSNorm(config, hidden_size=4, eps=1e-5) final = MuseGlimmerRMSNorm(config, hidden_size=4, eps=1e-5) @@ -294,6 +300,45 @@ def test_centered_and_final_norms_use_their_native_expressions() -> None: torch.testing.assert_close(final(hidden_states), expected_final) +def test_transformer_engine_backend_builds_all_affine_norms_with_te() -> 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 + + centered_norm_cls = get_muse_glimmer_hybrid_stack_spec( + config + ).submodules.attention_layer.submodules.input_layernorm + centered_norm = centered_norm_cls(config, config.hidden_size, eps=config.layernorm_epsilon) + final_norm = _standard_rms_norm(config, config.hidden_size, config.layernorm_epsilon) + vision_model = MuseGlimmerVisionModel(model_config.vision, config) + + assert isinstance(centered_norm, TERMSNorm) + assert centered_norm.zero_centered_gamma is True + assert isinstance(final_norm, TERMSNorm) + assert final_norm.zero_centered_gamma is False + 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(final_norm.state_dict()) == {"weight", "_extra_state"} + assert set(vision_model.ln_pre.state_dict()) == {"weight", "bias", "_extra_state"} + + registry = MuseGlimmerBridge().mapping_registry() + assert registry.megatron_to_hf_lookup("decoder.final_norm.weight") is not None + for name, _ in vision_model.named_parameters(): + assert registry.megatron_to_hf_lookup(f"vision_tower.{name}") is not None + + +def test_weightless_rms_norm_does_not_add_a_transformer_engine_parameter() -> None: + config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer + + norm = MuseGlimmerWeightlessRMSNorm(config, hidden_size=config.kv_channels) + + assert dict(norm.named_parameters()) == {} + assert norm.state_dict() == {} + + def test_centered_norm_sharded_state_dict_identifies_tensor_parallel_replicas() -> None: config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer norm = MuseGlimmerCenteredRMSNorm(config, hidden_size=4) @@ -336,8 +381,10 @@ def test_mapping_registry_covers_complete_checkpoint() -> None: def test_tiny_vision_model_preserves_expected_token_count() -> None: - vision_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).vision - model = MuseGlimmerVisionModel(vision_config) + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + model_config.transformer_impl = "local" + vision_config = model_config.vision + model = MuseGlimmerVisionModel(vision_config, model_config.transformer) pixel_values = torch.randn(16, 24) grid_thw = torch.tensor([[1, 4, 4]]) @@ -391,8 +438,10 @@ def test_vision_attention_matches_eager_reference_with_finite_backward() -> None def test_vision_layer_recomputation_is_training_only() -> None: - vision_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).vision - model = MuseGlimmerVisionModel(vision_config, recompute_layers=True) + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + model_config.transformer_impl = "local" + vision_config = model_config.vision + model = MuseGlimmerVisionModel(vision_config, model_config.transformer, recompute_layers=True) pixel_values = torch.randn(16, 24) grid_thw = torch.tensor([[1, 4, 4]]) @@ -413,6 +462,7 @@ 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 @@ -517,6 +567,9 @@ def test_checkpoint_schema_omits_only_empty_backend_norm_state(tiny_hybrid_model "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, } @@ -534,6 +587,9 @@ def test_checkpoint_schema_omits_only_empty_backend_norm_state(tiny_hybrid_model 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) @@ -550,7 +606,7 @@ def test_checkpoint_schema_rejects_stateful_backend_norm_state(tiny_hybrid_model "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.get_pg_rank", side_effect=[3, 4], ), - pytest.raises(ValueError, match="centered RMSNorm extra state must be empty"), + pytest.raises(ValueError, match="Transformer Engine norm extra state must be empty"), ): tiny_hybrid_model.sharded_state_dict() From 923f894c09181215d38f0856d40e5ce5e8eee8ca Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 14 Aug 2026 16:11:08 -0700 Subject: [PATCH 61/72] fix(muse-glimmer): decouple TE norms from attention backend Signed-off-by: yaoyu-33 --- skills/adding-model-support/SKILL.md | 3 ++ .../muse_glimmer/modeling_muse_glimmer.py | 19 +++++---- .../muse_glimmer/muse_glimmer_bridge.py | 2 +- .../models/muse_glimmer/test_muse_glimmer.py | 40 ++++++++++++++----- 4 files changed, 43 insertions(+), 21 deletions(-) diff --git a/skills/adding-model-support/SKILL.md b/skills/adding-model-support/SKILL.md index e1fe5ed997..e27928d962 100644 --- a/skills/adding-model-support/SKILL.md +++ b/skills/adding-model-support/SKILL.md @@ -185,6 +185,9 @@ Transformer Engine version pinned by Megatron-Core. Use the backend-provided `TE corresponding TE module) for affine LayerNorm and RMSNorm whenever TE implements the required semantics. 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 need not be coupled: a local attention spec +may still use `TENorm`. Retain a semantics-equivalent local fallback when model construction must +work in an environment where Transformer Engine itself is unavailable. 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 diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 7e5c2c3eec..30fc3fe7be 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -24,6 +24,7 @@ import torch.nn as nn import torch.nn.functional as F from megatron.core import parallel_state +from megatron.core.extensions.transformer_engine import HAVE_TE, 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 @@ -132,20 +133,18 @@ def forward(self, hidden_states: Tensor) -> Tensor: def _centered_rms_norm_cls(config: TransformerConfig) -> type[nn.Module]: - """Select the backend-native centered RMSNorm when it supports Muse semantics.""" - if config.transformer_impl == "local": - return MuseGlimmerCenteredRMSNorm - return get_backend(config.transformer_impl).layer_norm(rms_norm=True) + """Use TE for affine normalization whenever it is available.""" + return TENorm if HAVE_TE else MuseGlimmerCenteredRMSNorm def _standard_rms_norm(config: TransformerConfig, hidden_size: int, eps: float) -> nn.Module: - """Build a standard-gamma RMSNorm with the configured model backend.""" - if config.transformer_impl == "local": + """Build a standard-gamma RMSNorm with TE when it is available.""" + if not HAVE_TE: return MuseGlimmerRMSNorm(config, hidden_size, eps=eps) norm_config = copy(config) norm_config.layernorm_zero_centered_gamma = False - return get_backend(config.transformer_impl).layer_norm(rms_norm=True)( + return TENorm( norm_config, hidden_size, eps=eps, @@ -156,15 +155,15 @@ def _vision_layer_norm( config: MuseGlimmerVisionModelConfig, transformer_config: TransformerConfig, ) -> nn.Module: - """Build the vision LayerNorm with TE unless the local backend is explicit.""" - if transformer_config.transformer_impl == "local": + """Build the vision LayerNorm with TE when it is available.""" + if not HAVE_TE: return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) norm_config = copy(transformer_config) norm_config.normalization = "LayerNorm" norm_config.layernorm_zero_centered_gamma = False norm_config.sequence_parallel = False - return get_backend(transformer_config.transformer_impl).layer_norm()( + return TENorm( norm_config, config.hidden_size, eps=config.layer_norm_epsilon, diff --git a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index f27acbb408..64ced8aeec 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -216,7 +216,7 @@ def hf_config_to_model_config(self, hf_config: PretrainedConfig) -> MuseGlimmerM params_dtype=params_dtype, fp16=params_dtype == torch.float16, bf16=params_dtype == torch.bfloat16, - transformer_impl="transformer_engine", + transformer_impl="local", output_multiplier=float(text_config.output_multiplier), final_logit_softcapping=float(text_config.final_logit_softcapping), ) diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 733dbeb1f6..d85f26f551 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -113,7 +113,7 @@ def test_config_conversion_uses_model_config_path_only() -> None: assert isinstance(model_config, MuseGlimmerModelConfig) assert isinstance(model_config, HybridModelConfig) assert model_config.get_builder_cls() is MuseGlimmerModelBuilder - assert model_config.transformer_impl == "transformer_engine" + 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 @@ -209,7 +209,7 @@ def test_hybrid_attention_layers_use_full_transformer_flops() -> None: @pytest.mark.parametrize( ("transformer_impl", "expected_qkv", "expected_attention", "expected_norm"), [ - ("local", ColumnParallelLinear, DotProductAttention, MuseGlimmerCenteredRMSNorm), + ("local", ColumnParallelLinear, DotProductAttention, TENorm), ("transformer_engine", TEColumnParallelLinear, TEDotProductAttention, TENorm), ], ) @@ -358,6 +358,20 @@ def test_weightless_rms_norm_does_not_add_a_transformer_engine_parameter() -> No assert norm.state_dict() == {} +def test_affine_norms_retain_a_local_fallback_without_transformer_engine() -> None: + model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) + config = model_config.transformer + + with patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False): + stack_spec = get_muse_glimmer_hybrid_stack_spec(config) + final_norm = _standard_rms_norm(config, config.hidden_size, config.layernorm_epsilon) + vision_model = MuseGlimmerVisionModel(model_config.vision, config) + + assert stack_spec.submodules.attention_layer.submodules.input_layernorm is MuseGlimmerCenteredRMSNorm + assert isinstance(final_norm, MuseGlimmerRMSNorm) + assert isinstance(vision_model.ln_pre, torch.nn.LayerNorm) + + def test_centered_norm_sharded_state_dict_identifies_tensor_parallel_replicas() -> None: config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer norm = MuseGlimmerCenteredRMSNorm(config, hidden_size=4) @@ -403,11 +417,12 @@ 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" vision_config = model_config.vision - model = MuseGlimmerVisionModel(vision_config, model_config.transformer) pixel_values = torch.randn(16, 24) grid_thw = torch.tensor([[1, 4, 4]]) - output = model(pixel_values, grid_thw) + with patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False): + model = MuseGlimmerVisionModel(vision_config, model_config.transformer) + output = model(pixel_values, grid_thw) assert output.shape == (4, 64) @@ -460,14 +475,17 @@ 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" vision_config = model_config.vision - model = MuseGlimmerVisionModel(vision_config, model_config.transformer, recompute_layers=True) pixel_values = torch.randn(16, 24) grid_thw = torch.tensor([[1, 4, 4]]) - with patch( - "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.torch_checkpoint", - side_effect=lambda function, *args, **kwargs: function(*args), - ) as checkpoint: + with ( + patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False), + 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) training_output = model(pixel_values, grid_thw) model.eval() inference_output = model(pixel_values, grid_thw) @@ -524,7 +542,9 @@ def test_builder_constructs_native_hybrid_model_on_cpu(tiny_hybrid_model: MuseGl 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.final_norm, MuseGlimmerRMSNorm) + assert isinstance(model.decoder.final_norm, TERMSNorm) + 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) From cccdcbcd82da797efe668342405cc78efb8f2cc4 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 14 Aug 2026 16:42:06 -0700 Subject: [PATCH 62/72] docs(model): verify Muse Glimmer TE norm parity Signed-off-by: yaoyu-33 --- .../muse-glimmer-30b/card.yaml | 77 +++++++++++-------- .../test_validate_card.py | 5 ++ 2 files changed, 49 insertions(+), 33 deletions(-) diff --git a/examples/model_verification_cards/muse-glimmer-30b/card.yaml b/examples/model_verification_cards/muse-glimmer-30b/card.yaml index fc29c46da7..c41b1771a5 100644 --- a/examples/model_verification_cards/muse-glimmer-30b/card.yaml +++ b/examples/model_verification_cards/muse-glimmer-30b/card.yaml @@ -9,7 +9,10 @@ summary: > 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. + 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, @@ -38,7 +41,7 @@ model: min_transformers_version: "5.15.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.07-py3 - bridge_commit: 2bf9bbb49444954f6ce4361d7435d12021c842c2 # pragma: allowlist secret + bridge_commit: 923f894c09181215d38f0856d40e5ce5e8eee8ca # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -50,7 +53,7 @@ items: --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e --megatron-path work/model-verification/muse-glimmer-30b/cpu-megatron --torch-dtype bfloat16 --overwrite - last_verified: 2026-08-12 + 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 @@ -68,7 +71,7 @@ items: --hf-revision f84ecc3a0ea984a4c04542a84269e3d065350a6e --megatron-path work/model-verification/muse-glimmer-30b/gpu-megatron --torch-dtype bfloat16 --tp 8 --overwrite - last_verified: 2026-08-12 + 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 @@ -86,7 +89,7 @@ items: --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-12 + 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, @@ -105,7 +108,7 @@ items: --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-12 + 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 @@ -127,20 +130,19 @@ items: --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-12 + 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.999896, above the 0.99 gate; maximum and - mean absolute logit differences are 0.781250 and 0.139148, respectively. + 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 - bridge_commit: c7026ab2cbaac76408c23c87d9ab9a27123603d6 # pragma: allowlist secret command: > ./scripts/inference/infer.sh --nodes 1 --gpus-per-node 8 --task vlm-generation --hf_model_path meta-models/Muse-Glimmer-30B @@ -150,7 +152,7 @@ items: --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-12 + 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 @@ -190,18 +192,19 @@ items: 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 and 42.49 model TFLOP/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. + 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: @@ -229,15 +232,16 @@ items: 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 and 94.03 model TFLOP/s/GPU. The complete 32-shard - step-100 full-model checkpoint, metadata, train state, and resolved - config are saved and reloadable. + 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: @@ -301,6 +305,7 @@ items: 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 @@ -308,9 +313,10 @@ items: 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 and 16.07 model TFLOP/s/GPU. The - complete step-100 full-model checkpoint, metadata, train state, and - resolved config are saved and reloadable. + 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: @@ -337,19 +343,21 @@ items: 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 and 283.36 - model TFLOP/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. + 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: @@ -363,6 +371,7 @@ items: 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] @@ -395,6 +404,7 @@ items: 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 @@ -402,6 +412,7 @@ items: 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 and 428.24 model - TFLOP/s/GPU. The persisted post-setup config records the exact runtime, - and the process exits successfully without checkpoint output. + 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/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 61d17ebca3..135b6f1b1e 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 @@ -40,6 +40,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), From 9710d94e1573aca1117b16677b1e9c413b5f948d Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 15 Aug 2026 10:13:32 -0700 Subject: [PATCH 63/72] fix(muse-glimmer): declare recipe environments inline Signed-off-by: yaoyu-33 --- .../muse_glimmer/modeling_muse_glimmer.py | 2 +- .../muse_glimmer/h100/muse_glimmer.py | 31 ++++++++++++++----- .../recipes/muse_glimmer/h100/muse_glimmer.py | 13 +++++++- 3 files changed, 37 insertions(+), 9 deletions(-) diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 30fc3fe7be..bd6a801a7c 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -61,7 +61,7 @@ def _rms_norm(hidden_states: Tensor, eps: float) -> Tensor: """Apply the parameter-free RMS expression used by the HF reference. - Transformer Engine 2.17 RMSNorm always owns an affine weight, so Muse's + The pinned Transformer Engine RMSNorm always owns an affine weight, so Muse's scaleless Q/K, embedding, and perception norms cannot use it without changing the architecture and checkpoint schema. """ 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 index 985823bad7..a4f5cf6124 100644 --- a/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/perf_recipes/muse_glimmer/h100/muse_glimmer.py @@ -53,6 +53,13 @@ def _muse_glimmer_30b_pretrain_32gpu_h100_config() -> ConfigContainer: 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. @@ -63,23 +70,33 @@ def _muse_glimmer_30b_pretrain_32gpu_h100_config() -> ConfigContainer: "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_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") - 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 diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 613dcfb758..126216af5b 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -92,7 +92,6 @@ def _apply_model_and_data( cfg.train.micro_batch_size = 1 cfg.train.manual_gc = True cfg.train.manual_gc_interval = 10 - cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} def _set_optimizer(cfg: ConfigContainer, *, max_lr: float, warmup_iters: int) -> None: @@ -130,6 +129,9 @@ def muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() -> ConfigConta 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 @@ -153,6 +155,9 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_config() -> ConfigContainer: cfg.checkpoint.load = None cfg.checkpoint.save_optim = False cfg.checkpoint.save_rng = False + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } return cfg @@ -177,6 +182,9 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContaine cfg.checkpoint.load = None cfg.checkpoint.save_optim = False cfg.checkpoint.save_rng = False + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } return cfg @@ -207,6 +215,9 @@ def muse_glimmer_30b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora" _set_optimizer(cfg, max_lr=1e-4, warmup_iters=10) cfg.checkpoint.save_interval = 100 cfg.checkpoint.load = None + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } return cfg From feed3a55bfb3b9f1576ab8a4d7b089e82b5a895a Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 15 Aug 2026 10:19:01 -0700 Subject: [PATCH 64/72] docs(skill): clarify Transformer Engine API lookup Signed-off-by: yaoyu-33 --- skills/adding-model-support/SKILL.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/skills/adding-model-support/SKILL.md b/skills/adding-model-support/SKILL.md index e27928d962..ce2fdf20a0 100644 --- a/skills/adding-model-support/SKILL.md +++ b/skills/adding-model-support/SKILL.md @@ -180,11 +180,12 @@ than adding conditionals to shared code. #### Prefer Transformer Engine normalization -Before implementing a model-specific normalization with raw PyTorch operations, inspect the -Transformer Engine version pinned by Megatron-Core. Use the backend-provided `TENorm` (or the -corresponding TE module) for affine LayerNorm and RMSNorm whenever TE implements the required -semantics. Preserve `eps`, parameter dtype/device initialization, sequence-parallel metadata, and -`zero_centered_gamma`; standard-gamma and zero-centered-gamma norms are not interchangeable. +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. Use `TENorm` (or the corresponding TE module) for affine LayerNorm and RMSNorm +whenever that pinned TE implementation provides the required semantics. 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 need not be coupled: a local attention spec may still use `TENorm`. Retain a semantics-equivalent local fallback when model construction must work in an environment where Transformer Engine itself is unavailable. From 9a5bfe535c5694b30150bd18b9f0fbee0b644c17 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 15 Aug 2026 10:27:22 -0700 Subject: [PATCH 65/72] fix(models): support updated Transformers config APIs Signed-off-by: yaoyu-33 --- src/megatron/bridge/models/common/base.py | 7 +-- .../bridge/models/gemma/gemma4_bridge.py | 59 ++++++++++++++++--- .../qwen_omni/modeling_qwen3_omni/rope.py | 7 ++- .../modeling_qwen3_omni/thinker_model.py | 10 +++- .../models/gemma/test_gemma4_bridge.py | 2 + .../models/glm_moe_dsa/test_glm5_bridge.py | 2 +- 6 files changed, 67 insertions(+), 20 deletions(-) diff --git a/src/megatron/bridge/models/common/base.py b/src/megatron/bridge/models/common/base.py index 153facb33f..9c645d204b 100644 --- a/src/megatron/bridge/models/common/base.py +++ b/src/megatron/bridge/models/common/base.py @@ -121,9 +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.""" - result = deserialize_model_config(data) - if not isinstance(result, cls): - raise ValueError(f"Cannot deserialize: expected {cls.__name__}, got {type(result).__name__}") - return result + return deserialize_model_config(data) diff --git a/src/megatron/bridge/models/gemma/gemma4_bridge.py b/src/megatron/bridge/models/gemma/gemma4_bridge.py index a258a76f24..fd0f129fbd 100644 --- a/src/megatron/bridge/models/gemma/gemma4_bridge.py +++ b/src/megatron/bridge/models/gemma/gemma4_bridge.py @@ -98,6 +98,23 @@ 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) -> 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. The + # global value is still the correct fallback for the local-attention config. + return getattr(hf_config, "__dict__", {}).get(field_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 +198,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 +235,17 @@ 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, + "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, - "num_global_key_value_heads", - num_query_groups, + "full_attention", + "num_key_value_heads", + getattr(hf_config, "num_global_key_value_heads", num_query_groups), ) self._dense_num_attention_heads = num_attention_heads @@ -236,8 +264,8 @@ 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), num_global_query_groups=num_global_query_groups, seq_length=hf_config.max_position_embeddings, vocab_size=hf_config.vocab_size, @@ -260,6 +288,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 +302,18 @@ 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) + provider.num_global_key_value_heads = _attention_config_value( + hf_config, + "full_attention", + "num_key_value_heads", + getattr(hf_config, "num_global_key_value_heads", 2), + ) provider.attention_k_eq_v = getattr(hf_config, "attention_k_eq_v", False) rope_params = getattr(hf_config, "rope_parameters", {}) @@ -320,12 +359,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), 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/tests/unit_tests/models/gemma/test_gemma4_bridge.py b/tests/unit_tests/models/gemma/test_gemma4_bridge.py index 21d7d58b71..3a78d05513 100644 --- a/tests/unit_tests/models/gemma/test_gemma4_bridge.py +++ b/tests/unit_tests/models/gemma/test_gemma4_bridge.py @@ -219,6 +219,8 @@ 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 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 From a1ee7bf63b28ca19dc73b7256d02cf5830b02e4f Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 17 Aug 2026 10:26:33 -0700 Subject: [PATCH 66/72] fix(models): repair Transformers 5.15 compatibility Signed-off-by: yaoyu-33 --- .../bridge/models/gemma/gemma4_bridge.py | 51 +++++++++++++++---- .../models/gemma_vl/gemma4_vl_bridge.py | 25 +++++++-- .../models/gemma/test_gemma4_bridge.py | 2 + .../recipes/test_muse_glimmer_recipes.py | 2 + 4 files changed, 66 insertions(+), 14 deletions(-) diff --git a/src/megatron/bridge/models/gemma/gemma4_bridge.py b/src/megatron/bridge/models/gemma/gemma4_bridge.py index fd0f129fbd..0295413ec8 100644 --- a/src/megatron/bridge/models/gemma/gemma4_bridge.py +++ b/src/megatron/bridge/models/gemma/gemma4_bridge.py @@ -98,7 +98,14 @@ 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) -> Any: +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) @@ -110,9 +117,10 @@ def _attention_config_value(hf_config: Any, layer_type: str, field_name: str, de else: return getattr(per_layer_config[layer_index], field_name, default) - # Transformers 5.15 rejects direct reads of fields that vary by layer. The - # global value is still the correct fallback for the local-attention config. - return getattr(hf_config, "__dict__", {}).get(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]: @@ -245,7 +253,8 @@ def _build_dense_provider(self, hf_config) -> Gemma4DenseProvider: hf_config, "full_attention", "num_key_value_heads", - getattr(hf_config, "num_global_key_value_heads", num_query_groups), + num_query_groups, + legacy_field_name="num_global_key_value_heads", ) self._dense_num_attention_heads = num_attention_heads @@ -265,7 +274,13 @@ def _build_dense_provider(self, hf_config) -> Gemma4DenseProvider: num_attention_heads=num_attention_heads, num_query_groups=num_query_groups, 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), + 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, @@ -307,12 +322,19 @@ def _build_moe_provider(self, hf_config) -> Gemma4ModelProvider: provider.kv_channels = head_dim provider.qk_layernorm = True - provider.global_head_dim = _attention_config_value(hf_config, "full_attention", "head_dim", 512) + 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", - getattr(hf_config, "num_global_key_value_heads", 2), + 2, + legacy_field_name="num_global_key_value_heads", ) provider.attention_k_eq_v = getattr(hf_config, "attention_k_eq_v", False) @@ -447,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/tests/unit_tests/models/gemma/test_gemma4_bridge.py b/tests/unit_tests/models/gemma/test_gemma4_bridge.py index 3a78d05513..8ce379d2d0 100644 --- a/tests/unit_tests/models/gemma/test_gemma4_bridge.py +++ b/tests/unit_tests/models/gemma/test_gemma4_bridge.py @@ -221,6 +221,8 @@ def test_runtime_preserves_nonperiodic_attention_schedule(self, bridge): 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/recipes/test_muse_glimmer_recipes.py b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py index bcdfe87a61..93fecf530e 100644 --- a/tests/unit_tests/recipes/test_muse_glimmer_recipes.py +++ b/tests/unit_tests/recipes/test_muse_glimmer_recipes.py @@ -149,6 +149,8 @@ def test_muse_glimmer_performance_recipe_is_dense_decoder_only() -> None: "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, From fb901bd3f1e0efdb7f31f8c82692b44ec3d46b7e Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 17 Aug 2026 15:14:14 -0700 Subject: [PATCH 67/72] refactor(muse-glimmer): use native normalization modules Signed-off-by: yaoyu-33 --- skills/adding-model-support/SKILL.md | 22 +- skills/adding-model-support/vlm-patterns.md | 12 +- .../muse_glimmer/modeling_muse_glimmer.py | 218 +++++------------- .../muse_glimmer/muse_glimmer_bridge.py | 3 +- .../models/muse_glimmer/test_muse_glimmer.py | 135 +++-------- 5 files changed, 123 insertions(+), 267 deletions(-) diff --git a/skills/adding-model-support/SKILL.md b/skills/adding-model-support/SKILL.md index ce2fdf20a0..73bccd145c 100644 --- a/skills/adding-model-support/SKILL.md +++ b/skills/adding-model-support/SKILL.md @@ -182,13 +182,21 @@ than adding conditionals to shared code. 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. Use `TENorm` (or the corresponding TE module) for affine LayerNorm and RMSNorm -whenever that pinned TE implementation provides the required semantics. 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 need not be coupled: a local attention spec -may still use `TENorm`. Retain a semantics-equivalent local fallback when model construction must -work in an environment where Transformer Engine itself is unavailable. +`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 diff --git a/skills/adding-model-support/vlm-patterns.md b/skills/adding-model-support/vlm-patterns.md index 40dc457784..f0fbf4061f 100644 --- a/skills/adding-model-support/vlm-patterns.md +++ b/skills/adding-model-support/vlm-patterns.md @@ -121,10 +121,14 @@ from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM ### Normalization backend Thread the language model's `TransformerConfig` into custom Megatron vision modules so affine vision -LayerNorm/RMSNorm instances use the same Transformer Engine backend and initialization device as the -rest of the model. Do not instantiate `torch.nn.LayerNorm` unconditionally in a production TE model. -When the HF architecture specifies a weightless norm, first verify whether the pinned TE API supports -that exact parameter-free contract; do not add checkpoint parameters to simulate it. +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) diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index bd6a801a7c..86c1dc435d 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -17,14 +17,12 @@ from __future__ import annotations import functools -from copy import copy from typing import TYPE_CHECKING, Any import torch import torch.nn as nn import torch.nn.functional as F -from megatron.core import parallel_state -from megatron.core.extensions.transformer_engine import HAVE_TE, TENorm +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 @@ -37,10 +35,11 @@ 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, make_sharded_tensors_for_checkpoint +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 ( @@ -58,124 +57,22 @@ from megatron.core.packed_seq_params import PackedSeqParams -def _rms_norm(hidden_states: Tensor, eps: float) -> Tensor: - """Apply the parameter-free RMS expression used by the HF reference. - - The pinned Transformer Engine RMSNorm always owns an affine weight, so Muse's - scaleless Q/K, embedding, and perception norms cannot use it without - changing the architecture and checkpoint schema. - """ - normalized = hidden_states.float() * torch.pow( - hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + eps, - -0.5, - ) - return normalized.to(hidden_states.dtype) - - -def _centered_rms_norm(hidden_states: Tensor, eps: float) -> Tensor: - """Apply the rsqrt-based expression used by Muse's centered layer norms.""" - normalized = hidden_states.float() * torch.rsqrt(hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + eps) - return normalized.to(hidden_states.dtype) - - -class MuseGlimmerCenteredRMSNorm(nn.Module): - """Gemma2-style RMSNorm whose checkpoint weight is centered around zero.""" - - def __init__(self, config: TransformerConfig, hidden_size: int, eps: float = 1e-5) -> None: - super().__init__() - self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=config.params_dtype)) - setattr(self.weight, "sequence_parallel", bool(config.sequence_parallel)) - self.eps = eps - - def forward(self, hidden_states: Tensor) -> Tensor: - return _centered_rms_norm(hidden_states, self.eps) * (1.0 + self.weight.float()).to(hidden_states.dtype) - - def sharded_state_dict( - self, - prefix: str = "", - sharded_offsets: tuple[tuple[int, int, int], ...] = (), - metadata: dict[str, Any] | None = None, - ) -> dict[str, Any]: - """Represent the replicated weight with a distinct writer on each TP rank.""" - metadata = ensure_metadata_has_dp_cp_group(metadata) - return make_sharded_tensors_for_checkpoint( - self.state_dict(keep_vars=True), - prefix, - sharded_offsets=sharded_offsets, - tp_group=parallel_state.get_tensor_model_parallel_group(), - dp_cp_group=metadata["dp_cp_group"], - ) - - -class MuseGlimmerRMSNorm(nn.Module): - """Learnable RMSNorm used after the final Muse decoder layer.""" - - def __init__(self, config: TransformerConfig, hidden_size: int, eps: float = 1e-5) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=config.params_dtype)) - setattr(self.weight, "sequence_parallel", bool(config.sequence_parallel)) - self.eps = eps - - def forward(self, hidden_states: Tensor) -> Tensor: - return _rms_norm(hidden_states, self.eps) * self.weight.float().to(hidden_states.dtype) - - -class MuseGlimmerWeightlessRMSNorm(nn.Module): - """Scaleless RMSNorm used for Muse Q/K, embeddings, and perception features.""" - - def __init__(self, config: TransformerConfig, hidden_size: int, eps: float = 1e-5) -> None: - super().__init__() - del hidden_size - self.eps = eps - - def forward(self, hidden_states: Tensor) -> Tensor: - return _rms_norm(hidden_states, self.eps) - - -def _centered_rms_norm_cls(config: TransformerConfig) -> type[nn.Module]: - """Use TE for affine normalization whenever it is available.""" - return TENorm if HAVE_TE else MuseGlimmerCenteredRMSNorm - - -def _standard_rms_norm(config: TransformerConfig, hidden_size: int, eps: float) -> nn.Module: - """Build a standard-gamma RMSNorm with TE when it is available.""" - if not HAVE_TE: - return MuseGlimmerRMSNorm(config, hidden_size, eps=eps) - - norm_config = copy(config) - norm_config.layernorm_zero_centered_gamma = False - return TENorm( - norm_config, - hidden_size, - eps=eps, - ) - - -def _vision_layer_norm( - config: MuseGlimmerVisionModelConfig, - transformer_config: TransformerConfig, -) -> nn.Module: - """Build the vision LayerNorm with TE when it is available.""" - if not HAVE_TE: - return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) - - norm_config = copy(transformer_config) - norm_config.normalization = "LayerNorm" - norm_config.layernorm_zero_centered_gamma = False - norm_config.sequence_parallel = False - return TENorm( - norm_config, - config.hidden_size, - eps=config.layer_norm_epsilon, - ) - - 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) - self.post_layernorm = _centered_rms_norm_cls(config)( + 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)), @@ -203,7 +100,7 @@ def __init__( ffn_hidden_size=ffn_hidden_size, **kwargs, ) - self.post_layernorm = _centered_rms_norm_cls(config)( + self.post_layernorm = TENorm( config, config.hidden_size, eps=float(getattr(config, "post_norm_epsilon", 1e-8)), @@ -217,14 +114,13 @@ def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: 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) - centered_rms_norm = _centered_rms_norm_cls(config) return ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( attention_layer=ModuleSpec( module=TransformerLayer, submodules=TransformerLayerSubmodules( - input_layernorm=centered_rms_norm, + input_layernorm=TENorm, self_attention=ModuleSpec( module=MuseGlimmerSelfAttention, params={"attn_mask_type": AttnMaskType.causal}, @@ -232,12 +128,10 @@ def get_muse_glimmer_hybrid_stack_spec(config: TransformerConfig) -> ModuleSpec: linear_qkv=backend.column_parallel_linear(), core_attention=backend.core_attention(), linear_proj=backend.row_parallel_linear(), - q_layernorm=MuseGlimmerWeightlessRMSNorm, - k_layernorm=MuseGlimmerWeightlessRMSNorm, ), ), self_attn_bda=get_bias_dropout_add, - pre_mlp_layernorm=centered_rms_norm, + pre_mlp_layernorm=TENorm, mlp=functools.partial( MuseGlimmerMLP.as_mlp_submodule, submodules=MLPSubmodules( @@ -252,14 +146,6 @@ def get_muse_glimmer_hybrid_stack_spec(config: TransformerConfig) -> ModuleSpec: ) -class MuseGlimmerEmbeddingNormMixin(nn.Module): - """Apply the weightless Muse RMSNorm after word embedding lookup.""" - - def forward(self, *args: Any, **kwargs: Any) -> Tensor: - embeddings = super().forward(*args, **kwargs) - return _rms_norm(embeddings, self.config.layernorm_epsilon) - - class MuseGlimmerOutputLayerMixin(nn.Module): """Apply the Muse logit multiplier followed by tanh soft-capping.""" @@ -272,20 +158,6 @@ def forward(self, *args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | None]: return output, bias -def customize_muse_glimmer_language_model(model: HybridModel) -> None: - """Install Muse embedding, final-norm, and logit behavior on a Hybrid model.""" - if hasattr(model, "embedding"): - extend_instance(model.embedding, MuseGlimmerEmbeddingNormMixin) - if hasattr(model, "decoder") and getattr(model.decoder, "final_norm", None) is not None: - model.decoder.final_norm = _standard_rms_norm( - model.config, - model.config.hidden_size, - model.config.layernorm_epsilon, - ) - if hasattr(model, "output_layer"): - extend_instance(model.output_layer, MuseGlimmerOutputLayerMixin) - - 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) @@ -508,9 +380,24 @@ class MuseGlimmerVisionEncoderLayer(nn.Module): def __init__(self, config: MuseGlimmerVisionModelConfig, transformer_config: TransformerConfig) -> None: super().__init__() - self.norm1 = _vision_layer_norm(config, transformer_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.norm1 = TELayerNorm(**norm_kwargs) self.attn = MuseGlimmerVisionAttention(config) - self.norm2 = _vision_layer_norm(config, transformer_config) + self.norm2 = TELayerNorm(**norm_kwargs) self.mlp = MuseGlimmerVisionMLP(config) def forward( @@ -538,11 +425,26 @@ def __init__( self.recompute_layers = recompute_layers self.patch_embedder = MuseGlimmerVisionPatchEmbedder(config) self.rotary_emb = MuseGlimmerVisionRotaryEmbedding(config) - self.ln_pre = _vision_layer_norm(config, transformer_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 = _vision_layer_norm(config, transformer_config) + self.ln_post = TELayerNorm(**norm_kwargs) def _pixel_shuffle(self, hidden_states: Tensor, grid_thw: Tensor) -> Tensor: outputs = [] @@ -661,7 +563,8 @@ def __init__( vp_stage=vp_stage, ) self.model_config = config - customize_muse_glimmer_language_model(self) + if hasattr(self, "output_layer"): + extend_instance(self.output_layer, MuseGlimmerOutputLayerMixin) if pre_process: self.vision_tower = MuseGlimmerVisionModel( config.vision, @@ -743,7 +646,11 @@ def freeze( 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 _rms_norm(features, self.config.layernorm_epsilon) + 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: @@ -783,6 +690,11 @@ def forward( ) 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: @@ -835,11 +747,7 @@ def forward( __all__ = [ - "MuseGlimmerCenteredRMSNorm", "MuseGlimmerModel", - "MuseGlimmerRMSNorm", "MuseGlimmerVisionModel", - "MuseGlimmerWeightlessRMSNorm", - "customize_muse_glimmer_language_model", "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 index 64ced8aeec..16b3469809 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -31,6 +31,7 @@ GatedMLPMapping, MegatronParamMapping, ReplicatedMapping, + RMSNorm2ZeroCenteredRMSNormMapping, merge_qkv_weights, split_qkv_weights, ) @@ -352,7 +353,7 @@ def mapping_registry(self) -> MegatronMappingRegistry: megatron_param="output_layer.weight", hf_param="lm_head.weight", ), - ReplicatedMapping( + RMSNorm2ZeroCenteredRMSNormMapping( megatron_param="decoder.final_norm.weight", hf_param=f"{text_prefix}.norm.weight", ), diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index d85f26f551..68dc6dca74 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -21,7 +21,11 @@ from transformer_engine.pytorch import RMSNorm as TERMSNorm from megatron.bridge import AutoBridge -from megatron.bridge.models.conversion.param_mapping import merge_qkv_weights, split_qkv_weights +from megatron.bridge.models.conversion.param_mapping import ( + RMSNorm2ZeroCenteredRMSNormMapping, + merge_qkv_weights, + split_qkv_weights, +) from megatron.bridge.models.muse_glimmer import ( MuseGlimmerConfig, MuseGlimmerModel, @@ -32,12 +36,8 @@ MuseGlimmerVisionConfig, ) from megatron.bridge.models.muse_glimmer.modeling_muse_glimmer import ( - MuseGlimmerCenteredRMSNorm, - MuseGlimmerRMSNorm, MuseGlimmerVisionAttention, MuseGlimmerVisionModel, - MuseGlimmerWeightlessRMSNorm, - _standard_rms_norm, get_muse_glimmer_hybrid_stack_spec, ) from megatron.bridge.models.muse_glimmer.muse_glimmer_bridge import ( @@ -301,105 +301,35 @@ def test_full_head_gate_qkv_layout_round_trips() -> None: torch.testing.assert_close(restored_value, value) -def test_centered_and_final_norms_use_their_native_expressions() -> None: - config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer - config.transformer_impl = "local" - hidden_states = torch.tensor([[0.25, -0.5, 1.0, -2.0]], dtype=torch.float32) - centered = MuseGlimmerCenteredRMSNorm(config, hidden_size=4, eps=1e-5) - final = MuseGlimmerRMSNorm(config, hidden_size=4, eps=1e-5) - with torch.no_grad(): - centered.weight.fill_(0.25) - final.weight.fill_(1.25) - - mean_square = hidden_states.pow(2).mean(dim=-1, keepdim=True) - expected_centered = hidden_states * torch.rsqrt(mean_square + 1e-5) * 1.25 - expected_final = hidden_states * torch.pow(mean_square + 1e-5, -0.5) * 1.25 - - torch.testing.assert_close(centered(hidden_states), expected_centered) - torch.testing.assert_close(final(hidden_states), expected_final) - - -def test_transformer_engine_backend_builds_all_affine_norms_with_te() -> None: +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 - centered_norm_cls = get_muse_glimmer_hybrid_stack_spec( - config - ).submodules.attention_layer.submodules.input_layernorm - centered_norm = centered_norm_cls(config, config.hidden_size, eps=config.layernorm_epsilon) - final_norm = _standard_rms_norm(config, config.hidden_size, config.layernorm_epsilon) + 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(final_norm, TERMSNorm) - assert final_norm.zero_centered_gamma is False 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(final_norm.state_dict()) == {"weight", "_extra_state"} assert set(vision_model.ln_pre.state_dict()) == {"weight", "bias", "_extra_state"} registry = MuseGlimmerBridge().mapping_registry() - assert registry.megatron_to_hf_lookup("decoder.final_norm.weight") is not None + 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_weightless_rms_norm_does_not_add_a_transformer_engine_parameter() -> None: - config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer - - norm = MuseGlimmerWeightlessRMSNorm(config, hidden_size=config.kv_channels) - - assert dict(norm.named_parameters()) == {} - assert norm.state_dict() == {} - - -def test_affine_norms_retain_a_local_fallback_without_transformer_engine() -> None: - model_config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()) - config = model_config.transformer - - with patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False): - stack_spec = get_muse_glimmer_hybrid_stack_spec(config) - final_norm = _standard_rms_norm(config, config.hidden_size, config.layernorm_epsilon) - vision_model = MuseGlimmerVisionModel(model_config.vision, config) - - assert stack_spec.submodules.attention_layer.submodules.input_layernorm is MuseGlimmerCenteredRMSNorm - assert isinstance(final_norm, MuseGlimmerRMSNorm) - assert isinstance(vision_model.ln_pre, torch.nn.LayerNorm) - - -def test_centered_norm_sharded_state_dict_identifies_tensor_parallel_replicas() -> None: - config = MuseGlimmerBridge().hf_config_to_model_config(_tiny_hf_config()).transformer - norm = MuseGlimmerCenteredRMSNorm(config, hidden_size=4) - tp_group = object() - dp_cp_group = object() - - with ( - patch( - "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.parallel_state.get_tensor_model_parallel_group", - return_value=tp_group, - ), - patch( - "megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.make_sharded_tensors_for_checkpoint", - return_value={"norm.weight": object()}, - ) as make_sharded, - ): - result = norm.sharded_state_dict(prefix="norm.", metadata={"dp_cp_group": dp_cp_group}) - - assert set(result) == {"norm.weight"} - make_sharded.assert_called_once_with( - norm.state_dict(keep_vars=True), - "norm.", - sharded_offsets=(), - tp_group=tp_group, - dp_cp_group=dp_cp_group, - ) - - def test_mapping_registry_covers_complete_checkpoint() -> None: registry = MuseGlimmerBridge().mapping_registry() @@ -413,16 +343,17 @@ def test_mapping_registry_covers_complete_checkpoint() -> None: 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) - grid_thw = torch.tensor([[1, 4, 4]]) + 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.HAVE_TE", False): - model = MuseGlimmerVisionModel(vision_config, model_config.transformer) - output = model(pixel_values, grid_thw) + model = MuseGlimmerVisionModel(vision_config, model_config.transformer).cuda() + output = model(pixel_values, grid_thw) assert output.shape == (4, 64) @@ -471,21 +402,20 @@ def test_vision_attention_matches_eager_reference_with_finite_backward() -> None ) +@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) - grid_thw = torch.tensor([[1, 4, 4]]) - - with ( - patch("megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.HAVE_TE", False), - 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) + 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) @@ -542,7 +472,12 @@ def test_builder_constructs_native_hybrid_model_on_cpu(tiny_hybrid_model: MuseGl 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 From f1a9bd81911c5699b2126883f0470a2491e776c6 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 17 Aug 2026 17:15:02 -0700 Subject: [PATCH 68/72] test(muse-glimmer): restore attention backend environment Signed-off-by: yaoyu-33 --- .../models/muse_glimmer/test_muse_glimmer.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 68dc6dca74..509b4b8604 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -5,6 +5,7 @@ from __future__ import annotations import math +import os from collections.abc import Iterator from types import SimpleNamespace from unittest.mock import patch @@ -49,6 +50,20 @@ 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( From 6633d4748b7a951d20c982907340b53824dc696d Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 17 Aug 2026 17:29:46 -0700 Subject: [PATCH 69/72] docs(muse-glimmer): explain weightless QK norms Signed-off-by: yaoyu-33 --- .../bridge/models/muse_glimmer/modeling_muse_glimmer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py index 86c1dc435d..bdac50530b 100644 --- a/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py +++ b/src/megatron/bridge/models/muse_glimmer/modeling_muse_glimmer.py @@ -62,6 +62,8 @@ class MuseGlimmerSelfAttention(SelfAttention): 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, From 4719ba0cf353f18cd9288094367ffced21cc9468 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 24 Aug 2026 15:25:32 -0700 Subject: [PATCH 70/72] fix(conversion): preserve legacy model construction paths Signed-off-by: yaoyu-33 --- .../compare_hf_and_megatron/compare.py | 4 +- scripts/conversion/gpu_backend.py | 2 +- scripts/inference/vlm_generation.py | 2 +- .../bridge/models/conversion/auto_bridge.py | 2 +- .../bridge/models/conversion/model_bridge.py | 6 +++ .../bridge/models/conversion/utils.py | 7 ++-- .../muse_glimmer/muse_glimmer_bridge.py | 1 + src/megatron/bridge/training/checkpointing.py | 10 +++-- .../test_nemotron_omni_conversion.py | 13 +++++- .../conversion/launcher/test_gpu_backend.py | 2 + .../models/muse_glimmer/test_muse_glimmer.py | 1 + tests/unit_tests/models/test_auto_bridge.py | 12 +++--- .../models/test_conversion_utils.py | 22 +++++++++- .../models/test_gpt_model_config.py | 2 + .../scripts/inference/test_vlm_generation.py | 10 +++++ .../unit_tests/test_compare_mask_handling.py | 42 ++++++++++++++++++- .../unit_tests/training/test_checkpointing.py | 15 +++++++ 17 files changed, 131 insertions(+), 22 deletions(-) diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 051f3ab20b..ab85a57ed6 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -716,7 +716,7 @@ def _load_megatron_model(args): ), **_hf_revision_kwargs(args.hf_revision), ) - if getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None: + 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 @@ -756,7 +756,7 @@ def _load_megatron_model(args): ), **_hf_revision_kwargs(args.hf_revision), ) - if getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None: + 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 diff --git a/scripts/conversion/gpu_backend.py b/scripts/conversion/gpu_backend.py index 922e1ccd0b..b56006d12b 100644 --- a/scripts/conversion/gpu_backend.py +++ b/scripts/conversion/gpu_backend.py @@ -178,7 +178,7 @@ def _configure_model_provider( def _uses_model_builder(bridge: AutoBridge) -> bool: """Return whether the selected bridge supports native builder construction.""" - return getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None + return getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False) def _configure_model_config( diff --git a/scripts/inference/vlm_generation.py b/scripts/inference/vlm_generation.py index 02c9b20105..482c6fb1bd 100644 --- a/scripts/inference/vlm_generation.py +++ b/scripts/inference/vlm_generation.py @@ -151,7 +151,7 @@ def _hf_revision_kwargs(revision: str | None) -> dict[str, str]: def _uses_model_builder(bridge: AutoBridge) -> bool: """Return whether the selected bridge supports native builder construction.""" - return getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None + return getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False) def _enable_deterministic_execution() -> None: diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 9cb39d231b..433bb7aa0e 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -1484,7 +1484,7 @@ def import_ckpt( with model_context: # Prefer the native ModelConfig/ModelBuilder path for migrated model # families while preserving the provider path for legacy bridges. - if bridge._model_bridge.MODEL_CONFIG_CLASS is not None: + 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( diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index 884b9db3f5..8c8630d554 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -441,6 +441,12 @@ def mapping_registry(self) -> MegatronMappingRegistry: # transformer config class. MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = BridgeGPTModelConfig + # 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"]``. ADDITIONAL_FILE_PATTERNS = None 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/muse_glimmer/muse_glimmer_bridge.py b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py index 16b3469809..61097b3595 100644 --- a/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py +++ b/src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py @@ -139,6 +139,7 @@ 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]: diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 66053d6dbb..df7ca62744 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -177,8 +177,12 @@ class _CpuTorchDistSaveShardedStrategy(TorchDistSaveShardedStrategy): """Run MCore's synchronous torch-dist writer without a CUDA staging barrier.""" @staticmethod - def _run_cpu_finalize(finalize_fn: Callable[[], None]) -> None: - """Make MCore's failure-status tensor use the CPU for a Gloo save.""" + 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") @@ -204,7 +208,7 @@ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> No preload_fn=partial(preload_fn.func, *bound.args, **bound.kwargs), ) async_request = async_request._replace( - finalize_fns=[partial(self._run_cpu_finalize, finalize_fn) for finalize_fn in async_request.finalize_fns], + finalize_fns=[partial(self._run_finalize, finalize_fn) for finalize_fn in async_request.finalize_fns], ) async_request.execute_sync() 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..954b01a2e8 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,15 @@ # 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"] +# Transformers 5.15 serializes the same hybrid blocks under their canonical names. +_SERIALIZED_LLM_LAYER_TYPES = [ + "linear_attention", + "moe", + "linear_attention", + "moe", + "linear_attention", + "full_attention", +] _LLM_OVERRIDES = { "hybrid_override_pattern": "MEMEM*", "layer_types": _LLM_LAYER_TYPES, @@ -220,5 +229,5 @@ def test_nemotron_omni_conversion_roundtrip(self, nemotron_omni_toy_model_path, 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 + assert saved_config["llm_config"]["layer_types"] == _SERIALIZED_LLM_LAYER_TYPES + assert saved_config["llm_config"]["layers_block_type"] == _SERIALIZED_LLM_LAYER_TYPES diff --git a/tests/unit_tests/conversion/launcher/test_gpu_backend.py b/tests/unit_tests/conversion/launcher/test_gpu_backend.py index d1b6423b17..2e740a822b 100644 --- a/tests/unit_tests/conversion/launcher/test_gpu_backend.py +++ b/tests/unit_tests/conversion/launcher/test_gpu_backend.py @@ -115,6 +115,7 @@ 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"} @@ -191,6 +192,7 @@ def test_import_uses_builder_for_migrated_model(self, cli, monkeypatch): class BuilderModelBridge(_FakeModelBridge): MODEL_CONFIG_CLASS = object + USE_MODEL_CONFIG_FOR_CONVERSION = True class FakeBridge: _model_bridge = BuilderModelBridge() diff --git a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py index 509b4b8604..a9e79505dd 100644 --- a/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py +++ b/tests/unit_tests/models/muse_glimmer/test_muse_glimmer.py @@ -152,6 +152,7 @@ def test_model_config_owns_transformer_type_and_routes_flat_overrides() -> None: 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 diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index 124a853954..1ae372881d 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -1887,7 +1887,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.MODEL_CONFIG_CLASS = None + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = False mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model @@ -1919,7 +1919,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.MODEL_CONFIG_CLASS = None + mock_bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION = False mock_megatron_model = [Mock()] mock_bridge.to_megatron_model.return_value = mock_megatron_model @@ -1964,7 +1964,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.MODEL_CONFIG_CLASS = None + 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() @@ -1997,7 +1997,7 @@ def test_import_ckpt_with_low_memory_save( 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.MODEL_CONFIG_CLASS = object + 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 @@ -2038,7 +2038,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.MODEL_CONFIG_CLASS = None + 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 = {} @@ -2062,7 +2062,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.MODEL_CONFIG_CLASS = None + 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_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 c60fcaa1b3..24f94559d9 100644 --- a/tests/unit_tests/models/test_gpt_model_config.py +++ b/tests/unit_tests/models/test_gpt_model_config.py @@ -141,6 +141,8 @@ 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") diff --git a/tests/unit_tests/scripts/inference/test_vlm_generation.py b/tests/unit_tests/scripts/inference/test_vlm_generation.py index 15023c5441..720b8148e4 100644 --- a/tests/unit_tests/scripts/inference/test_vlm_generation.py +++ b/tests/unit_tests/scripts/inference/test_vlm_generation.py @@ -27,6 +27,16 @@ 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", diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index 68166ed5f7..6261964bf3 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -267,7 +267,7 @@ def test_checkpoint_load_uses_builder_config_when_available(self): model_config = SimpleNamespace(transformer=transformer, finalize=MagicMock()) model = MagicMock() bridge = MagicMock() - bridge._model_bridge = SimpleNamespace(MODEL_CONFIG_CLASS=object) + 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( @@ -316,7 +316,7 @@ def test_hf_conversion_uses_nested_builder_transformer_config(self): model_config = SimpleNamespace(transformer=transformer) model = MagicMock() bridge = MagicMock() - bridge._model_bridge = SimpleNamespace(MODEL_CONFIG_CLASS=object) + 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( @@ -351,6 +351,44 @@ def test_hf_conversion_uses_nested_builder_transformer_config(self): 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 a29483b42f..6914d71cc9 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -843,6 +843,21 @@ def write(_rank, write_buckets, _results_queue): 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" From 49502549d7958d21fd7ad9eaa6219829272967ec Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 24 Aug 2026 22:13:10 -0700 Subject: [PATCH 71/72] fix(models): support Transformers 5.15 conversions Signed-off-by: yaoyu-33 --- .../bridge/training/utils/config_utils.py | 6 +++- .../deepseek/test_deepseek_conversion.py | 4 ++- .../test_nemotron_omni_conversion.py | 34 +++++++++++++------ .../training/utils/test_config_utils.py | 19 +++++++++++ 4 files changed, 50 insertions(+), 13 deletions(-) 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/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 954b01a2e8..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,15 +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"] -# Transformers 5.15 serializes the same hybrid blocks under their canonical names. -_SERIALIZED_LLM_LAYER_TYPES = [ - "linear_attention", - "moe", - "linear_attention", - "moe", - "linear_attention", - "full_attention", -] +_LLM_LAYER_TYPE_ALIASES = { + "linear_attention": "mamba", + "full_attention": "attention", +} _LLM_OVERRIDES = { "hybrid_override_pattern": "MEMEM*", "layer_types": _LLM_LAYER_TYPES, @@ -72,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(): @@ -223,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"] == _SERIALIZED_LLM_LAYER_TYPES - assert saved_config["llm_config"]["layers_block_type"] == _SERIALIZED_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/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.""" From 4c59dc75859d81f040e960019c445ad79ce546c5 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 25 Aug 2026 15:17:01 -0700 Subject: [PATCH 72/72] refactor(recipes): flatten Muse Glimmer configs Signed-off-by: yaoyu-33 --- .../recipes/muse_glimmer/h100/muse_glimmer.py | 258 ++++++++++++------ 1 file changed, 182 insertions(+), 76 deletions(-) diff --git a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py index 126216af5b..84d712e314 100644 --- a/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py +++ b/src/megatron/bridge/recipes/muse_glimmer/h100/muse_glimmer.py @@ -33,28 +33,19 @@ _CORD_V2_REVISION = "7f0115a4b758a71d6473b8d085751692da2fef98" # pragma: allowlist secret -def _apply_model_and_data( - cfg: ConfigContainer, - *, - seq_length: int, - tensor_parallel_size: int, - pipeline_parallel_size: int, - context_parallel_size: int, -) -> None: - """Apply the exact model, multimodal data, and common execution settings.""" +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 = seq_length - cfg.model.tensor_model_parallel_size = tensor_parallel_size - cfg.model.pipeline_model_parallel_size = pipeline_parallel_size - cfg.model.pipeline_dtype = torch.bfloat16 if pipeline_parallel_size > 1 else None + 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 - if pipeline_parallel_size == 2: - cfg.model.hybrid_layer_pattern = f"{'*' * 20}|{'*' * 32}" - elif pipeline_parallel_size == 4: - cfg.model.hybrid_layer_pattern = "|".join("*" * layers for layers in (9, 15, 15, 13)) - cfg.model.context_parallel_size = context_parallel_size - if context_parallel_size > 1: - cfg.model.cp_comm_type = "a2a" + 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 @@ -67,7 +58,7 @@ def _apply_model_and_data( cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["core_attn"] - cfg.dataset.seq_length = seq_length + cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = _MODEL_ID cfg.dataset.hf_processor_kwargs = { "revision": _MODEL_REVISION, @@ -77,7 +68,11 @@ def _apply_model_and_data( 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.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() @@ -89,43 +84,21 @@ def _apply_model_and_data( 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 - - -def _set_optimizer(cfg: ConfigContainer, *, max_lr: float, warmup_iters: int) -> None: - """Set the bounded cosine schedule shared by verification workloads.""" cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=warmup_iters, - max_lr=max_lr, + lr_warmup_iters=40, + max_lr=3e-4, min_lr=0.0, ) cfg.scheduler.lr_decay_iters = 100 - - -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() - _apply_model_and_data( - cfg, - seq_length=4096, - tensor_parallel_size=8, - pipeline_parallel_size=2, - context_parallel_size=1, - ) - 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.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 - _set_optimizer(cfg, max_lr=3e-4, warmup_iters=40) cfg.optimizer.use_precision_aware_optimizer = True cfg.checkpoint.save_interval = 50 cfg.checkpoint.load = None @@ -138,19 +111,64 @@ def muse_glimmer_30b_pretrain_32gpu_h100_bf16_multimodal_config() -> ConfigConta 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() - _apply_model_and_data( - cfg, - seq_length=4096, - tensor_parallel_size=8, - pipeline_parallel_size=2, - context_parallel_size=1, - ) + + 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 - _set_optimizer(cfg, max_lr=5e-6, warmup_iters=10) + 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 @@ -164,20 +182,64 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_config() -> ConfigContainer: 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() - _apply_model_and_data( - cfg, - seq_length=8192, - tensor_parallel_size=1, - pipeline_parallel_size=4, - context_parallel_size=2, - ) + + 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 - _set_optimizer(cfg, max_lr=5e-6, warmup_iters=10) + 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 @@ -191,13 +253,40 @@ def muse_glimmer_30b_sft_32gpu_h100_bf16_long_context_config() -> ConfigContaine 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() - _apply_model_and_data( - cfg, - seq_length=8192, - tensor_parallel_size=8, - pipeline_parallel_size=1, - context_parallel_size=1, - ) + + 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 @@ -209,10 +298,27 @@ def muse_glimmer_30b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora" 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 - _set_optimizer(cfg, max_lr=1e-4, warmup_iters=10) + 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 = {