distributed#

Utility functions for using torch.distributed.

Classes

`DistributedProcessGroup`	A convenient wrapper around torch.distributed.ProcessGroup objects.
`DistributedState`	This process's place in a run: its rank, the run's size, and the device it computes on.
`DistributionPlan`	How a model is distributed over the ranks of a run.
`ParallelState`	A class to manage various parallel groups such as data parallel, tensor parallel etc.

Functions

`all_reduce`	Return `combine` of every rank's `fn()`, in rank order, on every rank.
`backend`	Returns the distributed backend.
`barrier`	Synchronizes all processes.
`broadcast_from`	Return `fn()`, run on `src` only, on every rank.
`distributed_run`	Own the run's process group for the block, and yield this process's state.
`distribution_plan`	How `model` is distributed: the plan recorded on it (`set_distribution_plan()`).
`fsdp2_wrap`	Auto-detect a HF causal-LM's decoder layers and FSDP2 `fully_shard` each one.
`gather_bytes`	Gathers byte strings to global rank `dst` only, without unpickling anything on the way.
`gather_state_dicts`	Gathers `state` from every rank of `group` to global rank `dst`, loaded with `weights_only`.
`is_available`	Returns whether the distributed package is available.
`is_distributed_launch`	True if launched distributed (WORLD_SIZE>1), readable BEFORE the process group is set up.
`is_distributed_run`	True if a real multi-rank job is running: the process group is initialized and world size > 1.
`is_fsdp2_model`	Return True if any submodule of `model` has been wrapped with FSDP2 `fully_shard`.
`is_initialized`	Returns whether the distributed package is initialized.
`is_master`	Returns whether the current process is the master process.
`rank`	Returns the rank of the current process.
`set_distribution_plan`	Record that `model` is distributed as `plan`, on it and on every submodule it has now.
`size`	Returns the number of processes.

class DistributedProcessGroup#

Bases: object

A convenient wrapper around torch.distributed.ProcessGroup objects.

__init__(group=None)#

Initialize the distributed process group.

Parameters:: group (ProcessGroup | int | None)

static get_dist_syncd_obj(obj, groups, op)#

Get the distributed synchronized object across the specified distributed groups.

Parameters:

obj (Any)
groups (DistributedProcessGroup | list[DistributedProcessGroup])
op (Callable)

is_initialized()#

Check if the distributed process group is initialized.

Return type:: bool

rank()#

Get the rank of the current process group.

Return type:: int

world_size()#

Get the world size of the current process group.

Return type:: int

class DistributedState#

Bases: object

This process’s place in a run: its rank, the run’s size, and the device it computes on.

distributed_run() yields it, for a single process too (rank 0 of a run of one).

__init__(rank, world_size, device)#

Parameters:

rank (int)
world_size (int)
device (device)

Return type:

None

device: device#

property is_main: bool#: Whether this is rank 0, the rank that writes the run’s outputs.

rank: int#

world_size: int#

class DistributionPlan#

Bases: object

How a model is distributed over the ranks of a run.

world_size – the ranks that hold the model; 1 for a single process.
ep_size – the expert-parallel degree: an MoE’s routed experts are sharded over groups of ep_size ranks. 1 shards none.
dense – how the rest of the model, all of it without expert parallelism, is placed: "replicate" gives every rank a whole copy, "fsdp" FSDP2-shards it.

Every rank calibrates on its own data, so data parallelism spans the whole world, and each expert shard is replicated edp = world_size / ep_size times.

__init__(world_size=1, ep_size=1, dense='replicate')#

Parameters:

world_size (int)
ep_size (int)
dense (Literal['replicate', 'fsdp'])

Return type:

None

dense: Literal['replicate', 'fsdp'] = 'replicate'#

property edp: int#

how many ranks hold a replica of each expert shard.

Type:: Expert-data-parallel degree

ep_size: int = 1#

property is_expert_parallel: bool#: Whether routed experts are sharded across ranks.

property is_fsdp_dense: bool#: Whether the dense part is FSDP2-sharded rather than replicated.

