core.dist_checkpointing.exchange_utils

Utilities for exchanging data between ranks.

Module Contents

Classes

ShardDistribution

Represents a distribution of ShardedTensors.

Functions

is_float8tensor	Check if a tensor is a Transformer Engine Float8Tensor
_shard_size	Returns size in bytes of a given sharded tensor.
_get_empty_tensor_for_exchange	Determines the empty tensor to use for exchange.
distribute_shards_to_ranks	Computes uniform distribution of workload across ranks, based on sizes.
_gather_shards_metadata	Runs the single expensive collective behind the load/save distribution.
_build_shard_distribution	Computes a ShardDistribution from already-gathered shard metadata.
_pg_dist_cache_file_path	Returns the cache file path holding the distributions for a given group.
_load_pg_dist_cache	Reads both pre-computed ShardDistributions for a group from disk.
_create_pg_dist_cache	Builds both distribution variants and persists them for the group.
determine_main_replica_uniform_distribution	Computes (or loads from cache) the load/save shard distribution.
exchange_loaded_tensors_gather_rounds	Exchange the tensors loaded by different ranks with several all_gather calls.
exchange_loaded_tensors_gather_object	Exchange the tensors loaded by different ranks with a simple all_gather_object call.
exchange_loaded_objects_gather_object	Exchange the objects loaded by different ranks with a simple all_gather_object call.
exchange_loaded_tensors_broadcast	Exchange the tensors loaded by different ranks by a series of broadcasts.
exchange_by_distribution	Exchange tensors loaded by different ranks using the specified exchange_algo.

Data

logger
T
PG_DIST_CACHE_FILE_PREFIX
_PG_DIST_CACHE

API

core.dist_checkpointing.exchange_utils.is_float8tensor(tensor: torch.Tensor) → bool

Check if a tensor is a Transformer Engine Float8Tensor

core.dist_checkpointing.exchange_utils.logger

‘getLogger(…)’

class core.dist_checkpointing.exchange_utils.ShardDistribution

Bases: typing.NamedTuple

Represents a distribution of ShardedTensors.

Given distribution is valid only for a specific parallelization group, which is implicit here (not referenced by this class).

Parameters:

• main_rank_for_shard (Dict[_ShardId, int]) – specifies which rank should hold the main replica for a given shard

• shards_in_this_group (Set[_ShardId]) – which shards have a main replica in this parallelization group

• shard_to_metadata (Dict[_ShardId, ShardedTensor]) – maps ShardedTensor identifier to the original ShardedTensor

• all_ranks_for_shard (Dict[_ShardId, List[int]]) – specifies which ranks need a given shard in a given parallelization group

main_rank_for_shard: Dict[core.dist_checkpointing.utils._ShardId, int]

None

shards_in_this_group: Set[core.dist_checkpointing.utils._ShardId]

None

shard_to_metadata: Dict[core.dist_checkpointing.utils._ShardId, core.dist_checkpointing.mapping.ShardedTensor]

None

all_ranks_for_shard: Dict[core.dist_checkpointing.utils._ShardId, List[int]]

None

core.dist_checkpointing.exchange_utils._shard_size(sh_ten: core.dist_checkpointing.mapping.ShardedTensor)

Returns size in bytes of a given sharded tensor.

core.dist_checkpointing.exchange_utils._get_empty_tensor_for_exchange(

shard_id: core.dist_checkpointing.utils._ShardId,

needed_shards: Dict[core.dist_checkpointing.utils._ShardId, core.dist_checkpointing.mapping.ShardedTensor],

unneeded_shards: Dict[core.dist_checkpointing.utils._ShardId, core.dist_checkpointing.mapping.ShardedTensor],

loaded_tensors: Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor],

) → Tuple[torch.Tensor, Optional[torch.device]]

Determines the empty tensor to use for exchange.

If shard_id is needed by this rank, it will be in the unloaded_shards. Otherwise, the metadata for this tensor can be found in shard_to_metadata

Parameters:

