nemo_automodel.components.models.gemma4_unified.state_dict_adapter

State-dict key translation for Hugging Face-native Gemma4 Unified models.

Module Contents

Classes

Name	Description
`Gemma4UnifiedStateDictAdapter`	Translate Gemma4 Unified keys between model and published HF names.

Functions

Name	Description
`_rename_key`	Rename one state-dict key while preserving an optional `model.` prefix.
`_rename_state_dict`	Return a state dict with renamed keys and unchanged tensor values.

Data

_HF_KEY_RENAMES

_MODEL_KEY_RENAMES

API

class nemo_automodel.components.models.gemma4_unified.state_dict_adapter.Gemma4UnifiedStateDictAdapter()

Bases: StateDictAdapter

Translate Gemma4 Unified keys between model and published HF names.

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

Rename one model tensor without changing its layout or storage.

Parameters:

fqn

str

Fully qualified native model tensor name.

tensor

Any

Tensor-like value of arbitrary rank. Its shape, dtype, device, strides, and storage are preserved exactly.

**kwargs

AnyDefaults to {}

Adapter-interface arguments. exclude_key_regex may contain an optional regular expression for keys to omit.

Returns: list[tuple[str, Any]]

An empty list when excluded, otherwise one (hf_fqn, tensor) pair

nemo_automodel.components.models.gemma4_unified.state_dict_adapter.Gemma4UnifiedStateDictAdapter.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]

Rename published Hugging Face FQNs back to model FQNs.

Parameters:

hf_state_dict

dict[str, Any]

Hugging Face state mapping. Tensor values may have arbitrary rank and retain their exact shape, dtype, device, and storage.

device_mesh

DeviceMesh | NoneDefaults to None

Unused; Gemma4 Unified renames keys without changing tensor placement.

**kwargs

AnyDefaults to {}

Unused adapter-interface arguments.

Returns: dict[str, Any]

A new model-keyed mapping whose tensor values alias the inputs.

nemo_automodel.components.models.gemma4_unified.state_dict_adapter.Gemma4UnifiedStateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex: str | None = None,
    kwargs: typing.Any = {}
) -> dict[str, typing.Any]

Rename model FQNs to published Hugging Face FQNs.

Parameters:

state_dict

dict[str, Any]

Native model state mapping. Tensor values may have arbitrary rank and retain their exact shape, dtype, device, and storage.

exclude_key_regex

str | NoneDefaults to None

Optional regular expression for keys to omit.

**kwargs

AnyDefaults to {}

Unused adapter-interface arguments.

Returns: dict[str, Any]

A new Hugging Face-keyed mapping whose tensor values alias the inputs.

nemo_automodel.components.models.gemma4_unified.state_dict_adapter._rename_key(
    key: str,
    renames: dict[str, str]
) -> str

Rename one state-dict key while preserving an optional model. prefix.

nemo_automodel.components.models.gemma4_unified.state_dict_adapter._rename_state_dict(
    state_dict: dict[str, typing.Any],
    renames: dict[str, str]
) -> dict[str, typing.Any]

Return a state dict with renamed keys and unchanged tensor values.

Parameters:

state_dict

dict[str, Any]

State mapping whose tensor values retain their original shape, dtype, device, and storage.

renames

dict[str, str]

Source-prefix to target-prefix mapping.

Returns: dict[str, Any]

A new mapping whose values alias the input values exactly.

Raises:

ValueError: If two source keys map to the same output key.

nemo_automodel.components.models.gemma4_unified.state_dict_adapter._HF_KEY_RENAMES = {'embed_vision.patch_ln1': 'vision_embedder.patch_ln1', 'embed_vision.patch_dens...

nemo_automodel.components.models.gemma4_unified.state_dict_adapter._MODEL_KEY_RENAMES = {hf_key: model_key for model_key, hf_key in (_HF_KEY_RENAMES.items())}