feat(model): add Muse Glimmer support

examples/conversion/compare_hf_and_megatron/compare.py

@@ -716,14 +716,26 @@ def _load_megatron_model(args):
             ),
             **_hf_revision_kwargs(args.hf_revision),
         )
-        model_provider = bridge.to_megatron_provider(load_weights=False)
-        model_provider.tensor_model_parallel_size = tp
-        model_provider.pipeline_model_parallel_size = pp
-        model_provider.expert_model_parallel_size = ep
-        model_provider.expert_tensor_parallel_size = etp
-        model_provider.pipeline_dtype = torch.bfloat16
-        model_provider.finalize()
-        model_provider.initialize_model_parallel(seed=0)
+        if getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False):
+            model_config = bridge.get_model_config()
+            transformer = model_config.transformer
+            transformer.tensor_model_parallel_size = tp
+            transformer.pipeline_model_parallel_size = pp
+            transformer.expert_model_parallel_size = ep
+            transformer.expert_tensor_parallel_size = etp
+            transformer.pipeline_dtype = torch.bfloat16
+            transformer.params_dtype = torch.bfloat16
+            model_config.finalize()
+            bridge._get_or_initialize_pg_collection(transformer)
+        else:
+            model_provider = bridge.to_megatron_provider(load_weights=False)
+            model_provider.tensor_model_parallel_size = tp
+            model_provider.pipeline_model_parallel_size = pp
+            model_provider.expert_model_parallel_size = ep
+            model_provider.expert_tensor_parallel_size = etp
+            model_provider.pipeline_dtype = torch.bfloat16
+            model_provider.finalize()
+            model_provider.initialize_model_parallel(seed=0)
         megatron_model = bridge.load_megatron_model(
             args.megatron_model_path,
             mp_overrides={
@@ -744,14 +756,29 @@ def _load_megatron_model(args):
             ),
             **_hf_revision_kwargs(args.hf_revision),
         )
-        model_provider = bridge.to_megatron_provider(load_weights=True)
-        model_provider.tensor_model_parallel_size = tp
-        model_provider.pipeline_model_parallel_size = pp
-        model_provider.expert_model_parallel_size = ep
-        model_provider.expert_tensor_parallel_size = etp
-        model_provider.pipeline_dtype = torch.bfloat16
-        model_provider.finalize()
-        megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False)
+        if getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False):
+            model_config = bridge.get_model_config()
+            transformer = model_config.transformer
+            transformer.tensor_model_parallel_size = tp
+            transformer.pipeline_model_parallel_size = pp
+            transformer.expert_model_parallel_size = ep
+            transformer.expert_tensor_parallel_size = etp
+            transformer.pipeline_dtype = torch.bfloat16
+            transformer.params_dtype = torch.bfloat16
+            megatron_model = bridge.get_model(
+                model_config,
+                wrap_with_ddp=False,
+                mixed_precision_wrapper=None,
+            )
+        else:
+            model_provider = bridge.to_megatron_provider(load_weights=True)
+            model_provider.tensor_model_parallel_size = tp
+            model_provider.pipeline_model_parallel_size = pp
+            model_provider.expert_model_parallel_size = ep
+            model_provider.expert_tensor_parallel_size = etp
+            model_provider.pipeline_dtype = torch.bfloat16
+            model_provider.finalize()
+            megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False)
 
     # Workaround: disable MTP for inference (causes hangs on NCCL collectives)
     for m in megatron_model:

examples/models/muse_glimmer/README.md

