`core.transformer.fsdp_dtensor_checkpoint`#

Module Contents#

Functions#

`get_ep_layer_offset`	Get the expert layer offset for the current model.
`get_total_num_experts`	Get the total number of experts for the current model.
`get_expert_index_from_key`	Extract expert index from various expert key formats.
`handle_experts_in_state_dict`	Rewrite expert keys in state dict.
`expert_param_local_key`	Get the module parameter corresponding to the key.
`handle_swiglu_in_state_dict`	Handle SWiGLU in model and optimizer state dicts.
`handle_gdn_in_state_dict`	Handle GDN (Gated DeltaNet) fused projections in model and optimizer state dicts.
`handle_fp8_extra_state_case`	Handle the case where FP8 extra state is present in the model state dict.
`flatten_state_dict`	Recursively flattens a nested state dict into a single-level dict with keys
`print_diff_in_state_dicts`	Print the differences between two state dicts: metadata state dict and load state dict. This function compares the keys and shapes of the tensors in both dicts.
`validate_loaded_state_dict`	Validate the loaded state dict against the expected structure and types.
`get_global_unique_param_name`	Get the global unique parameter name for a given model and parameter.

Data#

logger

API#

core.transformer.fsdp_dtensor_checkpoint.logger#: ‘getLogger(…)’

core.transformer.fsdp_dtensor_checkpoint.get_ep_layer_offset(num_experts: int | None = None) → int#

Get the expert layer offset for the current model.

Parameters:: num_experts – Total number of experts in the model. If None, returns 0.
Returns:: The expert layer offset for the current EP rank.

core.transformer.fsdp_dtensor_checkpoint.get_total_num_experts(num_experts: int | None = None) → int#

Get the total number of experts for the current model.

Parameters:: num_experts – Total number of experts in the model. If None, returns 0.
Returns:: The total number of experts.

core.transformer.fsdp_dtensor_checkpoint.get_expert_index_from_key(key)#

Extract expert index from various expert key formats.

Supported formats:

GroupedMLP: ‘mlp.experts.linear_fc1.weight0’, ‘mlp.experts.linear_fc2.weight0’
SequentialMLP: ‘mlp.experts.local_experts.0.linear_fc1.weight’, ‘mlp.experts.local_experts.0.linear_fc2.weight’

Returns:: Expert index if found, None otherwise.
Return type:: int

core.transformer.fsdp_dtensor_checkpoint.handle_experts_in_state_dict( state_dict, num_experts: int | None = None, )#

Rewrite expert keys in state dict.

Parameters:

state_dict – The state dictionary to process.
num_experts – Total number of experts in the model. If None, no expert processing occurs.

Returns:

The processed state dictionary with rewritten expert keys.

core.transformer.fsdp_dtensor_checkpoint.expert_param_local_key( key: str, num_experts: int | None = None, ) → str#

Get the module parameter corresponding to the key.

Parameters:

key – The parameter key to process.
num_experts – Total number of experts in the model. If None, no expert processing occurs.

Returns:

The local parameter key with adjusted expert indices.

core.transformer.fsdp_dtensor_checkpoint.handle_swiglu_in_state_dict( model, model_state_dict, optimizer_state_dict, )#

Handle SWiGLU in model and optimizer state dicts.

Only splits linear_fc1 parameters whose parent TransformerLayer has config.gated_linear_unit == True. This is critical for heterogeneous models (e.g. VLMs) where the vision encoder uses GELU while the language decoder uses SWiGLU — splitting non-SWiGLU fc1 weights would create _w/_v keys that don’t exist in the checkpoint, causing a load-time mismatch.

core.transformer.fsdp_dtensor_checkpoint.handle_gdn_in_state_dict( model, model_state_dict, optimizer_state_dict, )#

Handle GDN (Gated DeltaNet) fused projections in model and optimizer state dicts.

GDN layers fuse query/key/value/gate/beta/alpha projections into a single in_proj.weight ColumnParallelLinear, and query/key/value into conv1d (weight + optional bias). For FSDP checkpoints these fused tensors must be split back into their constituent sub-tensors so that each can be independently TP-sharded — otherwise loading a checkpoint written at TP=M into TP=N would slice across logical component boundaries.

This is analogous to :func:handle_swiglu_in_state_dict which splits linear_fc1 into weight_w / weight_v.

Sub-key naming follows GatedDeltaNet.sharded_state_dict()::

in_proj.weight  → .query / .key / .value / .z / .beta / .alpha   (6-way)
conv1d.weight   → .query / .key / .value                         (3-way)
conv1d.bias     → .query / .key / .value                         (3-way)

core.transformer.fsdp_dtensor_checkpoint.handle_fp8_extra_state_case(model_state_dict)#: Handle the case where FP8 extra state is present in the model state dict.

core.transformer.fsdp_dtensor_checkpoint.flatten_state_dict(obj, parent_key='', sep='.')#: Recursively flattens a nested state dict into a single-level dict with keys

core.transformer.fsdp_dtensor_checkpoint.print_diff_in_state_dicts( state_dict_metadata, load_state_dict, limit=100, )#: Print the differences between two state dicts: metadata state dict and load state dict. This function compares the keys and shapes of the tensors in both dicts.

core.transformer.fsdp_dtensor_checkpoint.validate_loaded_state_dict(state_dict, checkpoint_path)#: Validate the loaded state dict against the expected structure and types.

core.transformer.fsdp_dtensor_checkpoint.get_global_unique_param_name(model_chunks, param)#

Get the global unique parameter name for a given model and parameter.

Parameters:

model_chunks – List of model chunks to search for the parameter.
param – The parameter to find the name for.

Returns:

The global unique parameter name.

core.transformer.fsdp_dtensor_checkpoint#

Module Contents#

Functions#

Data#

API#

`core.transformer.fsdp_dtensor_checkpoint`#