Skip to content

Model State Mapper

About

The d9d.model_state.mapper package is a declarative, graph-based framework for transforming model states. You describe which checkpoint keys map to which model keys, and d9d runs the transformation while it streams the checkpoint.

Core Concept

Loading a large model is rarely a 1-to-1 key match. Common problems are:

  • Naming mismatches: Hugging Face uses model.layers.0, your model uses transformer.h.0.
  • Shape mismatches: The checkpoint stores q, k and v separately, but your model expects one stacked qkv tensor.
  • Scale: The checkpoint takes hundreds of GiB. You cannot load the whole dictionary on every GPU to process it.

A mapper does not loop over tensors and modify them. It treats the transformation as a directed acyclic graph (DAG). Each mapper declares its dependency groups: which input keys produce which output keys. d9d uses these groups to load, transform and save a checkpoint in a stream, without holding the whole checkpoint in memory.

Mappers come in three kinds:

  • Leaf mappers (d9d.model_state.mapper.leaf) transform individual tensors: rename, stack, chunk, transpose, distribute.
  • Composite mappers (d9d.model_state.mapper.compose) combine other mappers: in parallel, in sequence, under a key prefix or restricted to one shard.
  • Adapters (d9d.model_state.mapper.adapters) build identity mappers from a module or from another mapper.

Usage

Pass-Through Mapping for a PyTorch Module

Use identity_mapper_from_module when checkpoint keys match the model's state dict keys, as in standard load_state_dict. You still get d9d's streaming and sharding.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
import torch.nn as nn
from d9d.model_state.mapper.adapters import identity_mapper_from_module

model = nn.Sequential(
    nn.Linear(10, 10),
    nn.ReLU(),
    nn.Linear(10, 5)
)

# Creates identity mappers for "0.weight", "0.bias", "2.weight" and "2.bias".
mapper = identity_mapper_from_module(model)

Leaf Mappers

This example merges separate query, key and value tensors into a single tensor.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch
from d9d.model_state.mapper.leaf import ModelStateMapperStackTensors

stack_mapper = ModelStateMapperStackTensors(
    source_names=["attn.q.weight", "attn.k.weight", "attn.v.weight"],
    target_name="attn.qkv.weight",
    dim=0
)

# Show what this mapper needs and produces.
print(stack_mapper.state_dependency_groups())
# Output (abridged): frozenset({StateGroup(inputs=frozenset({'attn.q.weight', ...}), outputs=frozenset({'attn.qkv.weight'}))})

# Run the transformation.
dummy_data = {
    "attn.q.weight": torch.randn(64, 64),
    "attn.k.weight": torch.randn(64, 64),
    "attn.v.weight": torch.randn(64, 64),
}
result = stack_mapper.apply(dummy_data)
print(result["attn.qkv.weight"].shape)
# Output: torch.Size([3, 64, 64])

Composing Pipelines

A full model conversion processes many keys in parallel and can chain operations, e.g. rename, then stack.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
from d9d.model_state.mapper.compose import ModelStateMapperSequential, ModelStateMapperParallel
from d9d.model_state.mapper.leaf import ModelStateMapperRename, ModelStateMapperStackTensors

mapper = ModelStateMapperSequential([
    # Step 1: Rename keys to a short format.
    ModelStateMapperParallel([
        ModelStateMapperRename("bert.encoder.layer.0.attention.self.query.weight", "layer.0.q"),
        ModelStateMapperRename("bert.encoder.layer.0.attention.self.key.weight", "layer.0.k"),
        ModelStateMapperRename("bert.encoder.layer.0.attention.self.value.weight", "layer.0.v"),
    ]),
    # Step 2: Stack them into one attention tensor.
    ModelStateMapperStackTensors(
        source_names=["layer.0.q", "layer.0.k", "layer.0.v"],
        target_name="layer.0.qkv",
        dim=0
    )
])

API Reference

d9d.model_state.mapper

Core components of the model state mapping system.

ModelStateMapper

Bases: ABC

Abstract base class for all model state transformations.

A mapper separates the declaration of a transformation from its execution:

  1. Declarative (topology): state_dependency_groups() announces what the mapper does without touching data. This lets the system build execution graphs, validate chains, detect collisions and shard work before it allocates memory.
  2. Imperative (execution): apply() runs the PyTorch operations on model states.
apply(group) abstractmethod

Executes the transformation on the tensors of one dependency group.

The caller guarantees that group contains all keys listed in inputs of the active StateGroup. Implementations must return all keys listed in its outputs.

Parameters:

Name Type Description Default
group dict[str, Tensor]

The source tensors. Keys match StateGroup.inputs.

 required

Returns:

Type Description
dict[str, Tensor]

The transformed tensors. Keys must exactly match StateGroup.outputs.
state_dependency_groups() abstractmethod

