Skip to navigation
DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsMuse Sparknemo_automodel.components.models.muse_spark.layers

nemo_automodel.components.models.muse_spark.layers

View as Markdown

Text layers for the native Muse Spark Open implementation.

Module Contents

Classes

NameDescription
MuseSparkAttentionMuse GQA with Q/K and output RMSNorm, output gate, RoPE, SDPA, and TE.
MuseSparkDecoderLayerOne FP32 residual-stream Muse Spark Open decoder layer.
MuseSparkGateFP32 MetaP-scaled sigmoid router with bias used for selection only.
MuseSparkInputScaleRMSNormApply the learned gain before parameter-free RMS normalization.
MuseSparkMoECompressed Muse experts with per-expert norms and grouped-GEMM dispatch.
MuseSparkRMSNormFP32-accumulating RMSNorm with optional zero-centered checkpoint gain.
MuseSparkResidualLearned norm-preserving residual connection from the Muse checkpoint.
MuseSparkRotaryEmbeddingFP32 adjacent-pair rotary frequencies shared by all text layers.

Functions

NameDescription
_checkpoint_recomputed_mathKeep pure compiled math and autotuner scratch tensors outside SAC caching.
_compiled_norm_coreLazily share each of the five pure compiled cores within this process.
_input_scale_norm_coreApply the effective gain before FP32 RMS normalization.
_norm_residual_coreFuse a branch norm and residual blend without changing the branch cast.
_output_norm_gate_coreFuse parameter-free output RMSNorm and the checkpoint’s sigmoid gate.
_qk_norm_rope_coreNormalize and rotate Q/K while preserving the checkpoint’s dtype boundaries.
_repeat_kvRepeat BSHD key/value heads from [B,Nkv,S,D] to query-head count.
_residual_coreBlend the residual and branch using the checkpoint’s learned gate.
_rms_norm_coreNormalize in FP32, then apply the effective gain before the output cast.
apply_rotary_embeddingApply adjacent-pair RoPE with the checkpoint’s arithmetic precision.
build_muse_spark_moe_configTranslate the released Muse router and compressed-expert contract.
calculate_gain_centerReturn the checkpoint-defined zero-centered RMSNorm gain offset.
next_dtypeResolve the config’s model dtype without importing recipe utilities.

Data

_NormArgs

_NormOutput

API

