nemo_automodel.components.models.flint.state_dict_adapter

Strict Flint checkpoint conversion using the shared split-expert layout.

Module Contents

Classes

Name	Description
`FlintStateDictAdapter`	Convert split BF16 base or FP32 resume weights without losing source names.
`_LoadAlias`	Hold one destination’s storage alive and identify its exact tensor view.

Data

API

class nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter(
    config: nemo_automodel.components.models.flint.config.FlintConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.float32
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert split BF16 base or FP32 resume weights without losing source names.

HF experts are [intermediate, hidden] gate/up and [hidden, intermediate] down tensors. Native grouped weights are [experts, hidden, 2 * intermediate] (gate before up) and [experts, intermediate, hidden]. Shared expert paths and the FP32 sink holder are renamed; all other ordinary layouts are unchanged.

from_hf consumes its input mapping, as the shared merging implementation does. Runtime storage remains FP32 when requested. to_hf never casts training weights to BF16: consolidated-export dtype belongs to the shared writer’s original-checkpoint metadata. Export may allocate contiguous per-expert tensors; its memory use is distinct from direct loading.

_hf_shapes

= self._checkpoint_shapes()

_native_shapes

= self._native_tensor_shapes()

_pending_expert_aliases

dict[str, dict[str, _LoadAlias]] = {}

_pending_load_aliases

dict[str, _LoadAlias] = {}

_pending_load_shapes

dict[str, tuple[int, ...]] | None = None

_pipeline_native_keys

frozenset[str] | None = None

_source_hf_shapes

supports_low_memory_dcp_load

bool

Opt in only for the tested single-process model-backed Torch weights.

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._checkpoint_shapes(
    num_layers: int | None = None
) -> dict[str, tuple[int, ...]]

Enumerate source HF shapes without allocating any checkpoint tensors.

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._convert_lora_from_hf(
    tensors: collections.abc.Mapping[str, torch.Tensor]
) -> dict[str, torch.Tensor]

Validate and join complete expert PEFT families using the shared mixin.

Parameters:

tensors

Mapping[str, torch.Tensor]

Adapter-only HF tensors. Linear A/B weights have shape [rank,input]/[output,rank]; per-expert gate/up/down use those layouts independently. Meta tensors validate metadata without payload reads; materialized duplicated gate/up A values must be identical because they share one native adapter.

Returns: dict[str, torch.Tensor]

Native linear and grouped expert LoRA weights in layouts documented

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._hf_name(
    native_name: str
) -> str

staticmethod

Reverse only the two Flint-owned ordinary module path changes.

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._is_backend_extra_state(
    name: str
) -> bool

Identify only instantiated TE bookkeeping, never source HF weights.

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._native_lora_shapes(
    tensors: collections.abc.Mapping[str, torch.Tensor]
) -> dict[str, tuple[int, ...]]

Validate the shapes of existing native linear and grouped LoRA tensors.

Parameters:

tensors

Mapping[str, torch.Tensor]

Native BF16/FP32 tensors, optionally carrying the shared PEFT base_model.model. prefix. Linear A/B weights are [rank,input]/[output,rank]; grouped expert A/B weights are [experts,input,rank]/[experts,rank,output]. DTensor shapes are global and retain their expert-axis placements.

Returns: dict[str, tuple[int, ...]]

Exact allowed LoRA shapes derived from the corresponding known

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._native_name(
    hf_name: str
) -> str

staticmethod

Translate ordinary HF paths into the native shared-expert/sink paths.

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._native_tensor_shapes() -> dict[str, tuple[int, ...]]

Describe global native layouts, including the two grouped projections.

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter._validate_tensors(
    tensors: dict[str, torch.Tensor],
    shapes: dict[str, tuple[int, ...]]
) -> None

staticmethod

Reject unknown names, incompatible global shapes and unsupported dtypes.

Parameters:

tensors

dict[str, torch.Tensor]

Mapping of parameter names to BF16/FP32 tensors. Ordinary HF tensors use [out, in] matrices or one-dimensional norm/ sink scales. Native grouped tensors use [experts, hidden, 2 * intermediate] and [experts, intermediate, hidden]. DTensor shapes are global, retaining their original placements.

shapes

dict[str, tuple[int, ...]]

Required global shape for every allowed mapping key.

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

Convert one global native weight using the shared expert machinery.

Parameters:

fqn

str

Exact native model parameter name.

tensor

torch.Tensor

BF16/FP32 native tensor. Ordinary shapes are unchanged; grouped shapes are [experts, hidden, 2 * intermediate] or [experts, intermediate, hidden]. DTensors retain global shapes and the shared mixin’s adjusted rank-local placements.

**kwargs

AnyDefaults to {}

Shared conversion flags, including for_checkpoint_load.

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

HF name/tensor pairs; plain Torch load-mode tensors alias the

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

Validate and consume HF weights, joining only the owned expert rows.

Parameters:

hf_state_dict

dict[str, torch.Tensor]

Complete BF16 base or FP32 resume mapping in HF [out, in] layouts, consumed by this call. After to_hf in load mode, keys must exactly match that call’s DCP destinations. DTensors retain global shapes and shard placements.

device_mesh

DeviceMesh | NoneDefaults to None

Optional shared expert mesh selecting global expert IDs owned by this rank; no distributed capability is advertised.

**kwargs

AnyDefaults to {}

Compatibility arguments of the shared adapter interface.

Returns: dict[str, torch.Tensor]

Native mapping with FP32 sinks/head, unchanged ordinary layouts and

nemo_automodel.components.models.flint.state_dict_adapter.FlintStateDictAdapter.to_hf(
    state_dict: dict[str, torch.Tensor],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    for_checkpoint_load: bool = False,
    kwargs: typing.Any = {}
) -> dict[str, torch.Tensor]

Expose HF tensors, preserving resident precision and load aliases.

Parameters:

state_dict

dict[str, torch.Tensor]

Native weights with the global layouts documented on this adapter. Tensor dtype/device/placements are preserved; load destinations alias native storage on the plain Torch path.

exclude_key_regex

str | NoneDefaults to None

HF/native names to omit, typically backend state.

quantization

boolDefaults to False

Must be false: the supplied checkpoint is unquantized.

for_checkpoint_load

boolDefaults to False

Build a complete set of destination views for DCP. Save/export callers leave this false; grouped exports may allocate contiguous tensors without mutating source values.

**kwargs

AnyDefaults to {}

Shared checkpoint protocol arguments forwarded to the expert mixin, including device mesh and serialization options.

Returns: dict[str, torch.Tensor]

HF-keyed tensors in split-expert [out, in] layout. Plain Torch

class nemo_automodel.components.models.flint.state_dict_adapter._LoadAlias(
    storage: torch.UntypedStorage,
    offset: int,
    shape: tuple[int, ...],
    stride: tuple[int, ...],
    dtype: torch.dtype,
    device: torch.device
)

Dataclass

Hold one destination’s storage alive and identify its exact tensor view.

device

dtype

offset

int

shape

tuple[int, ...]

storage

UntypedStorage

stride

tuple[int, ...]

nemo_automodel.components.models.flint.state_dict_adapter._LoadAlias.capture(
    tensor: torch.Tensor
) -> nemo_automodel.components.models.flint.state_dict_adapter._LoadAlias | None

classmethod

Identify local storage without copying any parameter payload.

Parameters:

tensor

torch.Tensor

Ordinary HF-shaped destination or DTensor with that global shape; the local tensor retains its existing strided storage.

Returns: _LoadAlias | None

Storage/view identity, or None for a shape-only meta destination.

nemo_automodel.components.models.flint.state_dict_adapter._LoadAlias.matches(
    tensor: torch.Tensor | None
) -> bool

Check the exact destination storage, offset, layout, dtype and device.

Parameters:

tensor

torch.Tensor | None

Current HF expert tensor, possibly a different Python view of the original destination. None denotes an absent key.

Returns: bool

Whether this view still aliases the recorded destination exactly.

nemo_automodel.components.models.flint.state_dict_adapter._DTYPES = (torch.bfloat16, torch.float32)

nemo_automodel.components.models.flint.state_dict_adapter._EXPERT_KEY = re.compile('^model\\.layers\\.(\\d+)\\.mlp\\.experts\\.(\\d+)\\.(gate_proj|up_pr...

nemo_automodel.components.models.flint.state_dict_adapter._LORA_EXPERT_KEY = re.compile('^(?:base_model\\.model\\.)?model\\.layers\\.(\\d+)\\.mlp\\.experts\\...

nemo_automodel.components.models.flint.state_dict_adapter._LORA_LINEAR_KEY = re.compile('^(.+)\\.(lora_A|lora_B)\\.weight$')