world_size: int = 1#

class ParallelState#

Bases: object

A class to manage various parallel groups such as data parallel, tensor parallel etc.

Specify the parallel groups of type torch.distributed.ProcessGroup for the current module. If the parallel group is not used, it should be set to -1. if a parallel group is None, it will use the default PyTorch distributed process group which is the whole world.

__init__(data_parallel_group=None, tensor_parallel_group=-1, expert_model_parallel_group=-1)#

Initialize the parallel state.

Parameters:

data_parallel_group (ProcessGroup | int | None)
tensor_parallel_group (ProcessGroup | int | None)
expert_model_parallel_group (ProcessGroup | int | None)

all_reduce(label, fn, combine)#

Return combine of every rank’s fn(), in rank order, on every rank.

E.g. combine=min for a batch size each rank’s memory bounds. If fn fails on any rank, or its value does not round-trip (_encoded()), every rank raises instead of waiting in the gather: if any rank’s failure is not a SystemExit, a rank with such an error raises it and the others a RuntimeError naming the first; else each rank exits, with its own code or the first one’s. Collective, with the same label on every rank; ranks out of step raise. combine gets copies loaded from the bytes the ranks sent, with any tensor in them on the CPU; in a single process it gets fn()’s value as it is. The values must load with weights_only, so a class in them must be registered with torch.serialization.add_safe_globals. A rank of a launch whose process group is not set up raises.

Unlike allreduce(), it runs fn itself, so a rank whose fn fails still joins.

Parameters:

label (str)
fn (Callable[[], _V])
combine (Callable[[list[_V]], _V])

Return type:

backend()#

Returns the distributed backend.

Return type:: str | None

barrier(group=None)#

Synchronizes all processes.

Return type:: None

broadcast_from(label, fn, src=0)#

Return fn(), run on src only, on every rank.

For a value ranks computing it for themselves could disagree on (their environment, GPU or files). If fn raises or its value does not round-trip (_encoded()), every rank raises rather than wait for the broadcast: src its own error, the others a SystemExit with its code if it exited, else a RuntimeError. Every rank, src too, gets a copy loaded from the bytes sent, with any tensor in it on the CPU; the value must load with weights_only, so a class in it must be registered with torch.serialization.add_safe_globals. Collective: every rank calls it at the same point with the same label and src; a receiving rank whose label differs raises, but src itself cannot tell and returns its value. In a single process fn’s value comes back as it is; a rank of a launch whose process group is not set up raises.

Parameters:

label (str)
fn (Callable[[], _V])
src (int)

Return type:

distributed_run(timeout=None, device='cuda')#

Own the run’s process group for the block, and yield this process’s state.

On a distributed launch (WORLD_SIZE>1) this binds this rank’s CUDA device, and sets the process group up with timeout unless the caller set one up already. It yields this rank’s DistributedState, on that CUDA device, and refuses any other device. Without a launch it sets nothing up, and the state is on device. The state’s rank and size are always the live group’s, which the block’s collectives run on: rank 0 of a run of one for a single process, and this rank’s place in a group the caller set up without the launch’s environment. A timeout cannot apply to a group of several ranks the caller set up, so passing one then warns. A single process leaves it unused, without a warning, also in a one-rank group of the caller’s: no rank there waits on a peer, so an entry point may pass one for its launches only.

On exit, even if the body raises, it destroys the group the run set up: the launch’s, or one the body set up, such as the one-rank group a single process’s FSDP2 load shards through. A group that was already there when the run began is the caller’s, and is left as it is. A rank whose body raised, in a run of several ranks on a group that was not there when it began, exits through abort() instead, and the launcher tears the job down: destroying the group would wait on peers that may sit in a collective this rank never reaches. In a run that found the caller’s group, the error goes on to the caller, which owns the peers.

A SystemExit is not an error: the group is destroyed as on a normal exit, which waits on every rank, so the exit has to be every rank’s, as one raised inside broadcast_from() or all_reduce() is. A rank that exits alone waits in that teardown until the group times out.