• shard_id (_ShardId) – shard_id that will be exchanged

• needed_shards (Dict[_ShardId, ShardedTensor]) – mapping from shard ids to metadata for shards needed by this rank

• unneeded_shards (Dict[_ShardId, ShardedTensor]) – mapping from shard ids to metadata for shards that can be discarded after exchange

• loaded_tensors (Dict[_ShardId, torch.Tensor]) – mapping where useful tensors are placed in

Returns:

empty CUDA tensor to be exchanged, and the device of the original state dict tensor (if there was any)

Return type:

Tuple[torch.Tensor, Optional[torch.device]]

core.dist_checkpointing.exchange_utils.T

‘TypeVar(…)’

core.dist_checkpointing.exchange_utils.distribute_shards_to_ranks(

shard_to_ranks: Dict[core.dist_checkpointing.exchange_utils.T, List[int]],

shard_to_size: Dict[core.dist_checkpointing.exchange_utils.T, int],

num_ranks: int,

cross_parallelization_group_loads: Set[core.dist_checkpointing.exchange_utils.T],

) → Dict[core.dist_checkpointing.exchange_utils.T, int]

Computes uniform distribution of workload across ranks, based on sizes.

Currently, the assignment is greedy, based on:

1. Cross-parallelization group dependencies (shards with main rank in another group are assigned at the end to make sure the distribution for load and save is as similar as possible).

2. Secondly, the coverage of each shard (how many ranks the shard is available on; lower coverage is assigned first)

3. Then, the size of each shard (larger size is assigned first)

4. Finally, shard id for differentiation.

Last step is added because we rely on the fact that the assignment is deterministic on all ranks.

Parameters:

• shard_to_ranks (Dict[T, List[int]]) – mapping of rank access to shards

• shard_to_size (Dict[T, int]) – sizes of each shard

• num_ranks (int) – number of ranks in the parallelization group

• cross_parallelization_group_loads (Set[T]) – Shards to load that are not in the main replica

Returns (Dict[T, int]): assignment of shard to rank (which rank should do the work to achieve maximal uniformity)

core.dist_checkpointing.exchange_utils.PG_DIST_CACHE_FILE_PREFIX

‘pg_dist’

core.dist_checkpointing.exchange_utils._PG_DIST_CACHE: Dict[str, Dict[bool, core.dist_checkpointing.exchange_utils.ShardDistribution]]

None

core.dist_checkpointing.exchange_utils._gather_shards_metadata(

sharded_state_dict: core.dist_checkpointing.mapping.ShardedStateDict,

parallelization_group: torch.distributed.ProcessGroup,

) → List[List[core.dist_checkpointing.mapping.ShardedTensor]]

Runs the single expensive collective behind the load/save distribution.

Every rank in the group contributes the metadata (no tensor data) of the ShardedTensors it holds; the result is the list of every rank’s shard metadata, identical on all ranks of the group. This is the all_gather_object that the caching feature exists to eliminate on subsequent jobs.

Decision: data is stripped with without_data() before the gather. Only the metadata (key, offsets, shape, dtype, replica_id) drives the distribution algorithm, and gathering it keeps the pickled payload tiny.

Parameters:

• sharded_state_dict (ShardedStateDict) – state dict whose local shards are advertised to the group.

• parallelization_group (ProcessGroup) – group to gather within.

Returns:

per-rank lists of data-less ShardedTensors.

Return type:

List[List[ShardedTensor]]

core.dist_checkpointing.exchange_utils._build_shard_distribution(

all_shards: List[List[core.dist_checkpointing.mapping.ShardedTensor]],

ignore_groups: bool,

) → core.dist_checkpointing.exchange_utils.ShardDistribution

Computes a ShardDistribution from already-gathered shard metadata.

This is the pure, deterministic post-gather logic of determine_main_replica_uniform_distribution factored out so that:

• it can be run twice (for both ignore_groups variants) from a single gather when creating the cache, and

