Is this page helpful?

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor#

Distributed-checkpoint chunk metadata for Megatron-FSDP’s packed parameter buffers.

An :class:~.parameter_group.FsdpParameterGroup packs several parameters into one flat

class:

~.dbuffer.DBuffer with least-common-multiple row padding, so a parameter’s per-rank shard does not tile the way torch’s canonical Shard(0) does: a rank may own several rows of one parameter and none of the next. :meth:~.dbuffer.DBuffer.get_dtensor presents each parameter as a plain Shard(0) DTensor, which cannot express those offsets, so DCP’s default planner would mis-place every packed shard and silently corrupt the checkpoint.

DCP lets a local tensor override that placement through the __create_chunk_list__ and __create_write_items__ hooks, which :class:~torch.distributed.tensor.DTensor forwards to its local tensor. This module derives each shard’s true position from the

class:

~.layout.GlobalLayout that already backs the buffer – it records every logical tensor’s global element offset, and the buffer records this rank’s owned element range – and attaches those hooks. The derivation is pure rank-local arithmetic, so it needs no collective.

Module Contents#

Functions#

_dbuffer_chunk_metadata

Return the chunk of logical tensor index this rank holds in buffer.

chunk_metadata_by_fqn

Return this rank’s chunk of every Megatron-FSDP parameter in model, keyed by FQN.

attach_uneven_dtensor_metadata

Attach Megatron-FSDP chunk metadata to the sharded DTensors DCP will save or load.

_build_chunk_metadata

Build chunk metadata for a dim-0 row range of a tensor of shape shape.

_attach_by_fqn

Attach chunk metadata to the DTensors of a parameter-FQN-keyed state dict.

_attach_dcp_hooks

Override a DTensor’s DCP placement with chunk.

Data#

__all__

_OPTIMIZER_STATE_KEY

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor.__all__#

[‘attach_uneven_dtensor_metadata’, ‘chunk_metadata_by_fqn’]

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor._OPTIMIZER_STATE_KEY#

‘state’

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor._dbuffer_chunk_metadata(
buffer: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer,
index: int,
) → torch.distributed.checkpoint.metadata.ChunkStorageMetadata#

Return the chunk of logical tensor index this rank holds in buffer.

Flat placements shard the flat buffer, and every logical tensor is row-aligned within it, so a rank always owns whole dim-0 rows and the chunk is a dim-0 range of the global tensor.

Parameters:

  • buffer – Buffer holding the logical tensor.

  • index – Index of the logical tensor within buffer.

Returns:

This rank’s chunk offset and size, in global tensor coordinates.

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor.chunk_metadata_by_fqn(
model: torch.nn.Module,
) → dict[str, torch.distributed.checkpoint.metadata.ChunkStorageMetadata]#

Return this rank’s chunk of every Megatron-FSDP parameter in model, keyed by FQN.

The keys are the parameter FQNs that torch’s DCP state-dict helpers use for both the model state dict and the per-parameter entries of the optimizer state dict.

Parameters:

model – A module tree that has been sharded with :func:~.fully_shard.fully_shard.

Returns:

Chunk metadata for each sharded parameter, and nothing for parameters FSDP does not own. Tied parameters appear once per FQN, all sharing the one chunk of the buffer entry that backs them.

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor.attach_uneven_dtensor_metadata(
model: torch.nn.Module,
model_state_dict: dict[str, Any],
optimizer_state_dict: dict[str, Any],
) → None#

Attach Megatron-FSDP chunk metadata to the sharded DTensors DCP will save or load.

Optimizer state is allocated as zeros_like of a sharded parameter, so it shares that parameter’s chunk.

Parameters:

  • model – The sharded module tree the state dicts were taken from.

  • model_state_dict –

    Model state dict from

    func:

    ~torch.distributed.checkpoint.state_dict.get_model_state_dict.

  • optimizer_state_dict –

    Optimizer state dict from

    func:

    ~torch.distributed.checkpoint.state_dict.get_optimizer_state_dict.

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor._build_chunk_metadata(
row_offset: int,
rows: int,
shape: torch.Size,
) → torch.distributed.checkpoint.metadata.ChunkStorageMetadata#

Build chunk metadata for a dim-0 row range of a tensor of shape shape.

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor._attach_by_fqn(
state_dict: dict[str, Any],
metadata_by_fqn: dict[str, torch.distributed.checkpoint.metadata.ChunkStorageMetadata],
) → None#

Attach chunk metadata to the DTensors of a parameter-FQN-keyed state dict.

Each value is either a parameter’s tensor, as in a model state dict, or a mapping of that parameter’s state tensors, as in the per-parameter entries of an optimizer state dict.

core.distributed.fsdp.src.megatron_fsdp.experimental.uneven_dtensor._attach_dcp_hooks(
dtensor: torch.distributed.tensor.DTensor,
chunk: torch.distributed.checkpoint.metadata.ChunkStorageMetadata,
) → None#

Override a DTensor’s DCP placement with chunk.