class nemo_automodel.components.models.muse_spark.layers.MuseSparkAttention(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

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

_fused_output_norm_gate
_fused_qk_norm_rope
head_dim
= config.head_dim
k_proj
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_gate_proj
o_norm
o_proj
q_proj
qk_norm
rope_in_fp32
= config.model_type != 'muse_spark'
scaling
= self.head_dim ** -0.5
sliding_window
= config.layer_sliding_windows[layer_idx] or None
use_rope
= config.layer_uses_rope(layer_idx)
v_proj
nemo_automodel.components.models.muse_spark.layers.MuseSparkAttention._shape_qkv(
    hidden_states: torch.Tensor,
    is_thd: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Project [B,S,H] or [T,H] states into matching BSHD/THD QKV.

nemo_automodel.components.models.muse_spark.layers.MuseSparkAttention._te_window_size(
    is_thd: bool,
    max_real_seqlen: int | None
) -> tuple[int, int]

Return TE’s left/right window, treating short packed documents as full causal.

nemo_automodel.components.models.muse_spark.layers.MuseSparkAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Attend over ordinary or packed states using this layer’s CP transport.

Parameters:

hidden_states
torch.Tensor

CP-local states [batch, local_sequence, hidden] or packed [local_tokens, hidden].

position_embeddings
tuple[torch.Tensor, torch.Tensor]

Rotary cosine/sine tensors [batch, local_sequence, head_dim] or packed [local_tokens, head_dim].

attention_mask
torch.Tensor | NoneDefaults to None

Optional mask broadcastable to [batch, heads, sequence, sequence]; unused for packed THD.

kwargs
AnyDefaults to {}

Attention metadata, including qkv_format and global packed document boundaries cu_seqlens [documents + 1].

Returns: torch.Tensor

Attention output with the same local layout as hidden_states.

class nemo_automodel.components.models.muse_spark.layers.MuseSparkDecoderLayer(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: GradientCheckpointingLayer

One FP32 residual-stream Muse Spark Open decoder layer.

_fused_norm_residual
add_attn
add_ffn
feed_forward
layer_type
= config.layer_types[layer_idx]
mlp
MuseSparkMoE

Expose the checkpoint-named feed-forward through the generic MoE protocol.

post_attn_norm
post_ffn_norm
pre_ffn_norm
self_attn
= MuseSparkAttention(config, layer_idx, backend)
nemo_automodel.components.models.muse_spark.layers.MuseSparkDecoderLayer.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Transform batched or packed residual states through attention and MoE.

Parameters:

hidden_states
torch.Tensor

Residual activations shaped [B, S, H] or [T, H].

position_embeddings
tuple[torch.Tensor, torch.Tensor]

Cosine and sine tensors shaped [B, S, head_dim / 2].

attention_mask
torch.Tensor | NoneDefaults to None

Optional mask in the selected attention backend’s layout.

padding_mask
torch.Tensor | NoneDefaults to None

Optional token padding mask shaped [B, S] or [T].

**kwargs
AnyDefaults to {}

Additional arguments for the selected attention backend.

Returns: torch.Tensor

Independent residual states with the input shape, dtype and device.

nemo_automodel.components.models.muse_spark.layers.MuseSparkDecoderLayer.nemo_checkpoint_context_fn(
    context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
) -> collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None

Preserve norm compilation limits on the checkpoint recomputation thread.

Parameters:

context_fn
Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None

Existing forward/recompute context factory, if present.

Returns: Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None

A factory composing the existing contexts with the forward thread’s

class nemo_automodel.components.models.muse_spark.layers.MuseSparkGate(
    config: nemo_automodel.components.moe.config.MoEConfig,
    router_output_multiplier: float,
    division_eps: float
)

Bases: Gate

FP32 MetaP-scaled sigmoid router with bias used for selection only.

nemo_automodel.components.models.muse_spark.layers.MuseSparkGate.forward(
    hidden_states: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: typing.Any = None
) -> tuple[torch.Tensor, torch.Tensor, None]

Route local [tokens, hidden] states and return FP32 top-k weights.

class nemo_automodel.components.models.muse_spark.layers.MuseSparkInputScaleRMSNorm()

Bases: MuseSparkRMSNorm

Apply the learned gain before parameter-free RMS normalization.

_norm_core
= staticmethod(_input_scale_norm_core)
class nemo_automodel.components.models.muse_spark.layers.MuseSparkMoE(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: MoE

Compressed Muse experts with per-expert norms and grouped-GEMM dispatch.

gate
post_expert_norm
pre_expert_norm
nemo_automodel.components.models.muse_spark.layers.MuseSparkMoE.forward(
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
    cp_mesh: typing.Any = None
) -> torch.Tensor

Route [..., hidden] states through compressed, normalized experts.

class nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm(
    dim: int,
    eps: float,
    dtype: torch.dtype,
    has_weight: bool = True,
    zero_centered_gamma: bool = False,
    gain_center: float = 1.0,
    compile_norm: bool = False
)

Bases: Module

FP32-accumulating RMSNorm with optional zero-centered checkpoint gain.

_forward_core
_norm_core
= staticmethod(_rms_norm_core)
weight
nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm.effective_weight() -> torch.Tensor | None

Return the actual multiplicative gain with shape [hidden].

nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Normalize activations with the selected eager or compiled core.

Parameters:

hidden_states
torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

Returns: torch.Tensor

Independent normalized tensor with the input shape, dtype and device.

nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm.reset_parameters() -> None

Initialize the serialized gain without changing its effective value.

class nemo_automodel.components.models.muse_spark.layers.MuseSparkResidual(
    dim: int,
    dtype: torch.dtype,
    temperature: float = 0.3,
    eps: float = 0.001
)

Bases: Module

Learned norm-preserving residual connection from the Muse checkpoint.

gate
= nn.Parameter(torch.empty(dim, dtype=dtype))
nemo_automodel.components.models.muse_spark.layers.MuseSparkResidual.forward(
    stream: torch.Tensor,
    branch: torch.Tensor
) -> torch.Tensor

Blend activations while preserving the residual stream dtype.

Parameters:

stream
torch.Tensor

Residual tensor of shape […, hidden], ordinarily FP32.

branch
torch.Tensor

Tensor with the same shape and device as stream.

Returns: torch.Tensor

Independent tensor with the stream shape, dtype and device.

class nemo_automodel.components.models.muse_spark.layers.MuseSparkRotaryEmbedding(
    dim: int,
    theta: float
)

Bases: Module

FP32 adjacent-pair rotary frequencies shared by all text layers.

nemo_automodel.components.models.muse_spark.layers.MuseSparkRotaryEmbedding.forward(
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Return cosine/sine tensors shaped [batch, sequence, head_dim / 2].

nemo_automodel.components.models.muse_spark.layers.MuseSparkRotaryEmbedding.reset_parameters() -> None

Recreate FP32 frequencies after meta-device materialization.

Keep pure compiled math and autotuner scratch tensors outside SAC caching.

Muse’s selective checkpoint policy saves routing results, not norm or activation math. Autotuning a new shape during recomputation can execute scratch operations absent from the original forward. Exclude these pure islands from dispatch-mode accounting in both passes; autograd and the enclosing activation checkpoint still run.

Lazily share each of the five pure compiled cores within this process.

The cache retains callables, never model parameters or activations. Eager instances keep their original callables; enabling one model does not change another model’s dispatch. Compilation starts on the first real tensor call.

nemo_automodel.components.models.muse_spark.layers._input_scale_norm_core(
    hidden_states: torch.Tensor,
    weight: torch.Tensor | None,
    eps: float
) -> torch.Tensor

Apply the effective gain before FP32 RMS normalization.

Parameters:

hidden_states
torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

weight
torch.Tensor | None

Optional effective FP32 gain of shape [hidden], on the input device.

eps
float

Stabilizer added to the mean squared scaled activation.

Returns: torch.Tensor

Independent tensor of shape […, hidden], in the input dtype and device.

nemo_automodel.components.models.muse_spark.layers._norm_residual_core(
    stream: torch.Tensor,
    branch: torch.Tensor,
    weight: torch.Tensor | None,
    gate: torch.Tensor,
    input_scale: bool,
    norm_eps: float,
    temperature: float,
    residual_eps: float
) -> torch.Tensor

Fuse a branch norm and residual blend without changing the branch cast.

Parameters:

stream
torch.Tensor

Residual tensor of shape […, hidden], ordinarily FP32.

branch
torch.Tensor

Tensor of shape […, hidden], matching the stream shape and device. Normalization casts back to this dtype before residual blending.

weight
torch.Tensor | None

Optional effective FP32 norm gain of shape [hidden].

gate
torch.Tensor

Trainable residual gate of shape [hidden], on the stream device.

input_scale
bool

Whether the gain precedes rather than follows normalization.

norm_eps
float

RMS normalization stabilizer.

temperature
float

Residual gate temperature.

residual_eps
float

Lower bound on the squared residual coefficient.

Returns: torch.Tensor

Independent tensor of shape […, hidden], in the stream dtype and device.

nemo_automodel.components.models.muse_spark.layers._output_norm_gate_core(
    output: torch.Tensor,
    gate: torch.Tensor,
    eps: float
) -> torch.Tensor

Fuse parameter-free output RMSNorm and the checkpoint’s sigmoid gate.

Parameters:

output
torch.Tensor

Attention output [B,S,Hq,D] or packed [T,Hq,D].

gate
torch.Tensor

Gate projection with the same shape, dtype and device as output.

eps
float

FP32 RMS normalization stabilizer.

Returns: torch.Tensor

Gated output in the original dtype. Both normalization and sigmoid

nemo_automodel.components.models.muse_spark.layers._qk_norm_rope_core(
    query: torch.Tensor,
    key: torch.Tensor,
    cosine: torch.Tensor,
    sine: torch.Tensor,
    eps: float,
    query_scale: float,
    use_rope: bool,
    upcast_rope: bool
) -> tuple[torch.Tensor, torch.Tensor]

Normalize and rotate Q/K while preserving the checkpoint’s dtype boundaries.

Parameters:

query
torch.Tensor

Query tensor [B,S,Hq,D] or packed [T,Hq,D].

key
torch.Tensor

Key tensor with matching leading dimensions and Hkv heads.

cosine
torch.Tensor

Rotary cosines [B,S,D/2], or [1,T,D/2] for packed inputs.

sine
torch.Tensor

Rotary sines with the same shape and dtype as cosine.

eps
float

FP32 RMS normalization stabilizer.

query_scale
float

Multiplier applied after the normalized query’s dtype cast.

use_rope
bool

Whether this layer uses adjacent-pair rotary embeddings.

upcast_rope
bool

Whether rotation uses FP32 (legacy) or the Q/K dtype (v2).

Returns: torch.Tensor

Independent normalized query and key tensors in their original layouts

nemo_automodel.components.models.muse_spark.layers._repeat_kv(
    hidden_states: torch.Tensor,
    repeats: int
) -> torch.Tensor

Repeat BSHD key/value heads from [B,Nkv,S,D] to query-head count.

nemo_automodel.components.models.muse_spark.layers._residual_core(
    stream: torch.Tensor,
    branch: torch.Tensor,
    gate: torch.Tensor,
    temperature: float,
    eps: float
) -> torch.Tensor

Blend the residual and branch using the checkpoint’s learned gate.

Parameters:

stream
torch.Tensor

Residual tensor of shape […, hidden], ordinarily FP32.

branch
torch.Tensor

Tensor of shape […, hidden], matching the stream shape and device.

gate
torch.Tensor

Trainable gate of shape [hidden], on the stream device.

temperature
float

Gate temperature.

eps
float

Lower bound on the squared residual coefficient.

Returns: torch.Tensor

Independent tensor of shape […, hidden], in the stream dtype and device.

nemo_automodel.components.models.muse_spark.layers._rms_norm_core(
    hidden_states: torch.Tensor,
    weight: torch.Tensor | None,
    eps: float
) -> torch.Tensor

Normalize in FP32, then apply the effective gain before the output cast.

Parameters:

hidden_states
torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

weight
torch.Tensor | None

Optional effective FP32 gain of shape [hidden], on the input device.

eps
float

Stabilizer added to the mean squared activation.

Returns: torch.Tensor

Independent tensor of shape […, hidden], in the input dtype and device.

nemo_automodel.components.models.muse_spark.layers.apply_rotary_embedding(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    upcast: bool = True
) -> torch.Tensor

Apply adjacent-pair RoPE with the checkpoint’s arithmetic precision.

The original checkpoint computes the rotation in FP32. The HF release rounds both products in the Q/K dtype before adding them.

nemo_automodel.components.models.muse_spark.layers.build_muse_spark_moe_config(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    dtype: torch.dtype,
    overrides: dict[str, typing.Any] | None = None
) -> nemo_automodel.components.moe.config.MoEConfig

Translate the released Muse router and compressed-expert contract.

nemo_automodel.components.models.muse_spark.layers.calculate_gain_center(
    gain_center_type: str,
    residual_id: int
) -> float

Return the checkpoint-defined zero-centered RMSNorm gain offset.

nemo_automodel.components.models.muse_spark.layers.next_dtype(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig
) -> torch.dtype

Resolve the config’s model dtype without importing recipe utilities.

nemo_automodel.components.models.muse_spark.layers._NormArgs = ParamSpec('_NormArgs')
 
nemo_automodel.components.models.muse_spark.layers._NormOutput = TypeVar('_NormOutput')