Distributed¶

Torch-native distributed helpers. AgileRL is single-device by default; multi-GPU LLM training initialises torch.distributed from the standard launcher environment variables set by torchrun (or an orchestration layer such as Ray), and these helpers no-op on a single device.

How to launch, when to use data parallel vs FSDP2, and every FSDPConfig field: Multi-GPU LLM training.

agilerl.distributed.init_distributed(timeout_seconds: int = 1800) → bool¶

Initialise torch.distributed from launcher env vars.

No-op (returns False) on a single device with no launcher env. Safe to call repeatedly; if a process group already exists it is reused.

Parameters:: timeout_seconds (int) – Collective timeout for the process group.
Returns:: True when distributed training is active.
Return type:: bool

agilerl.distributed.is_distributed() → bool¶: Whether torch.distributed is available and initialised.

agilerl.distributed.get_rank() → int¶: Global rank (launcher RANK, or 0 on a single device).

agilerl.distributed.get_local_rank() → int¶: Rank within the node, used for device selection.

agilerl.distributed.get_world_size() → int¶: Number of processes (launcher WORLD_SIZE, or 1 on a single device).

agilerl.distributed.is_main_process() → bool¶: Whether this is rank 0 (always True on a single device).

agilerl.distributed.barrier() → None¶: Synchronise all processes; no-op on a single device.

agilerl.distributed.broadcast_object_list(objects: list, src: int = 0) → list¶

Broadcast a list of picklable objects from src to all ranks.

Mutates objects in place on non-source ranks and returns it. No-op on a single device.

Parameters:

objects (list) – Objects to broadcast (same length on every rank).
src (int) – Source rank.

Returns:

The broadcast list.

Return type:

list

agilerl.distributed.gather_tensor(tensor: Tensor | ndarray | float) → Tensor¶

Gather a tensor from every rank (identity on a single device).

Parameters:: tensor – Tensor (or array/scalar convertible to one) to gather.
Returns:: Stacked / concatenated tensors from all ranks.

agilerl.distributed.gather_objects(objects: list) → list¶

Flatten object lists gathered from every rank.

Identity on a single device. Each rank contributes a picklable list; the result is concatenation in rank order.

Parameters:: objects (list) – Local objects to gather.
Returns:: Flattened objects from all ranks.
Return type:: list

agilerl.distributed.allreduce_minmax_int(value: int) → tuple[int, int]¶

Return (min, max) of value across ranks ((value, value) locally).

Parameters:: value (int) – This rank’s value.
Returns:: (min, max) across ranks.
Return type:: tuple[int, int]

agilerl.distributed.aggregate_metrics_across_gpus(metric_tensor: Tensor | ndarray | float) → float¶

Average a metric across ranks (local mean on a single device).

Parameters:: metric_tensor (torch.Tensor | np.ndarray | float) – Metric values on this rank.
Returns:: Mean across all ranks.
Return type:: float

agilerl.distributed.aggregate_metrics_dict(metrics: dict[str, Tensor | ndarray | float], excluded: Collection[str] = (), maximized: Collection[str] = ()) → dict[str, float]¶

Reduce every metric across ranks in one all-reduce: the mean, or the max for maximized.

Every rank must pass the same keys in the same order. On a single device each metric is its local mean.

Parameters:

metrics (dict[str, torch.Tensor | np.ndarray | float]) – Metric values on this rank, by name.
excluded (Collection[str]) – Metrics this rank leaves out of their cross-rank mean. A metric every rank excludes averages to NaN.
maximized (Collection[str]) – Metrics reduced to their largest value on any rank; excluded does not apply to them.

Returns:

Mean of each metric across the ranks that count it, or its max across ranks for maximized.

Return type:

dict[str, float]

agilerl.distributed.sync_grads(params: Sequence[Parameter]) → None¶

Average gradients across data-parallel ranks.

One coalesced all-reduce of every .grad in params (SUM, then divide by world size so Gloo works). See all_reduce_grads() for params without a grad. No-op on a single device.

Parameters:: params – Optimizer parameters whose .grad should be averaged.

agilerl.distributed.materialize_dtensors(t0: torch.Tensor, t1: torch.Tensor | None, /) → _GeneratorContextManager[tuple[torch.Tensor, torch.Tensor | None], None, None]¶

agilerl.distributed.materialize_dtensors(*tensors: torch.Tensor | None) → _GeneratorContextManager[list[torch.Tensor | None], None, None]

All-gather DTensor shards to dense locals without swapping module params.

