kempnerforge.distributed.utils

Distributed training utilities.

Gradient clipping and data-parallel helpers that work correctly with FSDP2 and multi-dimensional parallelism.

Functions

agree_across_ranks(error, what, *[, group])

Barrier on group that raises on every rank if any rank's what failed.

clip_grad_norm_(model, max_norm[, foreach, ...])

Clip gradient norm across all parameters.

get_dp_info(device_mesh)

Get (dp_rank, dp_size) from the device mesh, accounting for PP/TP/EP.
kempnerforge.distributed.utils.get_dp_info(device_mesh)[source]

Get (dp_rank, dp_size) from the device mesh, accounting for PP/TP/EP.

Only dp_replicate and dp_shard split the data, so ranks that differ only in pp, ep or tp share a DP rank. Under HSDP the DP rank is the row-major index of the rank’s (dp_replicate, dp_shard) coordinates and the DP size is their product; no process group is created.

Parameters:

device_mesh (DeviceMesh | None) – Full DeviceMesh, or None for single-GPU.

Returns:

Tuple of (dp_rank, dp_world_size); (0, 1) when there is no mesh or the mesh has no DP dimension.

Return type:

tuple[int, int]

kempnerforge.distributed.utils.clip_grad_norm_(model, max_norm, foreach=True, pp_group=None)[source]

Clip gradient norm across all parameters.

Handles mixed DTensor meshes that arise when TP+FSDP produces parameters on different meshes (e.g., TP-sharded linears on (dp_shard, tp) vs FSDP-only norms on (dp_shard)). Groups gradients by mesh, computes per-group norms via stack, then combines across groups.

Falls back to torch.nn.utils.clip_grad_norm_ when there is only one mesh (pure FSDP or single-GPU) and neither pipeline nor expert parallelism, whose norms need a sum over their own group.

Under pipeline parallelism model holds one stage, so pass the stages’ process group as pp_group: the squared stage norms are summed over it, and every stage returns the whole model’s norm and scales by the same coefficient. Every stage must make the call, including one without gradients, since the sum is a collective.

Under expert parallelism each EP rank holds its own routed experts and a copy of every other parameter. The experts’ squared norms are summed over the EP group the MoEMLP layers record, and the copies’ are averaged over it, so every EP rank returns the whole model’s norm and scales by the same coefficient. Every EP rank must make the call.

Parameters:

  • model (torch.nn.Module) – Model whose gradients to clip.

  • max_norm (float) – Maximum gradient norm.

  • foreach (bool) – Use the faster foreach implementation.

  • pp_group (dist.ProcessGroup | None) – Process group of the pipeline stages, or None without pipeline parallelism.

Returns:

Total gradient norm (before clipping).

Return type:

torch.Tensor

kempnerforge.distributed.utils.agree_across_ranks(error, what, *, group=None)[source]

Barrier on group that raises on every rank if any rank’s what failed.

Every rank passes the error its own part of what raised, or None. A rank with an error re-raises it; when any rank has one, every other rank raises a RuntimeError naming the ranks that failed, so no rank is left waiting in a later collective for a rank that has raised. Without an initialized process group, error is raised if set.

Parameters:
Return type:

None