• it has no dependence on the local rank, which is what makes the result identical on every rank of the group and therefore safe to persist and reload across jobs.

Parameters:

• all_shards (List[List[ShardedTensor]]) – gathered per-rank shard metadata (output of _gather_shards_metadata).

• ignore_groups (bool) – if True (load mode) include shards whose main replica lives outside this group; if False (save mode) restrict to shards with a main replica in this group.

Returns:

the computed distribution.

Return type:

ShardDistribution

core.dist_checkpointing.exchange_utils._pg_dist_cache_file_path(

cache_path: str,

parallelization_group: torch.distributed.ProcessGroup,

) → str

Returns the cache file path holding the distributions for a given group.

The file is keyed by the minimum global rank of the parallelization group. Decision rationale:

• Parallelization groups partition the world, so min(global ranks) is a unique, collision-free identifier of each group.

• It is computed from purely local metadata (get_process_group_ranks), so the filename is derivable without any collective — essential for the zero-communication read path.

• Every member of a group maps to the same file (one file per group rather than one per rank): far fewer files at scale and a single read per rank, with duplicate reads absorbed by the FS page cache.

Parameters:

• cache_path (str) – directory holding the cache files.

• parallelization_group (ProcessGroup) – the group to key on.

Returns:

absolute-or-relative path of the group’s cache file.

Return type:

str

core.dist_checkpointing.exchange_utils._load_pg_dist_cache(

cache_file: str,

) → Dict[bool, core.dist_checkpointing.exchange_utils.ShardDistribution]

Reads both pre-computed ShardDistributions for a group from disk.

Returns the full {ignore_groups: ShardDistribution} map (not just one variant) so the caller can memoize both in the process-global cache and serve later save/load calls without re-opening the file.

Performs no existence/validity checks by design (see the feature docs): the user opts into the cache only when they guarantee the config and world size match the run that created it, trading safety for the lowest possible latency.

Parameters:

cache_file (str) – resolved path of the group’s cache file.

Returns:

the {ignore_groups: distribution} map.

Return type:

Dict[bool, ShardDistribution]

core.dist_checkpointing.exchange_utils._create_pg_dist_cache(

all_shards: List[List[core.dist_checkpointing.mapping.ShardedTensor]],

cache_path: str,

parallelization_group: torch.distributed.ProcessGroup,

) → Dict[bool, core.dist_checkpointing.exchange_utils.ShardDistribution]

Builds both distribution variants and persists them for the group.

Both the save (ignore_groups=False) and load (ignore_groups=True) distributions are derived from the same gathered metadata, so a single collective suffices to populate a cache usable by every later load and save — even if the creating job only ever loads (e.g. exit_after_loading_ckpt).

Only the group’s writer rank (global rank == min of the group) writes, keeping create-time IO to one file per group. The write is atomic (temp file + os.replace) so a concurrent/interrupted create can never leave a torn file that a reader would later trust blindly.

Parameters:

• all_shards (List[List[ShardedTensor]]) – gathered shard metadata.

• cache_path (str) – directory to write the cache file into.

• parallelization_group (ProcessGroup) – group being cached.

Returns:

the {ignore_groups: distribution} map, so the caller can return the variant it needs without recomputing.

Return type:

Dict[bool, ShardDistribution]

core.dist_checkpointing.exchange_utils.determine_main_replica_uniform_distribution(

sharded_state_dict: core.dist_checkpointing.mapping.ShardedStateDict,

parallelization_group: torch.distributed.ProcessGroup,

ignore_groups: bool = False,

pg_cache_path: Optional[str] = None,

pg_cache_create: bool = False,

) → Optional[core.dist_checkpointing.exchange_utils.ShardDistribution]

Computes (or loads from cache) the load/save shard distribution.

Should be used in conjunction with distribute_main_replicas_with_precomputed_distribution which applies the computed save distribution.

We rely on the fact that the assignment algorithm is deterministic on all ranks, so there is no extra communication needed after metadata exchange.

