`bridge.models.muse_glimmer.modeling_muse_glimmer`#

Megatron modeling components for Muse Glimmer.

Module Contents#

Classes#

`MuseGlimmerSelfAttention`	MCore self-attention with Muse’s residual-branch post norm.
`MuseGlimmerMLP`	MCore gated MLP with Muse’s residual-branch post norm.
`MuseGlimmerOutputLayerMixin`	Apply the Muse logit multiplier followed by tanh soft-capping.
`MuseGlimmerVisionPatchEmbedder`	Linear patch embedding plus exact half-pixel bilinear position lookup.
`MuseGlimmerVisionRotaryEmbedding`	Muse 2-D RoPE with [W, H, W, H] frequency interleaving.
`MuseGlimmerVisionAttention`	Bidirectional packed vision attention matching the HF parameter layout.
`MuseGlimmerVisionMLP`	Muse vision feed-forward network.
`MuseGlimmerVisionEncoderLayer`	Pre-normalized Muse vision transformer layer.
`MuseGlimmerVisionModel`	Muse vision tower with 3:1 window/full attention and 2x2 pixel shuffle.
`MuseGlimmerVisionAdapter`	Two-layer GELU adapter used before projection into the text width.
`MuseGlimmerModel`	Native MCore Hybrid model with the Muse vision modules attached.

Functions#

`get_muse_glimmer_hybrid_stack_spec`	Build the native Hybrid stack spec used by the Muse decoder.
`_vision_cu_seqlens`
`_vision_position_ids`
`_vision_window_index`
`_bilinear_indices_and_weights`
`_rotate_half`

Data#

__all__

API#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerSelfAttention(

config: megatron.core.transformer.transformer_config.TransformerConfig,

*args: Any,

**kwargs: Any,

)#

Bases: megatron.core.transformer.attention.SelfAttention

MCore self-attention with Muse’s residual-branch post norm.

Initialization

forward(

*args: Any,

**kwargs: Any,

) → tuple[torch.Tensor, torch.Tensor | None]#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerMLP(

config: megatron.core.transformer.transformer_config.TransformerConfig,

submodules: megatron.core.transformer.mlp.MLPSubmodules,

ffn_hidden_size: int | None = None,

**kwargs: Any,

)#

Bases: megatron.core.transformer.mlp.MLP

MCore gated MLP with Muse’s residual-branch post norm.

Initialization

forward(

*args: Any,

**kwargs: Any,

) → tuple[torch.Tensor, torch.Tensor | None]#

bridge.models.muse_glimmer.modeling_muse_glimmer.get_muse_glimmer_hybrid_stack_spec( config: megatron.core.transformer.transformer_config.TransformerConfig, ) → megatron.core.transformer.spec_utils.ModuleSpec#: Build the native Hybrid stack spec used by the Muse decoder.

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerOutputLayerMixin#

Bases: torch.nn.Module

Apply the Muse logit multiplier followed by tanh soft-capping.

forward(

*args: Any,

**kwargs: Any,

) → tuple[torch.Tensor, torch.Tensor | None]#

bridge.models.muse_glimmer.modeling_muse_glimmer._vision_cu_seqlens(grid_thw: torch.Tensor) → torch.Tensor#

bridge.models.muse_glimmer.modeling_muse_glimmer._vision_position_ids(grid_thw: torch.Tensor) → torch.Tensor#

bridge.models.muse_glimmer.modeling_muse_glimmer._vision_window_index( grid_thw: torch.Tensor, *, window_size: int, ) → tuple[torch.Tensor, torch.Tensor]#

bridge.models.muse_glimmer.modeling_muse_glimmer._bilinear_indices_and_weights( grid_thw: torch.Tensor, side: int, ) → tuple[torch.Tensor, torch.Tensor]#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionPatchEmbedder( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerVisionModelConfig, )#

Bases: torch.nn.Module

Linear patch embedding plus exact half-pixel bilinear position lookup.

Initialization

forward( pixel_values: torch.Tensor, grid_thw: torch.Tensor, ) → torch.Tensor#

bridge.models.muse_glimmer.modeling_muse_glimmer._rotate_half(hidden_states: torch.Tensor) → torch.Tensor#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionRotaryEmbedding( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerVisionModelConfig, )#

