distributed#

Utility functions for using torch.distributed.

Classes

`DistributedProcessGroup`	A convenient wrapper around torch.distributed.ProcessGroup objects.
`ParallelState`	A class to manage various parallel groups such as data parallel, tensor parallel etc.

Functions

`backend`	Returns the distributed backend.
`barrier`	Synchronizes all processes.
`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_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.
`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 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)

backend()#

Returns the distributed backend.

Return type:: str | None

barrier(group=None)#

Synchronizes all processes.

Return type:: None

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_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

size(group=None)#

Returns the number of processes.

Return type:: int