`core.dist_checkpointing.gpt_checkpoint_interop`#

Load GPT (pure transformer) distributed checkpoints into HybridModel runs.

A GPTModel decoder layer packs self-attention and an MLP into a single TransformerLayer, so a GPT checkpoint with L layers stores both sub-modules under decoder.layers.<i>.. HybridModel gives every sub-block its own layer: in a pattern such as M*-M*- each GPT layer corresponds to one attention (‘*’) position and one MLP (‘-’ dense or ‘E’ MoE) position, while SSM (‘M’) positions have no GPT counterpart.

Rather than rewriting the checkpoint on disk, the hybrid run’s own sharded state dict is retargeted at load time:

attention and MLP entries are rewritten to the GPT checkpoint’s canonical homogeneous-layer format: the layer index is dropped from the storage key (decoder.layers.<g>.mlp... -> decoder.layers.mlp...) and the matching GPT layer index becomes a prepended sharding axis, exactly mirroring TransformerBlock.sharded_state_dict with non_homogeneous_layers=False (the format GPTModel training saves);
decoder.final_norm is pointed at GPT’s decoder.final_layernorm;
HybridModel’s empty output_layer._extra_state entry stays local because GPT checkpoints intentionally omit that backward-compatibility key;
entries of layers without a GPT counterpart are wrapped in LocalNonpersistentObject so no storage read is attempted and the freshly initialized module values are kept (and remain visible to the subsequent strict load_state_dict).

The retargeted sharded state dict is then handed to the regular dist_checkpointing.load machinery, which reads the GPT checkpoint directly and reshards across any TP/PP/EP/ETP layout change on the way.

The same retargeting also applies to the distributed optimizer’s sharded state dict. In the model-space checkpoint formats (fully_reshardable / fully_sharded_model_space) every optimizer-state ShardedTensor is built by copying the corresponding model param’s metadata and prefixing its key with optimizer.state.<state>. (see DistributedOptimizer.sharded_param_state_*). Those entries therefore carry the same decoder.layers.<i>. keys and sharding as the model tensors, so

func:: retarget_sharded_state_dict_to_gpt_checkpoint rewrites them onto the GPT checkpoint identically – optimizer moments and fp32 master params for attention/MLP layers load from the GPT run, while fresh layers (e.g. Mamba) keep their freshly initialized optimizer state via LocalNonpersistentObject.

Module Contents#

Classes#

GPTCompatLayerMaps

Correspondence between hybrid layer indices and GPT layer indices.

Functions#

`gpt_compatible_layer_maps`	Derive hybrid->GPT layer index maps from a hybrid layer pattern.
`_prepend_gpt_layer_axis`	Add the GPT layer index as the leading sharding axis of an entry.
`retarget_sharded_state_dict_to_gpt_checkpoint`	Point a hybrid model’s sharded state dict at a GPT checkpoint, in place.
`_retarget_explicit_key_to_gpt_checkpoint`	Translate one explicit HybridModel state-dict key to its GPT key.
`retarget_fsdp_state_dict_to_gpt_checkpoint`	Return an `fsdp_dtensor` model or optimizer state dict under GPT keys.

Data#

`_GPT_SOURCED_SYMBOLS`
`_FRESH_INIT_SYMBOLS`
`_DECODER_LAYER_KEY_RE`
`_GPT_FINAL_NORM_KEY_MAP`
`_GPT_OMITTED_LOCAL_KEYS`

API#

core.dist_checkpointing.gpt_checkpoint_interop._GPT_SOURCED_SYMBOLS#: ()

core.dist_checkpointing.gpt_checkpoint_interop._FRESH_INIT_SYMBOLS#: ()

core.dist_checkpointing.gpt_checkpoint_interop._DECODER_LAYER_KEY_RE#: ‘compile(…)’

core.dist_checkpointing.gpt_checkpoint_interop._GPT_FINAL_NORM_KEY_MAP#: None

core.dist_checkpointing.gpt_checkpoint_interop._GPT_OMITTED_LOCAL_KEYS#: (‘output_layer._extra_state’,)

class core.dist_checkpointing.gpt_checkpoint_interop.GPTCompatLayerMaps#

Correspondence between hybrid layer indices and GPT layer indices.

.. attribute:: attention_to_gpt

hybrid global layer index of the i-th attention position -> GPT layer index i.

.. attribute:: mlp_to_gpt

hybrid global layer index of the i-th MLP-bearing position (‘-’ or ‘E’) -> GPT layer index i.

.. attribute:: fresh_init