PG-collective caching: The metadata exchange (all_gather_object) is by far the dominant cost of this function at scale. Its result is a pure deterministic function of the sharded state dict structure and the parallel layout, so for a fixed config and world size it is identical across jobs and can be cached on disk:

• pg_cache_path set, pg_cache_create=False (READ): skip the collective entirely and load the pre-computed distribution for this group and mode from disk.

• pg_cache_path set, pg_cache_create=True (CREATE): run the collective once, build both the save and load distributions from the single gather, persist them, and return the requested one.

• both unset (default): the original behaviour, unchanged.

Results are additionally memoized in a process-global cache (_PG_DIST_CACHE) keyed by the group’s cache file, so within a single job each group is created (gather + write) or read (file open) at most once — later calls (subsequent saves, the second replicate_local_replicas pass, or the save wrapper reusing the load wrapper’s entry) are served from memory.

No existence/validity checks are performed on the cache by design — the user opts in only when guaranteeing a matching config and world size.

Parameters:

• sharded_state_dict (ShardedStateDict) – state dict to compute the distribution of

• parallelization_group (ProcessGroup) – distribution will be computed within this process group

• ignore_groups (bool, optional) – whether the distribution defines groups. This option is primarily used during loading, as it ensures that all replicas, including non-main ones, are loaded by this parallelization group Defaults to False.

• pg_cache_path (str, optional) – directory of the PG-distribution cache. When set, enables the caching feature. Defaults to None (disabled).

• pg_cache_create (bool, optional) – when True (and pg_cache_path set), compute the distribution normally and write the cache; when False (and pg_cache_path set), read the cache and skip the collective. Defaults to False.

Returns (ShardDistribution, optional): distribution that can be used to apply the parallelization. Returns None if the process_group is trivial (1 rank)

core.dist_checkpointing.exchange_utils.exchange_loaded_tensors_gather_rounds(

loaded_tensors: Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor],

unloaded_shards: Dict[core.dist_checkpointing.utils._ShardId, core.dist_checkpointing.mapping.ShardedTensor],

shard_distribution: core.dist_checkpointing.exchange_utils.ShardDistribution = None,

parallelization_group: Optional[torch.distributed.ProcessGroup] = None,

) → Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor]

Exchange the tensors loaded by different ranks with several all_gather calls.

Groups tensors by dtype, divide tensors that will be exchanged into rounds and execute all_gather for tensors from each round.

Note: the loading is distributed across ranks based on total loaded size in bytes, so there is no guarantee that number of rounds needed for each rank will be similar, which might result in a lot of almost empty all_gathers. The solution would be to group all tensors into a one bytes tensor and do a single all_gather (with similarly sized messages).

Parameters:

• loaded_tensors (Dict[_ShardId, torch.Tensor]) – mapping from ShardedTensor shard ids to tensors already loaded by this rank.

• unloaded_shards (Dict[_ShardId, torch.Tensor]) – mapping from ShardedTensor shard ids to ShardedTensors that aren’t loaded yet.

• shard_distribution (ShardDistribution) – distribution of all shards

• parallelization_group (ProcessGroup, optional) – process group used for load distribution. Tensors will be exchanged within this group

Returns:

dictionary mapping shard ids to tensors needed by this rank to load a given state dict. Includes previously loaded tensors (from loaded_tensors input)

Return type:

Dict[_ShardId, torch.Tensor]

core.dist_checkpointing.exchange_utils.exchange_loaded_tensors_gather_object(

loaded_tensors: Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor],

unloaded_shards: Dict[core.dist_checkpointing.utils._ShardId, core.dist_checkpointing.mapping.ShardedTensor],

shard_distribution: core.dist_checkpointing.exchange_utils.ShardDistribution,

parallelization_group: Optional[torch.distributed.ProcessGroup] = None,

) → Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor]

Exchange the tensors loaded by different ranks with a simple all_gather_object call.

This version can be used for debugging purposes do to its simplistic implementation. Shouldn’t be used if performance is important.

Parameters:

