> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.muse_spark.model

Native AutoModel implementation of the Muse Spark Open Omni model.

## Module Contents

### Classes

| Name                                                                                                          | Description                                                         |
| ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`MuseSpark1xOmniForCausalLM`](#nemo_automodel-components-models-muse_spark-model-MuseSpark1xOmniForCausalLM) | Native Muse Spark Open text, image/video, and waveform-audio model. |
| [`MuseSparkTextModel`](#nemo_automodel-components-models-muse_spark-model-MuseSparkTextModel)                 | Muse Spark Open decoder backbone with an FP32 residual stream.      |

### Functions

| Name                                                                                              | Description                                                         |
| ------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`_causal_mask`](#nemo_automodel-components-models-muse_spark-model-_causal_mask)                 | Build an SDPA allow-mask for BSHD hidden states.                    |
| [`_residual_dtype`](#nemo_automodel-components-models-muse_spark-model-_residual_dtype)           | Resolve the checkpoint residual-stream dtype.                       |
| [`_select_cp_positions`](#nemo_automodel-components-models-muse_spark-model-_select_cp_positions) | Select the local round-robin positions from a BSHD position stream. |

### Data

[`ModelClass`](#nemo_automodel-components-models-muse_spark-model-ModelClass)

### API

```python
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](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-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`**

---

```python
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.

```python
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.

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

Return whether this module is a partial pipeline stage.

```python
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

```python
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.

```python
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.

```python
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.

```python
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.

```python
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 | None` — default: 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

```python
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 | None` — default: 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 | None` — default: None

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

---

**`position_ids`** `torch.Tensor | None` — default: None

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

---

**`attention_mask`** `torch.Tensor | None` — default: None

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

---

**`labels`** `torch.Tensor | None` — default: None

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

---

**`pixel_values`** `torch.Tensor | None` — default: None

Image patches \[patches, patch_features].

---

**`image_grid_thw`** `torch.Tensor | None` — default: None

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

---

**`pixel_values_videos`** `torch.Tensor | None` — default: None

Video patches \[patches, patch_features].

---

**`video_grid_thw`** `torch.Tensor | None` — default: None

Video patch grids \[videos, 3].

---

**`audio_values`** `list[torch.Tensor] | torch.Tensor | None` — default: None

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

---

**`logits_to_keep`** `int | torch.Tensor` — default: 0

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

---

**`kwargs`** `Any` — default: \{}

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

```python
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.

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

Return the text token embedding module.

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

Return the untied language-model head.

```python
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

```python
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.

```python
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`** `int` — default: 1

Number of independently scheduled PP microbatches.

---

**Returns:** `dict[str, Any]`

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

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

Replace the text token embedding module.

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

Replace the untied language-model head.

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

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

```python
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`**

---

```python
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.

```python
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 | None` — default: None

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

---

**`inputs_embeds`** `torch.Tensor | None` — default: None

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

---

**`position_ids`** `torch.Tensor | None` — default: None

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

---

**`attention_mask`** `torch.Tensor | None` — default: None

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

---

**`padding_mask`** `torch.Tensor | None` — default: None

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

---

**`output_hidden_states`** `bool` — default: False

Whether to return each layer's residual states.

---

**`**kwargs`** `Any` — default: \{}

Attention backend metadata, including packed sequence boundaries.

---

**Returns:** `BaseModelOutputWithPast`

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

```python
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

```python
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.

```python
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]`.

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