@@ -0,0 +1,55 @@
+# Muse Glimmer 30B
+
+[Muse Glimmer 30B](https://huggingface.co/meta-models/Muse-Glimmer-30B) is a dense
+vision-language model with a roughly 2B-parameter vision encoder and a 28B-parameter
+decoder. Megatron Bridge models the complete checkpoint: the vision tower, adapter,
+projection, and text decoder all participate in checkpoint conversion.
+
+## Modeling
+
+Muse Glimmer uses the builder-backed API. `MuseGlimmerBridge` translates the Hugging
+Face configuration directly into `MuseGlimmerModelConfig`, and
+`MuseGlimmerModelBuilder` constructs `MuseGlimmerModel` as a native Megatron-Core
+`HybridModel`; no model provider or `GPTModel` wrapper is involved.
+
+The text decoder uses a native Megatron-Core `HybridStack` whose layer pattern selects
+the Muse transformer-layer spec at every depth. It preserves GQA,
+tensor/pipeline/context parallelism, 3:1 sliding/full attention, per-layer NoPE,
+weightless Q/K RMSNorm, and the full-head sigmoid attention output gate. Muse-specific
+residual-branch RMSNorms wrap the attention and MLP outputs without replacing
+Megatron-Core's `TransformerLayer` execution path. The embedding norm, final scaled
+RMSNorm, output multiplier, and tanh logit soft-cap reproduce the published decoder
+math.
+
+The replicated vision path implements the published 50-layer encoder, half-pixel
+bilinear learned-position lookup with zero padding, 2-D RoPE, 3:1 window/full attention,
+2x2 pixel shuffle, two-layer GELU adapter, text-width projection, and weightless
+perception RMSNorm. Its parameters are marked for tensor-parallel gradient averaging.
+
+## Builder-backed configuration
+
+```python
+from megatron.bridge import AutoBridge
+
+bridge = AutoBridge.from_hf_pretrained(
+    "meta-models/Muse-Glimmer-30B",
+    revision="f84ecc3a0ea984a4c04542a84269e3d065350a6e",  # pragma: allowlist secret
+)
+model_config = bridge.get_model_config()
+
+# Apply deployment-specific parallelism before construction.
+model_config.tensor_model_parallel_size = 8
+models = bridge.get_model(model_config)
+```
+
+The published configuration needs Transformers 5.15 for native Hugging Face model
+execution. Megatron Bridge includes only a configuration compatibility class so that
+the lazy safetensors import and builder path continue to work with the repository's
+currently pinned Transformers release; it does not vendor or execute the Hugging Face
+model implementation.
+
+## Verification status
+
+See the [model verification card](../../model_verification_cards/muse-glimmer-30b/card.yaml)
+for the pinned source revision and current evidence. The card keeps import, export,
+forward correlation, inference, and training claims separate.

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",

scripts/conversion/gpu_backend.py

@@ -183,6 +183,33 @@ def _configure_model_provider(
         model_provider.use_cpu_initialization = True
 
 
+def _uses_model_builder(bridge: AutoBridge) -> bool:
+    """Return whether the selected bridge supports native builder construction."""
+    return getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False)
+
+
+def _configure_model_config(
+    model_config,
+    *,
+    tp: int,
+    pp: int,
+    ep: int,
+    etp: int,
+    dtype: torch.dtype,
+    use_cpu: bool = False,
+) -> None:
+    """Apply distributed parallelism and dtype settings to a builder config."""
+    transformer = model_config.transformer
+    transformer.tensor_model_parallel_size = tp
+    transformer.pipeline_model_parallel_size = pp
+    transformer.expert_model_parallel_size = ep
+    transformer.expert_tensor_parallel_size = etp
+    transformer.pipeline_dtype = dtype
+    transformer.params_dtype = dtype
+    if use_cpu:
+        transformer.use_cpu_initialization = True
+
+
 def _hf_tokenizer_kwargs(bridge: AutoBridge, *, trust_remote_code: bool) -> dict[str, object]:
     """Build tokenizer metadata for a saved Megatron checkpoint."""
     tokenizer_kwargs: dict[str, object] = {}
@@ -320,12 +347,22 @@ def import_checkpoint(
         torch_dtype=dtype,
         **revision_kwargs,
     )
-    model_provider = bridge.to_megatron_provider(load_weights=True)
-    _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype)
-    _maybe_generate_pipeline_layout(bridge, model_provider, pp)
-    model_provider.finalize()
-    model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False)
-    megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False)
+    if _uses_model_builder(bridge):
+        model_config = bridge.get_model_config()
+        _configure_model_config(model_config, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype)
+        _maybe_generate_pipeline_layout(bridge, model_config, pp)
+        megatron_model = bridge.get_model(
+            model_config,
+            wrap_with_ddp=False,
+            mixed_precision_wrapper=None,
+        )
+    else:
+        model_provider = bridge.to_megatron_provider(load_weights=True)
+        _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype)
+        _maybe_generate_pipeline_layout(bridge, model_provider, pp)
+        model_provider.finalize()
+        model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False)
+        megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False)
 
     bridge.save_megatron_model(
         megatron_model,
@@ -410,12 +447,20 @@ def export_checkpoint(
     # Preserve the reference wrapper's streaming state source and shard map while
     # exporting the checkpoint-derived architecture and vocabulary configuration.
     bridge.hf_pretrained.config = checkpoint_config_bridge.hf_pretrained
-    model_provider = bridge.to_megatron_provider(load_weights=False)
-    _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype, use_cpu=use_cpu)
-    _maybe_restore_pipeline_layout(bridge, model_provider, megatron_path, pp)
-    resolved_pipeline_layout = model_provider.pipeline_model_parallel_layout
-    model_provider.finalize()
-    model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False)
+    if _uses_model_builder(bridge):
+        model_config = bridge.get_model_config()
+        _configure_model_config(model_config, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype, use_cpu=use_cpu)
+        _maybe_restore_pipeline_layout(bridge, model_config, megatron_path, pp)
+        resolved_pipeline_layout = model_config.pipeline_model_parallel_layout
+        model_config.finalize()
+        bridge._get_or_initialize_pg_collection(model_config.transformer)
+    else:
+        model_provider = bridge.to_megatron_provider(load_weights=False)
+        _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype, use_cpu=use_cpu)
+        _maybe_restore_pipeline_layout(bridge, model_provider, megatron_path, pp)
+        resolved_pipeline_layout = model_provider.pipeline_model_parallel_layout
+        model_provider.finalize()
+        model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False)
 
     model_parallel_overrides: dict[str, object] = {
         "tensor_model_parallel_size": tp,
@@ -484,11 +529,21 @@ def roundtrip_checkpoint(
         torch_dtype=dtype,
     )
 
-    model_provider = bridge.to_megatron_provider(load_weights=True)
-    _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype)
-    _maybe_generate_pipeline_layout(bridge, model_provider, pp)
-    model_provider.finalize()
-    model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False)
-    megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False)
+    if _uses_model_builder(bridge):
+        model_config = bridge.get_model_config()
+        _configure_model_config(model_config, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype)
+        _maybe_generate_pipeline_layout(bridge, model_config, pp)
+        megatron_model = bridge.get_model(
+            model_config,
+            wrap_with_ddp=False,
+            mixed_precision_wrapper=None,
+        )
+    else:
+        model_provider = bridge.to_megatron_provider(load_weights=True)
+        _configure_model_provider(model_provider, tp=tp, pp=pp, ep=ep, etp=etp, dtype=dtype)
+        _maybe_generate_pipeline_layout(bridge, model_provider, pp)
+        model_provider.finalize()
+        model_provider.initialize_model_parallel(seed=0, create_gloo_process_groups=False)
+        megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False)
     _verify_roundtrip_weights(bridge, megatron_model)
     print_rank_0("GPU round-trip validation complete")