Prefer this for ephemeral matmuls (fused lm_head logprobs, Liger). Use gather_params() when in-module reads must see dense weights (state_dict / PEFT save_pretrained). All ranks must enter and exit together. Yields a list parallel to tensors.

Parameters:: tensors (torch.Tensor | None) – Tensors to densify; None passes through.

agilerl.distributed.gather_params(root: Module, params: Sequence[Tensor | None]) → Generator[list[Tensor | None], None, None]¶

Materialize full (unsharded) views of params for the duration of the context.

Plain tensors are left unchanged. For FSDP2 DTensor parameters, each tensor is all-gathered with full_tensor() and temporarily installed on its owning module so in-module reads (e.g. state_dict / PEFT save_pretrained) see dense weights. Original shards are restored on exit.

Yields a list parallel to params: dense locals for any gathered DTensor, and the original handles otherwise. Callers that hold pre-gather tensor references must use the yielded list for math — those references still point at the shard. For matmul-only gathers prefer materialize_dtensors() (no module Parameter install).

Gathered parameters are read-only: writes are discarded when the sharded DTensor is restored. Write into a sharded model with set_full_model_state_dict(). All ranks must enter and exit together.

Parameters:

root (nn.Module) – Module whose submodules own params.
params (Sequence[torch.Tensor | None]) – Parameters to gather; None entries pass through.

agilerl.distributed.full_shape_views(root: Module, params: Sequence[Tensor | None]) → Generator[None, None, None]¶

Expose global-shape views on FSDP2 DTensor params for shape-only reads.

Each sharded param is temporarily replaced on its owning module with a zero-storage view (a scalar expanded to the DTensor global shape) so shape, dtype and device reads see the full tensor without full_tensor(). Values must not be read inside the block; the original DTensor is restored on exit when the installed view is still the live parameter. Plain tensors pass through untouched. Duplicate references are processed once (by identity).

Parameters:

root (nn.Module) – Module whose submodules own params.
params (Sequence[torch.Tensor | None]) – Parameters to expose; None entries are skipped.

agilerl.distributed.resolve_device(requested: str | device | None = None) → str¶

Pick the training device.

Distributed runs are pinned to cuda:<local_rank>, or CPU without CUDA; otherwise the requested device (or the best available) is used.

Parameters:: requested (str | torch.device | None) – Device requested by the caller, if any.
Returns:: Device string.
Return type:: str

