nemo_automodel.components.models.muse_spark.model

Native AutoModel implementation of the Muse Spark Open Omni model.

Module Contents

Classes

Name	Description
`MuseSpark1xOmniForCausalLM`	Native Muse Spark Open text, image/video, and waveform-audio model.
`MuseSparkTextModel`	Muse Spark Open decoder backbone with an FP32 residual stream.

Functions

Name	Description
`_causal_mask`	Build an SDPA allow-mask for BSHD hidden states.
`_residual_dtype`	Resolve the checkpoint residual-stream dtype.
`_select_cp_positions`	Select the local round-robin positions from a BSHD position stream.

Data

ModelClass

API

class nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Native Muse Spark Open text, image/video, and waveform-audio model.

_keep_in_fp32_modules

= ['rotary_emb']

_keep_in_fp32_modules_strict

_packed_cp_attn_backends

= ('te', 'magi')

_pp_keep_self_forward

bool = True

_pp_return_hidden_states_supported

bool = True

audio_encoder

backend

= backend or BackendConfig()

lm_head

model

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

vision_model

= MuseSparkVisionModel(config)

vision_projection

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM._audio_features(
    audio_values: list[torch.Tensor] | torch.Tensor
) -> torch.Tensor

Encode waveform values into normalized [tokens,H] features.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM._embed_and_splice(
    input_ids: 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,
    audio_values: list[torch.Tensor] | torch.Tensor | None,
    local_indices: torch.Tensor | None,
    global_image_mask: torch.Tensor | None,
    global_video_mask: torch.Tensor | None,
    global_audio_mask: torch.Tensor | None
) -> torch.Tensor

Build normalized text embeddings and splice all supplied modalities.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM._is_pipeline_parallel_stage() -> bool

Return whether this module is a partial pipeline stage.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM._local_media_features(
    features: torch.Tensor,
    local_mask: torch.Tensor,
    global_mask: torch.Tensor | None,
    local_indices: torch.Tensor | None
) -> torch.Tensor

staticmethod

Select global media features for one packed CP rank.

Parameters:

features

torch.Tensor

Global feature stream shaped [global_media_tokens,H].

local_mask

torch.Tensor

Boolean mask over local packed tokens, shape [T_local].

global_mask

torch.Tensor | None

Boolean mask over the unsharded packed stream, shape [T_global]; required together with local_indices.

local_indices

torch.Tensor | None

Global indices for every local token, shape [T_local].

Returns: torch.Tensor

Feature rows corresponding to true entries in local_mask, in

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM._normalized_media(
    features: torch.Tensor
) -> torch.Tensor

Apply the checkpoint’s parameterless embedding norm to [tokens,H] media.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM._splice_features(
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    token_id: int,
    features: torch.Tensor,
    global_mask: torch.Tensor | None,
    local_indices: torch.Tensor | None
) -> torch.Tensor

Splice [media_tokens,H] features into BSHD or THD token embeddings.

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

Validate packed document lengths and return the real maximum length.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM._vision_features(
    pixel_values: torch.Tensor,
    grid_thw: torch.Tensor
) -> torch.Tensor

Encode flattened image/video patches into normalized [tokens,H] features.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]

Map generated PP FQNs to Muse Spark Open’s real multimodal layout.

Parameters:

module_names_per_stage

list[list[str]]

Auto-generated module FQNs for each stage.

layers_prefix

str

Prefix used for the text decoder, normally model..

text_model

nn.Module | NoneDefaults to None

Resolved text module supplied by the PP splitter; unused.

Returns: list[list[str]]

Stage module FQNs with the top-level vision/audio towers and

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    output_hidden_states: bool = False,
    return_dict: bool = True,
    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,
    audio_values: list[torch.Tensor] | torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | tuple[torch.Tensor, ...] | torch.Tensor

Run causal language modeling for BSHD or packed THD multimodal inputs.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Token ids [batch, sequence] or packed [tokens]. Under PP, stage zero receives global ids and later stages receive CP-local floating-point activations [batch, local_sequence, hidden] here.

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional embeddings [batch, sequence, hidden] or packed [tokens, hidden]; already local when received from a PP stage.

position_ids

torch.Tensor | NoneDefaults to None

Local positions [batch, local_sequence] or [local_tokens].

attention_mask

torch.Tensor | NoneDefaults to None

Optional attention mask broadcastable to [batch, heads, sequence, sequence].

labels

torch.Tensor | NoneDefaults to None

Optional local targets [batch, local_sequence] or [local_tokens].

pixel_values

torch.Tensor | NoneDefaults to None

Image patches [patches, patch_features].

image_grid_thw

torch.Tensor | NoneDefaults to None

Image patch grids [images, 3] in temporal/height/width order.

pixel_values_videos

torch.Tensor | NoneDefaults to None

Video patches [patches, patch_features].

video_grid_thw

torch.Tensor | NoneDefaults to None

Video patch grids [videos, 3].

audio_values

list[torch.Tensor] | torch.Tensor | NoneDefaults to None

Mono waveforms [samples] or [samples, 1], supplied as one tensor or a list of variable-length tensors.

logits_to_keep

