nemo_automodel.components.models.muse_spark.hf_state_dict_adapter

Checkpoint adapter for the nested Hugging Face Muse Spark Open release.

Module Contents

Classes

Name	Description
`MuseSparkHFStateDictAdapter`	Load the HF release through rank-local, predominantly direct DCP storage.

Functions

Name	Description
`_hf_key_to_native`	Map one HF release name to the native model hierarchy.
`_native_key_to_hf`	Map one native text, vision, or audio state name to the HF release.
`_permute_rope_rows`	Convert adjacent-pair and split-half RoPE projection rows.

Data

_AUDIO_BLOCK_SUFFIXES

_EXACT_KEYS

_HF_AUDIO_BLOCK_SUFFIXES

_HF_EXACT_KEYS

_HF_TEXT_SUFFIXES

_HF_VISION_SUFFIXES

_TEXT_SUFFIXES

_VISION_SUFFIXES

API

class nemo_automodel.components.models.muse_spark.hf_state_dict_adapter.MuseSparkHFStateDictAdapter()

Bases: MuseSparkStateDictAdapter

Load the HF release through rank-local, predominantly direct DCP storage.

Expert transposes are views into final native storage, including without EP. Only Q/K row permutations allocate, bounded by the Q/K projections owned by this stage and at most one gathered projection during conversion. No full expert tensor is rebuilt after a direct load.

supports_low_memory_dcp_load

bool

Enable direct destinations only for the supported dense expert layouts.

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter.MuseSparkHFStateDictAdapter._rope_head_dim(
    native_key: str
) -> int | None

Return the head width for text/vision Q/K weights and vision Q/K biases.

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

Convert one native tensor into an HF checkpoint tensor.

Parameters:

fqn

str

Native state name.

tensor

torch.Tensor

Expert [experts, in, out], router [experts, hidden], projection [out, in], bias/norm [out], or embedding [vocab, hidden]. DTensors retain their global shapes and parameter placements.

**kwargs

str | bool | NoneDefaults to {}

Checkpointer’s exclude_key_regex and for_checkpoint_load.

Returns: list[tuple[str, torch.Tensor]]

Zero or one (HF name, tensor) pairs. Expert [experts, out, in]

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter.MuseSparkHFStateDictAdapter.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 HF text, vision and audio tensors to the native checkpoint.

Parameters:

hf_state_dict

dict[str, torch.Tensor]

HF names mapped to experts [experts, out, in], routers [experts, hidden], projections [out, in], vectors [out], and embeddings [vocab, hidden]. DCP destinations may be strided model-storage views or DTensors with global shapes.

device_mesh

DeviceMesh | NoneDefaults to None

Optional EP mesh used to slice unsharded expert input. Already-sharded DTensor experts preserve their existing mesh.

**kwargs

str | bool | NoneDefaults to {}

Unused shared checkpoint adapter arguments.

Returns: dict[str, torch.Tensor]

Native names and tensor layouts. Expert transposes and unchanged

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._hf_key_to_native(
    key: str
) -> str

Map one HF release name to the native model hierarchy.

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._native_key_to_hf(
    key: str
) -> str | None

Map one native text, vision, or audio state name to the HF release.

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._permute_rope_rows(
    tensor: torch.Tensor,
    head_dim: int,
    to_hf: bool
) -> torch.Tensor

Convert adjacent-pair and split-half RoPE projection rows.

Parameters:

tensor

torch.Tensor

Projection [heads * head_dim, hidden] or bias [heads * head_dim]. A DTensor retains its global shape and arbitrary Shard/Replicate placements. Rank-local row fragments need not span complete heads.

head_dim

int

Width of a complete attention head.

to_hf

bool

Convert adjacent pairs to split halves when true; invert otherwise.

Returns: torch.Tensor

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

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._AUDIO_BLOCK_SUFFIXES = {None: {f'{native}.{part}': f'{hf}.{part}' for native, hf in (('conv.conv', 'con...

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._EXACT_KEYS = {'lm_head.weight': 'lm_head.weight', 'model.embed_tokens.weight': 'model.languag...

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._HF_AUDIO_BLOCK_SUFFIXES = {value: key for key, value in (_AUDIO_BLOCK_SUFFIXES.items())}

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._HF_EXACT_KEYS = {value: key for key, value in (_EXACT_KEYS.items())}

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._HF_TEXT_SUFFIXES = {value: key for key, value in (_TEXT_SUFFIXES.items())}

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._HF_VISION_SUFFIXES = {value: key for key, value in (_VISION_SUFFIXES.items())}

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._TEXT_SUFFIXES = {'self_attn.pre_qkv_proj_norm.weight': 'input_layernorm.weight', 'self_attn.o_ga...

nemo_automodel.components.models.muse_spark.hf_state_dict_adapter._VISION_SUFFIXES = {None: {f'attn.{proj}_proj.{part}': f'attn.{proj}_proj.{part}' for proj in ('q',...