class agilerl.distributed.FSDPConfig(reshard_after_forward: Annotated[bool, FieldInfo(annotation=NoneType, required=True, description="Free gathered parameters after each module's forward. False keeps them gathered: faster, more VRAM.")] = True, cpu_offload: Annotated[bool, FieldInfo(annotation=NoneType, required=True, description='Offload sharded parameters and gradients to CPU. Needs colocated vLLM. Mutually exclusive with optim_cpu_offload.')] = False, optim_cpu_offload: Annotated[bool | None, FieldInfo(annotation=NoneType, required=True, description='Keep Adam state on CPU and step it there with the default AdamW kernel. Parameters and gradients stay on GPU. False keeps Adam on GPU and steps it with fused AdamW. None picks GPU when the memory estimate fits, else CPU.')] | None = None, defer_grad_sync: Annotated[bool, FieldInfo(annotation=NoneType, required=True, description='Reduce-scatter gradients only on the last micro-batch of an optimizer step. Saves communication; holds unsharded grads in between.')] = True, param_dtype: Annotated[str, FieldInfo(annotation=NoneType, required=True, description="Mixed-precision parameter dtype, e.g. 'bfloat16'. A torch.dtype is accepted and stored by name.")] = 'bfloat16', reduce_dtype: Annotated[str, FieldInfo(annotation=NoneType, required=True, description='Dtype for gradient reduce-scatter and all-reduce.')] = 'float32', prefetch_units: Annotated[int, FieldInfo(annotation=NoneType, required=True, description='Neighbouring FSDP units to all-gather ahead during forward. 1 gathers unit i+1 while unit i runs.', metadata=[Ge(ge=1)])] = 1, backward_prefetch_units: Annotated[int, FieldInfo(annotation=NoneType, required=True, description='Neighbouring FSDP units to all-gather ahead during backward. 1 gathers unit i-1 while unit i runs backward.', metadata=[Ge(ge=1)])] = 1, checkpoint_skip_layer_types: Annotated[tuple[str, ...], FieldInfo(annotation=NoneType, required=True, description="Transformer block kinds left out of activation checkpointing. A block's kind is its block_type (hybrid models, e.g. 'linear_attention', 'full_attention', 'moe') or else its class name. Empty checkpoints every block. Applies when gradient_checkpointing is on.")] = (), checkpoint_every_n_blocks: Annotated[int, FieldInfo(annotation=NoneType, required=True, description='Checkpoint one in every n blocks not skipped by checkpoint_skip_layer_types. 1 checkpoints all of them.', metadata=[Ge(ge=1)])] = 1, wrap_every_n_blocks: Annotated[int, FieldInfo(annotation=NoneType, required=True, description='Consecutive transformer blocks per FSDP unit. Larger values mean fewer, bigger collectives.', metadata=[Ge(ge=1)])] = 1, param_persistence_threshold: Annotated[int, FieldInfo(annotation=NoneType, required=True, description='Parameters with fewer elements than this stay unsharded. 0 shards every parameter.', metadata=[Ge(ge=0)])] = 100000, ep: Annotated[int, FieldInfo(annotation=NoneType, required=True, description='Expert-parallel degree: packed MoE experts per layer are split across this many GPUs. 1 keeps data parallel plus FSDP sharding with no expert split.', metadata=[Ge(ge=1)])] = 1, ep_token_blocks: Annotated[int, FieldInfo(annotation=NoneType, required=True, description="Token blocks per routed MoE layer under expert parallel. Each block runs dispatch, experts, and combine; on GPU the next block's all-to-all overlaps this block's experts. 1 moves every token in one all-to-all.", metadata=[Ge(ge=1)])] = 1, routed_expert_chunk_mib: Annotated[int | None, FieldInfo(annotation=NoneType, required=True, description='MiB of the widest activation of one routed-expert row chunk under expert LoRA, of one fp32 chunk of the expert-parallel combine, and of one grouped LoRA GEMM output. Bigger chunks launch fewer kernels per MoE layer; on Super-VL the learn peak grows ~0.37 GiB per GiB of chunk. None picks the largest of 64/128/256/512 the memory estimate fits.', metadata=[Ge(ge=1)])] | None = None, tp: Annotated[int, FieldInfo(annotation=NoneType, required=True, description='Tensor-parallel degree for dense layers. Ranks in one TP group share a batch shard. 1 gives every rank its own batch.', metadata=[Ge(ge=1)])] = 1, shard_group_size: Annotated[int | None, FieldInfo(annotation=NoneType, required=True, description='Ranks that shard weights between them; groups replicate (HSDP). Set to GPUs per node so weight gathers stay inside a node. None shards across all trainer ranks.', metadata=[Ge(ge=1)])] | None = None)¶: Settings for sharding the LLM actor with PyTorch FSDP2 (fully_shard).

agilerl.distributed.apply_fsdp2(model: Module, config: FSDPConfig | None = None, mesh: DeviceMesh | None = None, expert_mesh: DeviceMesh | None = None, gradient_checkpointing: bool = False) → Module¶

Shard model with FSDP2: blocks, PLE, embed/(untied) lm_head, root, prefetch.

Parameters become DTensors in place, so any optimizer must be (re)built after this call. Callers must not pass a dense full-model replica on CUDA — use materialize_fsdp2_from_cpu_state() so weights stay on CPU/meta until only local shards are allocated on the compute device.

When mesh is provided, non-expert fully_shard uses that mesh; when None, FSDP uses the default process group (flat path). Packed expert modules are sharded on expert_mesh when that is provided.

When gradient_checkpointing is on, the transformer blocks chosen by config.checkpoint_skip_layer_types and config.checkpoint_every_n_blocks are wrapped with checkpoint_wrapper before fully_shard so the checkpoint boundary sits inside the FSDP unit. Blocks holding packed experts reuse their forward routing in the recompute; their listed routers (ROUTER_FORWARDS) get out-of-place top-k renormalization.

Parameters:

model (nn.Module) – Model to shard (CPU or meta parameters).
config (FSDPConfig | None) – Sharding settings; defaults to FSDPConfig’s defaults.
mesh (DeviceMesh | None) – Optional FSDP device mesh for non-expert units. Ignored when None.
expert_mesh (DeviceMesh | None) – Optional FSDP mesh for packed-expert modules (the leftover data-parallel axis). Ignored when None.
gradient_checkpointing (bool) – Wrap the transformer blocks selected by config with non-reentrant activation checkpointing before sharding.

Returns:

The sharded model (same object).

Return type:

nn.Module