Bases: torch.nn.Module

Muse 2-D RoPE with [W, H, W, H] frequency interleaving.

Initialization

forward( hidden_states: torch.Tensor, position_ids: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionAttention( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerVisionModelConfig, )#

Bases: torch.nn.Module

Bidirectional packed vision attention matching the HF parameter layout.

Initialization

forward( hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], ) → torch.Tensor#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionMLP( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerVisionModelConfig, )#

Bases: torch.nn.Module

Muse vision feed-forward network.

Initialization

forward(hidden_states: torch.Tensor) → torch.Tensor#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionEncoderLayer( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerVisionModelConfig, transformer_config: megatron.core.transformer.transformer_config.TransformerConfig, )#

Bases: torch.nn.Module

Pre-normalized Muse vision transformer layer.

Initialization

forward( hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], ) → torch.Tensor#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionModel( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerVisionModelConfig, transformer_config: megatron.core.transformer.transformer_config.TransformerConfig, *, recompute_layers: bool = False, )#

Bases: torch.nn.Module

Muse vision tower with 3:1 window/full attention and 2x2 pixel shuffle.

Initialization

_pixel_shuffle( hidden_states: torch.Tensor, grid_thw: torch.Tensor, ) → torch.Tensor#

forward( pixel_values: torch.Tensor, grid_thw: torch.Tensor, ) → torch.Tensor#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionAdapter( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerModelConfig, )#

Bases: torch.nn.Module

Two-layer GELU adapter used before projection into the text width.

Initialization

forward(hidden_states: torch.Tensor) → torch.Tensor#

class bridge.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerModel( config: megatron.bridge.models.muse_glimmer.muse_glimmer_config.MuseGlimmerModelConfig, hybrid_stack_spec: megatron.core.transformer.spec_utils.ModuleSpec, vocab_size: int, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, *, pre_process: bool, post_process: bool, vp_stage: int | None, )#

Bases: megatron.core.models.hybrid.hybrid_model.HybridModel

Native MCore Hybrid model with the Muse vision modules attached.

Initialization

_EMPTY_NORM_EXTRA_STATE_SUFFIXES#: (‘input_layernorm._extra_state’, ‘self_attention.post_layernorm._extra_state’, ‘pre_mlp_layernorm._e…

_CENTERED_NORM_WEIGHT_SUFFIXES#: (‘input_layernorm.weight’, ‘self_attention.post_layernorm.weight’, ‘pre_mlp_layernorm.weight’, ‘mlp….

sharded_state_dict( prefix: str = '', sharded_offsets: tuple[tuple[int, int, int], ...] = (), metadata: dict[str, Any] | None = None, ) → dict[str, Any]#

Return a backend-stable Muse checkpoint schema.

Transformer Engine affine norms add empty _extra_state objects that the explicit local fallback does not expose. Omitting only those empty objects keeps checkpoints backend-interchangeable without relaxing strict loading for parameters or stateful Transformer Engine modules.

freeze( *, freeze_language_model: bool, freeze_vision_model: bool, freeze_vision_projection: bool, ) → None#

_vision_features( pixel_values: torch.Tensor, grid_thw: torch.Tensor, ) → torch.Tensor#

static _scatter_features( inputs_embeds: torch.Tensor, input_ids: torch.Tensor, features: torch.Tensor, token_id: int, ) → torch.Tensor#

forward( input_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.Tensor | None = None, pixel_values_videos: torch.Tensor | None = None, video_grid_thw: torch.Tensor | None = None, labels: torch.Tensor | None = None, inference_context: megatron.core.inference.contexts.BaseInferenceContext | None = None, runtime_gather_output: bool | None = None, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None = None, *, inference_params: megatron.core.inference.contexts.BaseInferenceContext | None = None, loss_mask: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None, ) → torch.Tensor | tuple[torch.Tensor, torch.Tensor]#

bridge.models.muse_glimmer.modeling_muse_glimmer.__all__#: [‘MuseGlimmerModel’, ‘MuseGlimmerVisionModel’, ‘get_muse_glimmer_hybrid_stack_spec’]

bridge.models.muse_glimmer.modeling_muse_glimmer#

Module Contents#

Classes#

Functions#

Data#

API#

`bridge.models.muse_glimmer.modeling_muse_glimmer`#