nemo_automodel.components.models.qwen3_8_flash_next.state_dict_adapter

View as Markdown

State-dict conversion for Qwen3.8-Flash-Next and its owner-sharded Engram table.

Dense, grouped-MoE, shared-expert, and GatedDeltaNet parameters reuse the Qwen3.5-MoE checkpoint layouts. The PLE table is special: the checkpoint stores split_ngram_parts physical shard tensors while the native module registers one contiguous rank-local row range. to_hf exposes the local range as narrow views of the native parameter, so the checkpoint reader writes directly into final model storage and the 51.2B-parameter global table is never materialized. MTP checkpoint keys are ignored in both directions because the SFT target does not construct MTP.

Module Contents

Classes

NameDescription
Qwen3_8_FlashNextStateDictAdapterConvert Qwen3.8-Flash-Next checkpoints without gathering the global PLE table.

API

class nemo_automodel.components.models.qwen3_8_flash_next.state_dict_adapter.Qwen3_8_FlashNextStateDictAdapter(
config: nemo_automodel.components.models.qwen3_8_flash_next.config.Qwen3_8_FlashNextTextConfig,
moe_config: nemo_automodel.components.moe.layers.MoEConfig,
backend: nemo_automodel.components.models.common.BackendConfig,
engram_table: nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding,
dtype: torch.dtype = torch.bfloat16,
pretrained_model_name_or_path: str | None = None
)

Bases: Qwen3_5MoeStateDictAdapter

Convert Qwen3.8-Flash-Next checkpoints without gathering the global PLE table.

_end_local_checkpoint_shard
= end // self._rows_per_checkpoint_shard
_first_local_checkpoint_shard
= start // self._rows_per_checkpoint_shard
_rows_per_checkpoint_shard
_table_hf_key_pattern
_table_hf_prefix
= self._table_native_key.removesuffix('.weight')
_table_native_key
_view_loaded_native_keys
set[str] = set()
split_ngram_parts
= int(config.split_ngram_parts)
view_loaded_native_keys
set[str]

Return native parameters already populated through checkpoint views.

nemo_automodel.components.models.qwen3_8_flash_next.state_dict_adapter.Qwen3_8_FlashNextStateDictAdapter.convert_single_tensor_to_hf(
fqn: str,
tensor: typing.Any,
kwargs: typing.Any = {}
) -> list[tuple[str, typing.Any]]

Convert one native tensor, specializing the PLE and MTP entries.

nemo_automodel.components.models.qwen3_8_flash_next.state_dict_adapter.Qwen3_8_FlashNextStateDictAdapter.from_hf(
hf_state_dict: dict[str, typing.Any],
device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
kwargs: typing.Any = {}
) -> dict[str, typing.Any]

Drop PLE and MTP checkpoint keys; the table was written through views.

The PLE entries returned by :meth:to_hf alias the native parameter, so by the time DCP hands them back the table is already populated. view_loaded_native_keys records the native key as loaded.

nemo_automodel.components.models.qwen3_8_flash_next.state_dict_adapter.Qwen3_8_FlashNextStateDictAdapter.get_hf_state_dict_keys(
state_dict: dict[str, typing.Any]
) -> list[str]

Return the rank-independent global HF key set without gathering PLE.

Consolidated checkpoint planning requires the same global key list on every rank, so the one local PLE weight is replaced by all split_ngram_parts physical shard names. No data is touched.