Returns the set of independent dependency groups this mapper handles.

Returns:

Type Description
frozenset[StateGroup]

A frozenset of StateGroup objects. Each group is a disjoint operation. For example, a mapper that
frozenset[StateGroup]

renames ten independent tensors returns ten distinct groups, so they can be sharded or processed
frozenset[StateGroup]

individually.

StateGroup dataclass

Atomic unit of dependency in the model state transformation graph.

A StateGroup binds a set of input keys (source) to a set of output keys (destination).

Attributes:

Name Type Description
inputs frozenset[str]

The complete set of keys required from the source state dictionary to satisfy this dependency.
outputs frozenset[str]

The complete set of keys that will be produced as a result of this transformation.

d9d.model_state.mapper.adapters

Adapters that build simple ModelStateMapper instances from PyTorch modules or other mappers.

identity_mapper_from_mapper_outputs(mapper)

Creates an identity mapper for every output key of the given mapper.

Parameters:

Name Type Description Default
mapper ModelStateMapper

The mapper whose state_dependency_groups() outputs are passed through.

 required

Returns:

Type Description
ModelStateMapper

A composite mapper that passes through every key produced by mapper.

identity_mapper_from_module(module)

Creates an identity mapper for every key in the state dict of a PyTorch module.

Use it when the checkpoint keys exactly match the module's state dict keys, as in standard load_state_dict behavior.

Parameters:

Name Type Description Default
module Module

The instantiated PyTorch module to inspect.

 required

Returns:

Type Description
ModelStateMapper

A composite identity mapper for the module's state.

d9d.model_state.mapper.compose

ModelStateMapper implementations that compose other mappers.

ModelStateMapperParallel

Bases: ModelStateMapper

Mapper that executes a list of state mappers independently of each other.

No two mappers can consume the same input key or produce the same output key. apply() routes each group to the mapper that declared it.

__init__(mappers)

Constructs the ModelStateMapperParallel object.

Parameters:

Name Type Description Default
mappers Sequence[ModelStateMapper]

The mappers to run side by side. Mappers without inputs and outputs are dropped.

 required

Raises:

Type Description
ValueError

If two mappers share an input key or an output key.

ModelStateMapperPrefixScope

Bases: ModelStateMapper

Mapper that runs a child mapper under key prefixes.

Use it to apply a mapper written for a submodule (e.g. one operating on "in_proj") to the state dict of a parent module. Input (source) and output (target) prefixes are independent.

__init__(mapper, source_prefix='', target_prefix='')

Constructs the ModelStateMapperPrefixScope object.

Parameters:

Name Type Description Default
mapper ModelStateMapper

The child mapper to run within the scope.

 required
source_prefix str

The prefix added to the input keys of mapper.

 ''
target_prefix str

The prefix added to the output keys of mapper.

 ''

ModelStateMapperSequential

Bases: ModelStateMapper

Mapper that executes a list of mappers one after another.

Intermediate states stay hidden: the mapper exposes only the inputs of the first stage and the outputs of the last stage.

  1. Gap filling: If a tensor must pass a stage unchanged to reach a later stage or the final output, an identity mapper is added for it.
  2. Group merging: The mapper reports the net dependency graph. If stage A maps x to y and stage B maps y to z, the mapper reports a single group {x} -> {z}.
__init__(mappers)

Constructs the ModelStateMapperSequential object.

Parameters:

Name Type Description Default
mappers list[ModelStateMapper]

The stages in execution order. Mappers without inputs and outputs are dropped.

 required

Raises:

Type Description
ValueError

If no mapper with inputs or outputs is given.

ModelStateMapperShard

Bases: ModelStateMapper

Mapper that restricts another mapper to one shard of its dependency groups.

Use it to split model loading across processes. Give each process a different current_shard so that each one loads only its part of the checkpoint.

__init__(sub_mapper, total_shards, current_shard)

Constructs the ModelStateMapperShard object.

Parameters:

Name Type Description Default
sub_mapper ModelStateMapper

The mapper whose dependency groups are split.

 required
total_shards int

The total number of shards.

 required
current_shard int

The index of the shard this mapper handles, in [0, total_shards).

 required

filter_empty_mappers(mappers)

Filters out mappers that have no effect (no inputs and no outputs).

Parameters:

Name Type Description Default
mappers Sequence[ModelStateMapper]

The list of mappers to filter.

 required

Returns:

Type Description
list[ModelStateMapper]

A new list containing only active mappers.

d9d.model_state.mapper.leaf

Leaf mappers that transform individual tensors.

ModelStateMapperChunkTensors

Bases: ModelStateMapper

Mapper that chunks a single input tensor into multiple output tensors along a specified dimension.

__init__(source_name, target_names, dim)

