nemo_automodel.components.models.muse_spark

Native Muse Spark Open Omni model support.

Submodules

Package Contents

Classes

Name	Description
`MuseSpark1xOmniConfig`	Configuration matching the released Muse Spark Open Omni dump.
`MuseSparkProcessor`	Combine the Muse Spark Open tokenizer with native multimodal preprocessing.

Functions

Name	Description
`build_muse_spark_processor`	Build the native processor directly from a Muse Spark Open checkpoint.

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.

class nemo_automodel.components.models.muse_spark.processing.MuseSparkProcessor(
    image_processor: nemo_automodel.components.models.muse_spark.image_processing.MuseSparkImageProcessor,
    tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase,
    chat_template: str | None = None,
    image_token: str = '<image>',
    patch_token: str = '<|patch|>',
    image_start_token: str = '<|image_start|>',
    image_end_token: str = '<|image_end|>',
    video_token: str = '<video>',
    video_patch_token: str = '<|video|>',
    video_start_token: str = '<|vid_start|>',
    video_end_token: str = '<|vid_end|>',
    video_separator_token: str = '<|vid_frame_separator|>',
    audio_token: str = '<audio>',
    audio_patch_token: str = '<|audio|>',
    audio_start_token: str = '<|soft_audio_start|>',
    audio_end_token: str = '<|soft_audio_end|>',
    timestamp_format: str = 'Time: {s:.1f}s'
)

Bases: ProcessorMixin

Combine the Muse Spark Open tokenizer with native multimodal preprocessing.

audio_token_id

= tokenizer.convert_tokens_to_ids(audio_patch_token)

image_token_id

= tokenizer.convert_tokens_to_ids(patch_token)

model_input_names

list[str]

Return model inputs produced by this processor.

video_token_id

= tokenizer.convert_tokens_to_ids(video_patch_token)

nemo_automodel.components.models.muse_spark.processing.MuseSparkProcessor.__call__(
    images: typing.Any | None = None,
    text: str | list[str] | None = None,
    videos: typing.Any | None = None,
    audio: typing.Any | None = None,
    kwargs: typing.Any = {}
) -> transformers.feature_extraction_utils.BatchFeature

Tokenize prompts after expanding image/video/audio placeholders.

nemo_automodel.components.models.muse_spark.processing.build_muse_spark_processor(
    pretrained_model_name_or_path: str,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.muse_spark.processing.MuseSparkProcessor

Build the native processor directly from a Muse Spark Open checkpoint.