Model State I/O

from pathlib import Path
from d9d.model_state.io import read_model_state
from d9d.model_state.mapper.adapters import identity_mapper_from_module

model = ...

# Load every state of the model unchanged.
mapper = identity_mapper_from_module(model)

# src_dir must contain either sharded .safetensors files with model.safetensors.index.json,
# or a single unindexed model.safetensors.
loader_stream = read_model_state(
    src_dir=Path("./checkpoint"),
    mapper=mapper,
    device="cpu"  # Or "cuda:0"
)

state_dict = {}
for name, tensor in loader_stream:
    print(f"Loaded and transformed: {name} -> {tensor.shape}")
    state_dict[name] = tensor

from pathlib import Path
from d9d.model_state.io import write_model_state_local
from d9d.model_state.mapper.adapters import identity_mapper_from_module

model = ...

# The writer creates the shard files and the index file.
write_model_state_local(
    dest_dir=Path("./output_checkpoint"),
    mapper=identity_mapper_from_module(model),
    state_generator=model.state_dict().items(),
    shard_size_gb=1.0
)

import torch.distributed as dist
from pathlib import Path
from d9d.model_state.io import read_model_state, write_model_state_distributed
from d9d.model_state.mapper.compose import ModelStateMapperShard
from d9d.model_state.mapper.adapters import identity_mapper_from_mapper_outputs

dist.init_process_group("nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()

# Describe how the ENTIRE model is converted,
# e.g. "stack q, k, v", "rename MLP", "load everything else as is".
mapper = build_my_custom_mapper()

# Each rank processes only its share of the dependency groups,
# so no two ranks load, process or save the same tensors.
local_work_mapper = ModelStateMapperShard(
    sub_mapper=mapper,
    total_shards=world_size,
    current_shard=rank
)

# read_model_state yields tensors that already have their target names and shapes,
# so the writer only passes them through.
writer_mapper = identity_mapper_from_mapper_outputs(local_work_mapper)

# The reader loads and transforms the source tensors, and the writer saves them.
# Rank 0 then writes model.safetensors.index.json for all ranks and gives the files their final names.
write_model_state_distributed(
    dest_dir=Path("./converted_checkpoint"),
    mapper=writer_mapper,
    state_generator=read_model_state(
        src_dir=Path("./original_checkpoint"),
        mapper=local_work_mapper,
        device="cuda",
        show_progress=False  # Hide reader progress bars to keep the output readable
    ),
    process_group=dist.group.WORLD,
    shard_size_gb=4.0,
    show_progress=True
)

from pathlib import Path
from d9d.model_state.io import load_model_state
from d9d.model_state.mapper.adapters import identity_mapper_from_module

model = ...

# Load every key of the model as is.
mapper = identity_mapper_from_module(model)

load_model_state(
    src_dir=Path("./checkpoints/v1"),
    mapper=mapper,
    device="cuda",
    model=model
)

from pathlib import Path

from torch.distributed.device_mesh import init_device_mesh
from d9d.model_state.io import save_model_state_pipeline_parallel
from d9d.model_state.mapper.compose import ModelStateMapperParallel
from d9d.model_state.mapper.adapters import identity_mapper_from_module

# pp=2 (pipeline), dp=2 (data), tp=2 (tensor).
mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("pp", "dp", "tp"))

# Each pipeline rank holds two stages of the model.
my_stages = [TransformerStage(...), TransformerStage(...)]

# Combine the states of all local stages into one mapper.
mapper = ModelStateMapperParallel([
    identity_mapper_from_module(stage) for stage in my_stages
])

save_model_state_pipeline_parallel(
    dest_dir=Path("./checkpoint"),
    mapper=mapper,
    device_mesh=mesh,
    pipeline_dim_name="pp",
    models=my_stages,
    shard_size_gb=4.0
)