kempnerforge.checkpoint.expert_layout

Expert-parallel experts under their global identity in a DCP state dict.

apply_expert_parallel leaves each EP rank num_experts / ep experts of every MoE layer: an unpacked layer keeps them in a ModuleList numbered from 0, and a packed layer keeps their rows of each packed weight. Every EP rank’s state dict then names and shapes its experts as if they were the layer’s only ones, and DCP, which writes one copy of each name and offset, would save one rank’s experts in place of all of them.

ExpertLayout translates a state dict so that each rank presents its experts as part of all num_experts: an unpacked expert under its global index, and a packed weight as a DTensor over ep and the FSDP mesh whose rows cover every expert. A checkpoint then holds every expert in the layout an ep = 1 run saves, so DCP loads it at any ep. The translated tensors share storage with the rank’s own, so nothing is copied or gathered, and a load into them loads the model and optimizer.

check_saved_experts refuses a checkpoint that holds fewer experts than the model, as one saved with ep > 1 before this layout existed does.

Functions

check_saved_experts(model, saved, where)	Refuse a checkpoint whose MoE layers hold fewer experts than model's.
expert_layout(model)	The layout of model's expert-parallel layers; None when no layer splits its experts.
moe_layers(model)	Each MoEMLP in model with the name its state-dict entries start with.

Classes

ExpertLayout

Translates a DCP state dict between one EP rank's experts and all of the layer's.

kempnerforge.checkpoint.expert_layout.moe_layers(model)

Each MoEMLP in model with the name its state-dict entries start with.

The name is empty when model is itself the MoEMLP. It leaves out the activation-checkpoint and torch.compile wrappers, as get_model_state_dict does.

Parameters:

model (torch.nn.Module)

Return type:

Iterator[tuple[str, MoEMLP]]

kempnerforge.checkpoint.expert_layout.expert_layout(model)

The layout of model’s expert-parallel layers; None when no layer splits its experts.

Parameters:

model (torch.nn.Module)

Return type:

ExpertLayout | None

class kempnerforge.checkpoint.expert_layout.ExpertLayout

Bases: object

Translates a DCP state dict between one EP rank’s experts and all of the layer’s.

to_global translates the "model" and "optimizer" entries of a state dict, and the names in every other entry that maps parameter names to values, such as the never-stepped record. to_local undoes it once a load has filled the result, so one layout serves one save or one load.

Parameters:

layers – The split MoEMLP layers, each with the name its state-dict entries start with, as moe_layers gives them.

__init__(layers)

Parameters:

layers (list[tuple[str, MoEMLP]])

Return type:

None

to_global(state)

state with each rank’s experts placed among all of the layer’s experts.

Returns a new dict and leaves state as it is. Each tensor in the result is the one it replaces or a view of it, so a load into the result loads them.

Raises:

ValueError – The model state lacks the expert weights of a split layer, or the optimizer holds a tensor for a packed weight in a shape other than the weight’s, so the checkpoint could not hold every rank’s experts.

Parameters:

state (Mapping[str, Any])

Return type:

dict[str, Any]

to_local(state)

Undo to_global on the state dict that a load filled.

Names go back to this rank’s, and each view to_global made goes back to the tensor it views, which holds what the load wrote. Each optimizer param group lists this rank’s parameters again and keeps the loaded settings.

Parameters:

state (Mapping[str, Any])

Return type:

dict[str, Any]

kempnerforge.checkpoint.expert_layout.check_saved_experts(model, saved, where)

Refuse a checkpoint whose MoE layers hold fewer experts than model’s.

saved is the checkpoint’s DCP index, its state_dict_metadata. A checkpoint saved with ep > 1 before the experts had a global identity holds num_experts / ep experts of each layer: every EP rank wrote its own under the same names, and the checkpoint kept one rank’s. Loading it at the same ep would give every rank that rank’s experts, and at another ep it fails on the experts it lacks, so it is refused here with what it holds. A layer the checkpoint holds no expert of is left to the load.

Raises:

ValueError – A layer of model has more experts than the checkpoint holds for it.

Parameters:

• model (torch.nn.Module)

• saved (Mapping[str, Any])

• where (object)

Return type:

None