int | torch.TensorDefaults to 0

Final token count, or indices [kept_tokens]; zero selects every position.

kwargs

AnyDefaults to {}

Attention/packing metadata and model-owned media coordinates. _thd_local_indices [local_tokens] selects from the current microbatch’s global padded stream; global media masks cover that full stream. PP may add a leading singleton microbatch dimension.

Returns: CausalLMOutputWithPast | tuple[torch.Tensor, ...] | torch.Tensor

Causal LM output with logits [batch, local_sequence, vocab] and hidden

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.from_config(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
) -> 'MuseSpark1xOmniForCausalLM'

classmethod

Construct the native model from a resolved checkpoint config.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.get_input_embeddings() -> torch.nn.Module

Return the text token embedding module.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.get_output_embeddings() -> torch.nn.Module

Return the untied language-model head.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]

Build exact PP input/output metadata without a probe forward.

Parameters:

is_first

bool

Whether this is the first physical pipeline stage.

microbatch_size

int

Number of samples in one pipeline microbatch.

seq_len

int

Global token length presented to stage zero, for both ordinary BSHD and packed THD batches.

dtype

torch.dtype

Framework-selected compute dtype for final logits or fused loss inputs. Intermediate PP activations follow the model’s residual dtype, independently of the parameter/compute dtype. FSDP must preserve these natural model output dtypes.

Returns: tuple[torch.Tensor, ...]

A pair (inputs_meta, outputs_meta). Token ids have layout

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

Initialize every parameter for checkpoint-free construction.

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

Select BSHD round-robin or packed TE THD context-parallel preparation.

Parameters:

batch

dict[str, Any]

Token tensors of shape [batch, sequence], packed real/padded lengths of shape [batch, documents], and optional media payloads.

num_chunks

intDefaults to 1

Number of independently scheduled PP microbatches.

Returns: dict[str, Any]

Model-owned sharder and metadata. Global media masks have shape

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

Replace the text token embedding module.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.set_output_embeddings(
    value: torch.nn.Module
) -> None

Replace the untied language-model head.

nemo_automodel.components.models.muse_spark.model.MuseSpark1xOmniForCausalLM.update_moe_gate_bias() -> None

Keep the checkpoint-trained, non-updating router correction bias fixed.

class nemo_automodel.components.models.muse_spark.model.MuseSparkTextModel(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None
)

Bases: Module

Muse Spark Open decoder backbone with an FP32 residual stream.

embed_norm

embed_tokens

layers

m_emb

moe_config

norm

rotary_emb

nemo_automodel.components.models.muse_spark.model.MuseSparkTextModel.embed(
    input_ids: torch.Tensor
) -> torch.Tensor

Embed [B,S] or packed [T] ids and apply the parameterless norm.

nemo_automodel.components.models.muse_spark.model.MuseSparkTextModel.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    output_hidden_states: bool = False,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Run BSHD or packed THD decoder activations.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Local token ids [batch, sequence] or packed [tokens].

inputs_embeds

torch.Tensor | NoneDefaults to None

Local embeddings [batch, sequence, hidden] or [tokens, hidden].

position_ids

torch.Tensor | NoneDefaults to None

Local positions [batch, sequence] or packed [tokens].

attention_mask

torch.Tensor | NoneDefaults to None

Token validity [batch, sequence] (True/1 means valid), or a prepared attention mask broadcastable to [batch, heads, sequence, sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Local token padding [batch, sequence] or [tokens], with True meaning padding. Takes precedence over attention_mask.

output_hidden_states

boolDefaults to False

Whether to return each layer’s residual states.

**kwargs

AnyDefaults to {}

Attention backend metadata, including packed sequence boundaries.

Returns: BaseModelOutputWithPast

Decoder output with last_hidden_state [batch, sequence, hidden] or

nemo_automodel.components.models.muse_spark.model._causal_mask(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None,
    window: int | None
) -> torch.Tensor

Build an SDPA allow-mask for BSHD hidden states.

Parameters:

hidden_states

torch.Tensor

Input activations shaped [B,S,H].

attention_mask

torch.Tensor | None

Optional padding mask [B,S] or ready SDPA mask broadcastable to [B,heads,S,S].

window

int | None

Number of causal key positions visible to each query, or None for full causal attention.

Returns: torch.Tensor

Boolean SDPA mask shaped [B,1,S,S] (or the caller-provided

nemo_automodel.components.models.muse_spark.model._residual_dtype(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig
) -> torch.dtype

Resolve the checkpoint residual-stream dtype.

nemo_automodel.components.models.muse_spark.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

Select the local round-robin positions from a BSHD position stream.

Parameters:

position_ids

torch.Tensor | None

Full [B,S] or already-local [B,S/cp] positions.

full_seq_len

int

Unpadded global sequence length.

local_indices

torch.Tensor

Global positions retained by this CP rank, shape [padded_sequence/cp].

batch_size

int

BSHD batch size.

device

torch.device

Device on which default positions are constructed.

Returns: torch.Tensor

Local position ids with shape [B,padded_sequence/cp].

nemo_automodel.components.models.muse_spark.model.ModelClass = MuseSpark1xOmniForCausalLM