nemo_automodel.components.models.muse_glimmer.model

Native AutoModel implementation of the complete dense MuseGlimmer VLM.

The parameter hierarchy intentionally matches the checkpoint’s Hugging Face implementation. Text attention is backend-native: PyTorch SDPA is used for the ordinary backend and Transformer Engine’s DotProductAttention is constructed directly for TE BSHD/THD execution.

Module Contents

Classes

Name	Description
`MuseGlimmerAttention`	MuseGlimmer GQA with Q/K RMSNorm, optional RoPE, output gate, SDPA, and TE.
`MuseGlimmerDecoderLayer`	One MuseGlimmer language decoder layer.
`MuseGlimmerFinalRMSNorm`	RMSNorm whose checkpoint weight stores the actual output gain.
`MuseGlimmerForConditionalGeneration`	Native complete MuseGlimmer VLM with causal language-modeling head.
`MuseGlimmerMLP`	Bias-free SwiGLU language MLP.
`MuseGlimmerModel`	Complete MuseGlimmer vision-language backbone.
`MuseGlimmerPreTrainedModel`	Hugging Face-compatible base class for native MuseGlimmer.
`MuseGlimmerRMSNorm`	RMSNorm whose checkpoint weight stores an offset from one.
`MuseGlimmerRotaryEmbedding`	MuseGlimmer split-half rotary embedding matching the canonical HF implementation.
`MuseGlimmerScalelessRMSNorm`	Parameter-free RMSNorm used for embedding and Q/K normalization.

Functions

Name	Description
`_select_cp_positions`	Align full or already-sharded position IDs with model-owned CP embeddings.
`apply_rotary_emb`	Apply split-half RoPE to native BSHD or THD tensors.
`repeat_kv`	Repeat B,Hkv,S,D keys/values to query-head count.

Data

ModelClass

