Gradient utilities

Two primitives govern how gradients move through the training loop: maybe_no_sync controls when they get reduced, clip_grad_norm_ controls their magnitude.

maybe_no_sync

from kempnerforge.training.grad import maybe_no_sync

for micro_step in range(grad_accum_steps):
    with maybe_no_sync(model, micro_step, grad_accum_steps):
        logits = model(input_ids, doc_ids=doc_ids)
        loss = loss_fn(logits, labels)
        (loss / grad_accum_steps).backward()

Context manager from kempnerforge/training/grad.py. On all microbatches except the last:

model.set_requires_gradient_sync(False)
try: yield
finally: model.set_requires_gradient_sync(True)

Why:

  • FSDP2 fires a reduce-scatter collective at the end of every backward pass. Without maybe_no_sync, a step with grad_accum_steps = 8 fires 8 collectives per optimizer step. With it, the first 7 accumulate locally and the 8th fires one collective that covers all accumulated gradients.

  • The method is FSDP2-specific (set_requires_gradient_sync). On non-FSDP models (DDP, single GPU), hasattr(model, "set_requires_gradient_sync") is False and the context is a no-op — safe to use unconditionally.

  • On the last microbatch (micro_step + 1 == grad_accum_steps), the context skips straight to yield so the final backward triggers a normal reduce-scatter.

Not used under pipeline parallel — the PP schedule manages its own sync internally. See Training loop § PP step.

clip_grad_norm_

from kempnerforge.distributed import clip_grad_norm_

grad_norm = clip_grad_norm_(model, tc.grad_clip_norm)

Lives in kempnerforge/distributed/utils.py (re-exported from kempnerforge.distributed). Wraps torch.nn.utils.clip_grad_norm_ with an extra path for mixed DTensor meshes.

The problem it solves: when a model combines TP and FSDP, some parameters are DTensors on the (dp_shard, tp) mesh (TP-sharded linears) and others are on (dp_shard,) alone (norm scales, biases). PyTorch’s clip_grad_norm_ doesn’t know how to combine gradients living on different meshes — it would produce an incorrect norm.

Algorithm:

  1. Collect all non-None .grad tensors from the model.

  2. Build a mesh_key per gradient (DTensor’s _spec.mesh id, or 0 for plain tensors).

  3. If only one mesh (pure FSDP, single GPU, plain tensors) and no pipeline or expert parallelism: call torch.nn.utils.clip_grad_norm_ directly — the fast path. Under either of those, the norm is summed over the stages’ or the EP group as described below.

  4. If multiple meshes: group gradients by mesh and compute each group’s sum-of-squares from per-gradient norms taken in at least fp32 (fp16 and bf16 are promoted, fp64 stays fp64). A DTensor norm is partial over its shards, and squaring it runs the all-reduce that replicates it, so .full_tensor() only unwraps each group’s sum. The groups’ sums are added and combined with a plain sqrt. In fp16, squaring a norm above about 256 would overflow to inf and the clip would zero every gradient.

  5. Scale every gradient by clip_coef = max_norm / (total_norm + 1e-6), clamped to ≤ 1.

Returns the pre-clip total norm as a scalar tensor — the value you log as grad_norm in metrics.

Under pipeline parallelism

Each PP rank holds one stage, and a stage’s own norm is not the model’s. Pass the stages’ process group as pp_group:

grad_norm = clip_grad_norm_(stage, tc.grad_clip_norm, pp_group=device_mesh["pp"].get_group())

Each stage computes its sum of squares as in step 4, the sums are added over pp_group, and every stage scales by the coefficient of that whole-model norm. The clipped model’s norm is then at most max_norm, and every stage returns, and logs, the same pre-clip norm. Clipping each stage on its own instead leaves the model at sqrt(pp) × max_norm when every stage is above the threshold, and the different coefficients change the update’s direction, not only its length. Every stage must make the call, including one without gradients, since the sum is a collective. pipeline_step in training/loop.py passes the group.

Under expert parallelism

With ep > 1 each EP rank holds its own slice of the routed experts and a copy of every other parameter, so a rank’s own norm is not the model’s either. No argument is needed: clip_grad_norm_ reads the EP group from the model’s MoEMLP layers (expert_parallel_params in kempnerforge/distributed/expert_parallel.py) and, in one all-reduce over that group, sums the routed experts’ squared norms and averages the squared norms of the copies. Every EP rank then returns the whole model’s norm and scales by the same coefficient, so its copies of the shared parameters stay equal to the other ranks’ as long as every EP rank computes them identically; the norm is the same on every EP rank even when the copies differ in the last bits. Clipping each rank by its own norm instead scales the copies by different coefficients, and nothing reduces them over ep afterwards, so they drift apart. On GPU the copies drift apart even with one coefficient, with or without clipping, and nothing re-syncs them yet (#482). Every EP rank must make the call, since the sum is a collective. How the expert gradients themselves are averaged over the group is in Expert parallelism § Gradients and clipping.

foreach parameter

clip_grad_norm_(..., foreach=True) is the default — faster on CUDA via the fused foreach implementation. Only touched if your model contains some exotic parameter type that the foreach path can’t handle.

Gradient accumulation in practice

The full microbatching pattern in training/loop.py is:

for micro_step in range(tc.grad_accum_steps):
    with maybe_no_sync(model, micro_step, tc.grad_accum_steps):
        logits = model(input_ids, doc_ids=doc_ids)
        loss   = loss_fn(logits, labels)
        scaled_loss = loss / tc.grad_accum_steps
        scaled_loss.backward()
        total_loss += loss.item()

grad_norm = clip_grad_norm_(model, tc.grad_clip_norm)
optimizer.step()
optimizer.zero_grad()

Two invariants:

  • Loss scaling: scaled_loss = loss / grad_accum_steps keeps the effective LR independent of the accumulation factor. If you change grad_accum_steps=4 → 8, the optimizer step size stays the same per-token.

  • Clip after accumulate: clip_grad_norm_ runs after the microbatch loop, not inside it — the clip sees the final accumulated gradient.

Config fields

[train]
grad_accum_steps = 8                    # microbatches per optimizer step
grad_clip_norm = 1.0                    # max grad L2 norm

Both live in TrainConfig. TrainConfig.__post_init__ requires grad_clip_norm > 0 — there is no “disable clipping” escape hatch; 1.0 is a near-free safety margin and what every shipped config uses.

See also