Constructs the ModelStateMapperChunkTensors object.

Parameters:

Name Type Description Default
source_name str

Name of the input tensor to read.

 required
target_names list[str]

Names of the resulting chunked output tensors.

 required
dim int

Dimension along which to chunk the tensor.

 required

ModelStateMapperConcatenateTensors

Bases: ModelStateMapper

Mapper that concatenates multiple input tensors into a single output tensor along a specified dimension.

It is the inverse of ModelStateMapperChunkTensors.

__init__(source_names, target_name, dim)

Constructs the ModelStateMapperConcatenateTensors object.

Parameters:

Name Type Description Default
source_names list[str]

Names of the input tensors to read.

 required
target_name str

Name of the resulting concatenated output tensor.

 required
dim int

Dimension along which to concatenate the tensors.

 required

ModelStateMapperDistribute

Bases: ModelStateMapper

Mapper that converts a single full local tensor into a DTensor.

Every rank must hold the same full tensor. No communication happens.

__init__(name, device_mesh, placements)

Constructs the ModelStateMapperDistribute object.

Parameters:

Name Type Description Default
name str

The name of the tensor to distribute.

 required
device_mesh DeviceMesh | None

The device mesh of the resulting DTensor. Passed to distribute_tensor.

 required
placements Sequence[Placement] | None

The placements of the resulting DTensor. Passed to distribute_tensor.

 required

ModelStateMapperGatherFullTensor

Bases: ModelStateMapper

Mapper that gathers a single DTensor into a full local tensor.

__init__(name)

Constructs the ModelStateMapperGatherFullTensor object.

Parameters:

Name Type Description Default
name str

The name of the tensor to gather.

 required

ModelStateMapperIdentity

Bases: ModelStateMapper

Mapper that passes a single state tensor through unchanged.

__init__(name)

Constructs the ModelStateMapperIdentity object.

Parameters:

Name Type Description Default
name str

The name of the tensor to pass through.

 required

ModelStateMapperRename

Bases: ModelStateMapper

Mapper that renames a single state tensor.

__init__(name_from, name_to)

Constructs the ModelStateMapperRename object.

Parameters:

Name Type Description Default
name_from str

The source key.

 required
name_to str

The target key.

 required

ModelStateMapperSelectChildModules

Bases: ModelStateMapper

Mapper that selects the keys of a child module and strips the module prefix.

It is a batch rename that moves parameters from a submodule scope to the current scope.

__init__(base_names, parent_name)

Constructs the ModelStateMapperSelectChildModules object.

Parameters:

Name Type Description Default
base_names list[str]

The keys relative to the child module, e.g. "weight".

 required
parent_name str

The name of the child module, e.g. "in_proj". Input keys are "{parent_name}.{base_name}".

 required

ModelStateMapperSqueeze

Bases: ModelStateMapper

Mapper that squeezes an input tensor along a specified dimension or all dimensions of size 1.

__init__(name, dim=None)

Constructs the ModelStateMapperSqueeze object.

Parameters:

Name Type Description Default
name str

Name of the tensor to operate on.

 required
dim int | None

Dimension to squeeze. If None, squeezes all dimensions of size 1.

 None

ModelStateMapperStackTensors

Bases: ModelStateMapper

Mapper that stacks multiple input tensors into a single output tensor along a new dimension.

__init__(source_names, target_name, dim)

Constructs the ModelStateMapperStackTensors object.

Parameters:

Name Type Description Default
source_names list[str]

Names of the input tensors to read.

 required
target_name str

Name of the resulting stacked output tensor.

 required
dim int

Dimension along which to stack the tensors.

 required

ModelStateMapperTranspose

Bases: ModelStateMapper

Mapper that transposes an input tensor along two specified dimensions.

__init__(name, dims)

Constructs the ModelStateMapperTranspose object.

Parameters:

Name Type Description Default
name str

Name of the tensor to operate on.

 required
dims tuple[int, int]

Dimensions to transpose.

 required

ModelStateMapperUnsqueeze

Bases: ModelStateMapper

Mapper that unsqueezes an input tensor along a specified dimension.

__init__(name, dim)

Constructs the ModelStateMapperUnsqueeze object.

Parameters:

Name Type Description Default
name str

Name of the tensor to operate on.

 required
dim int

Dimension to unsqueeze.

 required

ModelStateMapperUnstackTensors

Bases: ModelStateMapper

Mapper that unstacks a single input tensor into multiple output tensors along a specified dimension.

__init__(source_name, target_names, dim)

Constructs the ModelStateMapperUnstackTensors object.

Parameters:

Name Type Description Default
source_name str

Name of the input tensor to read.

 required
target_names list[str]

Names of the resulting unstacked output tensors.

 required
dim int

Dimension along which to unstack the tensor.

 required