API

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerAttention(
    config: nemo_automodel.components.models.muse_glimmer.config.MuseGlimmerConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

MuseGlimmer GQA with Q/K RMSNorm, optional RoPE, output gate, SDPA, and TE.

attention_dropout

= config.attention_dropout

head_dim

= config.head_dim

k_proj

= nn.Linear(config.hidden_size, kv_dim, bias=False)

layer_type

= config.layer_types[layer_idx]

num_heads

= config.num_attention_heads

num_key_value_groups

= self.num_heads // self.num_key_value_heads

num_key_value_heads

= config.num_key_value_heads

o_proj

output_gate_proj

q_proj

qk_norm

scale_query_by

= config.scale_query_by

scaling

= self.head_dim ** -0.5

sliding_window

use_output_gate

= config.use_attn_output_gate

use_qk_norm

= config.use_qk_norm

use_rope

= config.no_rope_layers[layer_idx] == 1

v_proj

= nn.Linear(config.hidden_size, kv_dim, bias=False)

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerAttention._shape_qkv(
    hidden_states: torch.Tensor,
    is_thd: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerAttention._te_window_size(
    is_thd: bool,
    max_seqlen: int | None
) -> tuple[int, int]

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    cache_position: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerDecoderLayer(
    config: nemo_automodel.components.models.muse_glimmer.config.MuseGlimmerConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: GradientCheckpointingLayer

One MuseGlimmer language decoder layer.

input_layernorm

mlp

= MuseGlimmerMLP(config)

post_attention_layernorm

post_attn_norm

post_ffn_norm

self_attn

= MuseGlimmerAttention(config, layer_idx, backend)

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerDecoderLayer.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    cache_position: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerFinalRMSNorm(
    dim: int,
    eps: float = 1e-05
)

Bases: Module

RMSNorm whose checkpoint weight stores the actual output gain.

weight

= nn.Parameter(torch.zeros(dim))

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerFinalRMSNorm.forward(
    x: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration(
    config: nemo_automodel.components.models.muse_glimmer.config.MuseGlimmerConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)

Bases: HFCheckpointingMixin, MuseGlimmerPreTrainedModel, GenerationMixin

Native complete MuseGlimmer VLM with causal language-modeling head.

_keep_in_fp32_modules

= ['rotary_emb']

_tied_weights_keys

= []

_tp_plan

= {'lm_head': 'colwise_rep'}

backend

= backend or BackendConfig()

lm_head

model

= MuseGlimmerModel(config, self.backend)

state_dict_adapter

= MuseGlimmerStateDictAdapter(config)

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

vocab_size

= config.vocab_size

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration._set_te_cp_transport(
    comm_type: str
) -> None

Switch native TE modules between ordinary and packed CP transport.

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration._validate_thd_documents(
    batch: dict[str, typing.Any]
) -> int

Record real lengths and validate the TE p2p path when CP is active.

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.forward(
    input_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    inputs_embeds: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    cache_position: 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,
    vision_mask: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.from_config(
    config: nemo_automodel.components.models.muse_glimmer.config.MuseGlimmerConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
) -> 'MuseGlimmerForConditionalGeneration'

classmethod

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.get_decoder() -> torch.nn.Module

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.get_input_embeddings() -> torch.nn.Module

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.get_output_embeddings() -> torch.nn.Module

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.prepare_inputs_for_generation(
    input_ids,
    past_key_values = None,
    attention_mask = None,
    inputs_embeds = None,
    cache_position = None,
    pixel_values = None,
    image_grid_thw = None,
    pixel_values_videos = None,
    video_grid_thw = None,
    vision_mask = None,
    kwargs = {}
)

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]

Select native MuseGlimmer CP preparation for BSHD or packed TE THD.

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.set_decoder(
    decoder: torch.nn.Module
) -> None

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.set_input_embeddings(
    value: torch.nn.Module
) -> None

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerForConditionalGeneration.set_output_embeddings(
    new_embeddings: torch.nn.Module
) -> None

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerMLP(
    config: nemo_automodel.components.models.muse_glimmer.config.MuseGlimmerConfig
)

Bases: Module

Bias-free SwiGLU language MLP.

down_proj

gate_proj

up_proj

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerMLP.forward(
    x: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerModel(
    config: nemo_automodel.components.models.muse_glimmer.config.MuseGlimmerConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: MuseGlimmerPreTrainedModel

Complete MuseGlimmer vision-language backbone.

embed_norm

embed_tokens

has_vision

= config.has_vision

layers

norm

perception_emb_norm

rotary_emb

vision_adapter

= MuseGlimmerVisionAdapter(config)

vision_encoder

= MuseGlimmerVisionEncoder(config)

vision_projection

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerModel._embed_vision(
    input_ids: torch.Tensor,
    hidden_states: torch.Tensor,
    pixel_values: torch.Tensor | None,
    image_grid_thw: torch.Tensor | None,
    pixel_values_videos: torch.Tensor | None,
    video_grid_thw: torch.Tensor | None,
    vision_mask: torch.Tensor | None,
    global_vision_mask: torch.Tensor | None = None,
    thd_local_indices: torch.Tensor | None = None
) -> torch.Tensor

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerModel.forward(
    input_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    inputs_embeds: torch.Tensor | None = None,
    use_cache: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    cache_position: 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,
    vision_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerPreTrainedModel()

Bases: PreTrainedModel

Hugging Face-compatible base class for native MuseGlimmer.

_no_split_modules

= ['MuseGlimmerDecoderLayer']

_skip_keys_device_placement

= ['past_key_values']

base_model_prefix

= 'model'

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerRMSNorm(
    dim: int,
    eps: float = 1e-05
)

Bases: Module

RMSNorm whose checkpoint weight stores an offset from one.

weight

= nn.Parameter(torch.zeros(dim))

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerRMSNorm.forward(
    x: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerRotaryEmbedding(
    dim: int,
    max_position_embeddings: int,
    theta: float = 500000.0
)

Bases: Module

MuseGlimmer split-half rotary embedding matching the canonical HF implementation.

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerRotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

class nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerScalelessRMSNorm(
    dim: int,
    eps: float = 1e-05
)

Bases: Module

Parameter-free RMSNorm used for embedding and Q/K normalization.

nemo_automodel.components.models.muse_glimmer.model.MuseGlimmerScalelessRMSNorm.forward(
    x: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.muse_glimmer.model._select_cp_positions(
    position_ids: torch.Tensor | None,
    full_seq_len: int,
    local_indices: torch.Tensor,
    batch_size: int,
    device: torch.device
) -> torch.Tensor

Align full or already-sharded position IDs with model-owned CP embeddings.

nemo_automodel.components.models.muse_glimmer.model.apply_rotary_emb(
    x: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor

Apply split-half RoPE to native BSHD or THD tensors.

nemo_automodel.components.models.muse_glimmer.model.repeat_kv(
    x: torch.Tensor,
    n_rep: int
) -> torch.Tensor

Repeat B,Hkv,S,D keys/values to query-head count.

nemo_automodel.components.models.muse_glimmer.model.ModelClass = MuseGlimmerForConditionalGeneration