hybrid global layer indices with no GPT counterpart; their modules keep the run’s fresh initialization.

.. attribute:: num_gpt_layers

number of layers the source GPT checkpoint must have.

attention_to_gpt: Mapping[int, int]#: None

mlp_to_gpt: Mapping[int, int]#: None

fresh_init: frozenset#: None

num_gpt_layers: int#: None

core.dist_checkpointing.gpt_checkpoint_interop.gpt_compatible_layer_maps( hybrid_layer_pattern: str, ) → core.dist_checkpointing.gpt_checkpoint_interop.GPTCompatLayerMaps#

Derive hybrid->GPT layer index maps from a hybrid layer pattern.

Parameters:: hybrid_layer_pattern – the run’s unified hybrid layer pattern (pipeline ‘|’ separators allowed).
Returns:: GPTCompatLayerMaps for retargeting a sharded state dict.
Raises:: ValueError – if the pattern cannot be paired one-to-one with a GPT checkpoint layout (MTP present, non-translatable symbols, mixed dense/MoE positions, or unbalanced ‘*’ vs MLP counts).

core.dist_checkpointing.gpt_checkpoint_interop._prepend_gpt_layer_axis( entry, gpt_layer_idx: int, num_gpt_layers: int, )#

Add the GPT layer index as the leading sharding axis of an entry.

Mirrors what TransformerBlock.sharded_state_dict does for homogeneous layers by passing sharded_offsets=[(0, layer_idx, num_layers)] down to make_sharded_tensors_for_checkpoint:

ShardedTensor: one more prepended axis of size num_gpt_layers at position 0, this shard sitting at gpt_layer_idx;
ShardedObject: (1,)/(0,) placeholder offsets (from _get_extra_state_offsets with no offsets) are replaced by the layer axis, otherwise the layer axis is prepended (e.g. before an expert axis);
ShardedTensorFactory: the built sub-entries get the same treatment.

core.dist_checkpointing.gpt_checkpoint_interop.retarget_sharded_state_dict_to_gpt_checkpoint( sharded_state_dict: megatron.core.dist_checkpointing.mapping.ShardedStateDict, layer_maps: core.dist_checkpointing.gpt_checkpoint_interop.GPTCompatLayerMaps, ) → None#

Point a hybrid model’s sharded state dict at a GPT checkpoint, in place.

Only the storage lookup metadata (key and sharding axes) of each ShardedBase entry is rewritten into the GPT checkpoint’s homogeneous layer format; the nested state dict structure (used by the subsequent load_state_dict) keeps the hybrid model’s own names. Entries of layers with no GPT counterpart are replaced by LocalNonpersistentObject so the loaded state dict returns their current (freshly initialized) values.

The same routine handles the distributed optimizer’s sharded state dict: its per-parameter entries embed the model key (optimizer.state.<state>.decoder. layers.<i>...) and mirror the model param’s sharding, so they retarget the same way, and fresh-layer optimizer state is likewise kept local.

Parameters:

sharded_state_dict – one model chunk’s sharded state dict (as produced by model.sharded_state_dict()) or the matching optimizer sharded state dict.
layer_maps – maps from :func:gpt_compatible_layer_maps derived from the same pattern the model was built with.

core.dist_checkpointing.gpt_checkpoint_interop._retarget_explicit_key_to_gpt_checkpoint( key: Any, layer_maps: core.dist_checkpointing.gpt_checkpoint_interop.GPTCompatLayerMaps, checkpoint_keys: Iterable[str] | None = None, ) → Any | None#

Translate one explicit HybridModel state-dict key to its GPT key.

fsdp_dtensor checkpoints store explicit parameter names rather than homogeneous-layer ShardedTensor metadata. Returning None omits a fresh-only or GPT-omitted entry from the DCP load plan while leaving its existing HybridModel value untouched.

core.dist_checkpointing.gpt_checkpoint_interop.retarget_fsdp_state_dict_to_gpt_checkpoint( state_dict: Mapping[Any, Any], layer_maps: core.dist_checkpointing.gpt_checkpoint_interop.GPTCompatLayerMaps, checkpoint_keys: Iterable[str] | None = None, checkpoint_prefix: str = '', ) → dict[Any, Any]#

Return an fsdp_dtensor model or optimizer state dict under GPT keys.

FSDP model state is a flat parameter-name mapping. Distributed-optimizer state can contain nested state and param_to_group_meta mappings (and chained-optimizer integer keys), so the translation recursively rewrites every parameter-name key while preserving the DTensor leaves.

core.dist_checkpointing.gpt_checkpoint_interop#