nemo_automodel.components.models.muse_spark.state_dict_adapter

Checkpoint conversion for the released Muse Spark Open dump.

Module Contents

Classes

Name	Description
`MuseSparkStateDictAdapter`	Convert raw grouped Muse experts and multimodal namespaces to AutoModel.

Functions

Name	Description
`_native_to_raw_key`	Map one native tensor key to the released checkpoint hierarchy.
`_raw_to_native_key`	Map one serialized checkpoint key to the native module hierarchy.

Data

_AUDIO_PREFIX

_LEGACY_VISION_PREFIX

_NATIVE_EXPERT_RE

_RAW_EXPERT_RE

_RAW_LAYER_RE

_REDUNDANT_VISION_PREFIX

_VISION_PROJECTION

API

class nemo_automodel.components.models.muse_spark.state_dict_adapter.MuseSparkStateDictAdapter(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)

Bases: StateDictAdapter

Convert raw grouped Muse experts and multimodal namespaces to AutoModel.

The checkpoint stores expert matrices as [E,out,in] while the reusable grouped-GEMM implementation consumes [E,in,out]. Under expert parallelism, conversion first selects this rank’s contiguous expert range, transposes only that slice, and wraps the local result as an EP DTensor.

nemo_automodel.components.models.muse_spark.state_dict_adapter.MuseSparkStateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    kwargs: str | bool | None = {}
) -> list[tuple[str, torch.Tensor]]

Convert a native tensor, preserving grouped expert tensors as views.

nemo_automodel.components.models.muse_spark.state_dict_adapter.MuseSparkStateDictAdapter.from_hf(
    hf_state_dict: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    kwargs: str | bool | None = {}
) -> dict[str, torch.Tensor]

Convert released keys and slice grouped experts for the EP owner rank.

nemo_automodel.components.models.muse_spark.state_dict_adapter.MuseSparkStateDictAdapter.to_hf(
    state_dict: dict[str, torch.Tensor],
    exclude_key_regex: str | None = None,
    kwargs: str | bool | None = {}
) -> dict[str, torch.Tensor]

Convert native tensors to released Muse Spark Open checkpoint keys.

nemo_automodel.components.models.muse_spark.state_dict_adapter._native_to_raw_key(
    key: str
) -> str | None

Map one native tensor key to the released checkpoint hierarchy.

nemo_automodel.components.models.muse_spark.state_dict_adapter._raw_to_native_key(
    key: str
) -> str | None

Map one serialized checkpoint key to the native module hierarchy.

nemo_automodel.components.models.muse_spark.state_dict_adapter._AUDIO_PREFIX = 'modalities.soft_audio_perception.interaction_encoder.'

nemo_automodel.components.models.muse_spark.state_dict_adapter._LEGACY_VISION_PREFIX = 'modalities.vision_perception.encoder.vision_module.'

nemo_automodel.components.models.muse_spark.state_dict_adapter._NATIVE_EXPERT_RE = re.compile('model\\.layers\\.(\\d+)\\.feed_forward\\.experts\\.(gate_and_up_proj...

nemo_automodel.components.models.muse_spark.state_dict_adapter._RAW_EXPERT_RE = re.compile('model\\.layers\\.(\\d+)\\.feed_forward\\.experts\\.(w13_weight|w2_we...

nemo_automodel.components.models.muse_spark.state_dict_adapter._RAW_LAYER_RE = re.compile('model\\.layers\\.(\\d+)\\.')

nemo_automodel.components.models.muse_spark.state_dict_adapter._REDUNDANT_VISION_PREFIX = 'modalities.vision_perception.encoder.vision_model.'

nemo_automodel.components.models.muse_spark.state_dict_adapter._VISION_PROJECTION = 'modalities.vision_perception.encoder.vision_projection.weight'