Parameters:

timeout (timedelta | None)
device (device | str)

Return type:

Iterator[DistributedState]

distribution_plan(model)#

How model is distributed: the plan recorded on it (set_distribution_plan()).

A model with no plan recorded is a single process’s, unless it shows otherwise: one that is FSDP2-sharded, or held in a run of several ranks, raises rather than be taken for one. That includes a module added to a model after its plan was recorded: a module does not know the model it was added to, so read the plan off that model. A run is its live process group: before a launch sets its group up, each of its processes runs alone with a model of its own, so a model with no plan reads as a single process’s there.

Parameters:: model (Module)
Return type:: DistributionPlan

fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload=False)#

Auto-detect a HF causal-LM’s decoder layers and FSDP2 fully_shard each one.

By default (shard_root=True) the root module is wrapped too, so embed/lm_head/norm are sharded instead of replicated per rank; pass shard_root=False to leave the root replicated (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the detection result.

Parameters whose dtype differs from the model’s dominant one are excluded from the wrap (see _off_dtype_params()), since FSDP2 rejects a shard group that mixes dtypes. They stay replicated and are moved onto the shards’ device, which fully_shard does not do for the params it ignores.

Parameters:: cpu_offload (bool)

gather_bytes(data, dst, group=None, chunk_bytes=268435456)#

Gathers byte strings to global rank dst only, without unpickling anything on the way.

Unlike torch.distributed.gather_object it never runs pickle, and it moves chunk_bytes per rank at a time to bound device memory. Returns the per-rank bytes on dst and None elsewhere.

Parameters:

data (bytes)
dst (int)
chunk_bytes (int)

Return type:

list[bytearray] | None

gather_state_dicts(state, dst, group=None)#

Gathers state from every rank of group to global rank dst, loaded with weights_only.

Yields one rank’s dict at a time on dst (to bound host memory) and nothing elsewhere. The collective runs eagerly, so every rank must call this before iterating.

Parameters:

state (dict)
dst (int)

Return type:

Iterator[dict]

is_available()#

Returns whether the distributed package is available.

Return type:: bool

is_distributed_launch()#

True if launched distributed (WORLD_SIZE>1), readable BEFORE the process group is set up.

Env-based, so parse-time code can detect the launch mode before distributed_run() sets the process group up. Once it has, is_distributed_run() reads the live group.

Return type:: bool

is_distributed_run()#

True if a real multi-rank job is running: the process group is initialized and world size > 1.

Runtime state (checks the live process group), unlike is_distributed_launch() which reads only the launch env. A single process or a 1-rank group returns False. Backend-agnostic – true for any backend that initialized torch.distributed (HF EP/DP loader, Megatron, plain DDP).

Return type:: bool

is_fsdp2_model(model)#

Return True if any submodule of model has been wrapped with FSDP2 fully_shard.

Return type:: bool

is_initialized()#

Returns whether the distributed package is initialized.

Return type:: bool

is_master(group=None)#

Returns whether the current process is the master process.

Return type:: bool

rank(group=None)#

Returns the rank of the current process.

Return type:: int

set_distribution_plan(model, plan)#

Record that model is distributed as plan, on it and on every submodule it has now.

Code handed a submodule, such as a VLM’s language model, then reads the same plan as code handed the whole model; a module added afterwards records none. Code that distributes a model records its plan here before calibrating or exporting the model. The plan says how the model is placed over the run; the ParallelState a quantized module holds names the process groups its collectives run on.

A model is distributed once, and its plan recorded once: recording a plan on a model that, or any of whose submodules, already records one raises, the same plan included, and leaves the model as it was. No code reloads or re-distributes a model it loaded.

In a run of several ranks, a plan for another number of ranks raises too, and records nothing: a single process’s plan there would be taken at its word, and the model replicated. A plan recorded before the run’s group is set up is not checked, nor one in a run of one, where a rank can stand in for several.

Parameters:

model (Module)
plan (DistributionPlan)

Return type:

None

size(group=None)#

Returns the number of processes.

Return type:: int