nemo_automodel.components.models.muse_spark.config

Configuration for Muse Spark Open Omni checkpoints.

Module Contents

Classes

Name	Description
`MuseSpark1xOmniConfig`	Configuration matching the released Muse Spark Open Omni dump.

Functions

Name	Description
`expand_global_attention_schedule`	Expand a periodic attention-window schedule to one entry per layer.

API

class nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig(
    vocab_size: int = 202048,
    hidden_size: int = 8192,
    num_hidden_layers: int = 62,
    schedule_num_hidden_layers: int | None = None,
    num_attention_heads: int = 128,
    num_key_value_heads: int = 16,
    head_dim: int = 64,
    max_position_embeddings: int = 1048576,
    rms_norm_eps: float = 1e-05,
    rms_post_norm_eps: float = 1e-08,
    rope_theta: float = 500000.0,
    use_qk_norm: bool = True,
    use_rope: bool = True,
    use_sqrt_attn_scaling: bool = False,
    every_n_layers_nope: int | None = 4,
    global_attn_cfg: list[int] | None = None,
    use_attn_o_gate: bool = True,
    use_attn_o_norm: bool = True,
    ffn_dim_multiplier: float = 2.0,
    non_linearity: str = 'swiglu',
    use_ffn_learnable_scales: bool = True,
    norm_affine: bool = True,
    post_norm_affine: bool = False,
    norm_zero_centered_gamma: bool = True,
    post_norm_gain_center_type: str = 'one',
    output_norm_gain_center_type: str = 'zero',
    normalize_tok_embeddings: bool = True,
    use_residual_gate: bool = True,
    residual_gate_temperature: float = 0.3,
    residual_dtype: str = 'fp32',
    output_soft_cap_temp: float | None = 20.0,
    use_fp32_for_output: bool = False,
    moe: dict[str, typing.Any] | None = None,
    metap: dict[str, typing.Any] | None = None,
    numerics_args: dict[str, typing.Any] | None = None,
    vision_config: dict[str, typing.Any] | None = None,
    audio_perception_modality_args: dict[str, typing.Any] | None = None,
    use_audio_perception: bool = True,
    image_token_id: int = 200090,
    patch_token_id: int = 200092,
    video_token_id: int = 200091,
    audio_token_id: int = 200038,
    initializer_range: float = 0.02,
    attention_dropout: float = 0.0,
    attention_bias: bool = False,
    tie_word_embeddings: bool = False,
    pad_token_id: int | None = None,
    bos_token_id: int = 200000,
    eos_token_id: int = 200001,
    kwargs: typing.Any = {}
)

Bases: PretrainedConfig

Configuration matching the released Muse Spark Open Omni dump.

audio_perception_modality_args

= dict(audio_perception_modality_args)

expert_hidden_size

int

Return the pooler-compressed expert input/output width.

global_attn_cfg

= list(global_attn_cfg)

intermediate_size

= int(hidden_size * ffn_dim_multiplier)

keys_to_ignore_at_inference

= ['past_key_values']

layer_sliding_windows

= [0] * num_hidden_layers

layer_types

= ['full_attention'] * num_hidden_layers

metap

= dict(metap)

model_type

= 'musespark1x_omni'

moe

= dict(moe)

moe_intermediate_size

int

Return each routed expert’s SwiGLU intermediate width.

no_rope_layers

num_experts_per_tok

int

Return the number of selected experts per token.

num_local_experts

int

Return the total routed-expert count before EP sharding.

numerics_args

= dict(numerics_args or {})

output_multiplier

float

Return the checkpoint’s language-model output multiplier.

patch_token_id

router_output_multiplier

float

Return the MetaP multiplier applied to logits before sigmoid routing.

scale_query_by

float

Return the post-QK-normalization query scale.

schedule_num_hidden_layers

= schedule_num_hidden_layers or num_hidden_layers

sliding_window

= max(self.global_attn_cfg, default=0)

vision_config

= dict(vision_config)

vision_navit_config

dict[str, Any]

Return the flattened NaViT configuration mapping.

nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig._validate() -> None

Reject checkpoint settings unsupported by the native implementation.

nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig.layer_uses_rope(
    layer_idx: int
) -> bool

Return whether decoder layer layer_idx applies text RoPE.

nemo_automodel.components.models.muse_spark.config.expand_global_attention_schedule(
    schedule: list[int] | tuple[int, ...] | None,
    num_hidden_layers: int,
    every_n_layers_nope: int | None
) -> list[int] | None

Expand a periodic attention-window schedule to one entry per layer.

Parameters:

schedule

list[int] | tuple[int, ...] | None

Periodic window sizes. Zero denotes global attention.

num_hidden_layers

int

Number of decoder layers.

every_n_layers_nope

int | None

Period used to align no-RoPE layers from the end.

Returns: list[int] | None

A list of num_hidden_layers window sizes, or None when no