• loaded_tensors (Dict[_ShardId, torch.Tensor]) – mapping from ShardedTensor shard ids to tensors already loaded by this rank.

• unloaded_shards (Dict[_ShardId, torch.Tensor]) – mapping from ShardedTensor shard ids to ShardedTensors that aren’t loaded yet.

• shard_distribution (ShardDistribution) – distribution of all shards

• parallelization_group (ProcessGroup, optional) – process group used for load distribution. Tensors will be exchanged within this group

Returns:

dictionary mapping shard ids to tensors needed by this rank to load a given state dict. Includes previously loaded tensors (from loaded_tensors input)

Return type:

Dict[_ShardId, torch.Tensor]

core.dist_checkpointing.exchange_utils.exchange_loaded_objects_gather_object(

loaded_objects: Dict[core.dist_checkpointing.utils._ShardId, Any],

) → Dict[core.dist_checkpointing.utils._ShardId, Any]

Exchange the objects loaded by different ranks with a simple all_gather_object call.

Parameters:

loaded_objects (Dict[_ShardId, Any]) – mapping from shard ids to objects already loaded by this rank.

Returns:

dictionary mapping shard ids to objects needed by this rank to load a given state dict.

Return type:

Dict[_ShardId, Any]

core.dist_checkpointing.exchange_utils.exchange_loaded_tensors_broadcast(

loaded_tensors: Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor],

unloaded_shards: Dict[core.dist_checkpointing.utils._ShardId, core.dist_checkpointing.mapping.ShardedTensor],

shard_distribution: core.dist_checkpointing.exchange_utils.ShardDistribution,

parallelization_group: Optional[torch.distributed.ProcessGroup] = None,

) → Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor]

Exchange the tensors loaded by different ranks by a series of broadcasts.

For each rank for each loaded tensor do a broadcast to the whole group. A reasonable tradeoff in terms of performance and simplicity.

Parameters:

• loaded_tensors (Dict[_ShardId, torch.Tensor]) – mapping from ShardedTensor shard ids to tensors already loaded by this rank.

• unloaded_shards (Dict[_ShardId, ShardedTensor]) – mapping from ShardedTensor shard ids to ShardedTensors that aren’t loaded yet.

• shard_distribution (ShardDistribution) – distribution of all shards

• parallelization_group (ProcessGroup, optional) – process group used for load distribution. Tensors will be exchanged within this group

Returns:

dictionary mapping shard ids to tensors needed by this rank to load a given state dict. Includes previously loaded tensors (from loaded_tensors input)

Return type:

Dict[_ShardId, torch.Tensor]

core.dist_checkpointing.exchange_utils.exchange_by_distribution(

loaded_tensors: Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor],

unloaded_shards: Dict[core.dist_checkpointing.utils._ShardId, core.dist_checkpointing.mapping.ShardedTensor],

shard_distribution: core.dist_checkpointing.exchange_utils.ShardDistribution,

parallelization_group: Optional[torch.distributed.ProcessGroup] = None,

exchange_algo='broadcast',

) → Dict[core.dist_checkpointing.utils._ShardId, torch.Tensor]

Exchange tensors loaded by different ranks using the specified exchange_algo.

Parameters:

• loaded_tensors (Dict[_ShardId, torch.Tensor]) – mapping from ShardedTensor shard ids to tensors already loaded by this rank.

• unloaded_shards (Dict[_ShardId, ShardedTensor]) – mapping from ShardedTensor shard ids to ShardedTensors that aren’t loaded yet.

• shard_distribution (ShardDistribution) – distribution of all shards

• parallelization_group (ProcessGroup, optional) – process group used for load distribution. Tensors will be exchanged within this group

• exchange_algo (str) – The algorithm used for performing exchanges. Defaults to ‘broadcast’.

Returns:

dictionary mapping shard ids to tensors needed by this rank to load a given state dict. Includes previously loaded tensors (from loaded_tensors input)

Return type:

Dict[_ShardId, torch.Tensor]