DDP Loss Normalization with Unequal Token Counts
DDP synchronizes gradients; it does not automatically turn an average of per-rank mean losses into a mean over all valid tokens. With unequal valid-token counts, averaging local means weights ranks equally rather than tokens equally.
Contents
- Write the intended objective first
- One rank has one token; another has three
- Why the local loss needs the W factor
- Nine checks compare the gradients
- No local labels does not mean a rank can exit alone
- Accumulation extends the denominator to the update unit
- Separate reporting scalars from backward scalars
- Use partition invariance to test the objective
- Sources and backfill note
For a token-mean objective, let N be the valid-token count for the update and W the number of processes participating in gradient averaging. Backpropagating W·Sᵣ/N on each rank, where Sᵣ is its local loss sum, compensates for that averaging under the assumptions below. Exact checks cover unequal partitions, equal partitions and an empty local partition. 中文版.
Write the intended objective first
Let rank r have nᵣ valid tokens and local loss sum Sᵣ. The intended objective is L=(ΣᵣSᵣ)/(Σᵣnᵣ), with N=Σᵣnᵣ. Validity is a task definition, such as positions with supervised labels. Padded tensor length need not equal that count.
The DDP documentation describes averaging gradients across processes. This derivation assumes a fixed participating group, ordinary gradient averaging and coordinated execution of one update. It does not assign the same divisor to custom communication hooks, dynamic join behavior or another parallel wrapper.
N is independent of model parameters here. Weighted targets may require a different normalization mass according to their actual loss definition. A variable named num_tokens does not override a criterion that already uses another weighted-mean denominator.
The comparison also assumes each token has the same loss function and value after repartitioning. Batch-dependent layers or changed randomness can violate that premise. Control those effects in the reference test before attributing a difference to normalization alone.
One rank has one token; another has three
Use one parameter θ and teaching loss ℓ(θ,y)=½(θ−y)². At θ=0, the first rank has target zero and derivative zero. The second has three targets equal to two, each with derivative minus two, totaling minus six.
Local means are zero and minus two. Averaging them gives minus one. The mean derivative across all four valid tokens is instead (0−2−2−2)/4=−1.5. The first calculation gives each rank half the weight; the second gives each token one-quarter. This is an objective mismatch, not communication roundoff.
Repartition the same targets into two tokens per rank and local-mean averaging gives minus 1.5 in this example. The incorrect implementation changes its gradient when only partitioning changes. That supplies a sharper regression test than asking whether a distributed training curve appears plausible.
Why the local loss needs the W factor
Choose the local backward scalar Lᵣ=W·Sᵣ/N. Ordinary DDP averaging gives (1/W)Σᵣ∇Lᵣ=(1/W)Σᵣ(W/N)∇Sᵣ=(Σᵣ∇Sᵣ)/N. The result is the gradient of the global token mean. First use a common denominator, then compensate for the known averaging divisor.
If a local mean Sᵣ/nᵣ already exists and nᵣ is positive, multiplying it by W·nᵣ/N is equivalent. An explicit sum is easier to define for a zero-valid-token partition: it avoids constructing an undefined local mean and attempting to repair it afterward by multiplying by zero.
If communication returns a sum instead of a mean, that W compensation would over-scale the gradient. Re-derive the reduction when using a custom hook, changing process groups or selecting another wrapper. W means the size of the group participating in this reduction, not a guessed machine count or stale configuration value.
Nine checks compare the gradients
The code directly computes the analytical derivative θ−y of the square loss using Fraction. It starts no processes, calls no all_reduce and runs no PyTorch. It checks normalization algebra rather than communication correctness or GPU performance.
# Exact gradient algebra; no distributed framework is executed.
from fractions import Fraction as F
def global_gradient(groups, theta=F(0)):
values=[v for group in groups for v in group]
if not values:
raise ValueError("no valid tokens")
return sum((theta-y for y in values),F(0))/len(values)
def averaged_local_means(groups, theta=F(0)):
return sum((sum((theta-y for y in g),F(0))/len(g)
for g in groups),F(0))/len(groups)
def scaled_sum_gradient(groups, theta=F(0)):
total=sum(map(len,groups))
if total==0:
raise ValueError("no valid tokens")
world=len(groups)
local=[F(world,total)*sum((theta-y for y in g),F(0)) for g in groups]
return sum(local,F(0))/world
groups=[[F(0)],[F(2),F(2),F(2)]]
assert global_gradient(groups)==F(-3,2)
assert averaged_local_means(groups)==-1
assert scaled_sum_gradient(groups)==F(-3,2)
assert scaled_sum_gradient([[F(0),F(2)],[F(2),F(2)]])==F(-3,2)
assert averaged_local_means([[F(0),F(2)],[F(2),F(2)]])==F(-3,2)
assert scaled_sum_gradient([[],[F(2),F(2)]])==-2
assert scaled_sum_gradient(groups,F(1))==global_gradient(groups,F(1))
assert F(1,2)!=F(1,4)
try:
scaled_sum_gradient([[],[]])
except ValueError:
pass
else:
raise AssertionError("expected zero-count rejection")
print("9 DDP normalization checks passed")
All nine checks passed locally. Unequal partitions expose the local-mean bias, while equal partitions provide a control. An empty local partition has a defined zero-sum contribution when the global count is positive. An all-empty update is explicitly rejected. Changing θ checks that the comparison is not an accident specific to zero.
No local labels does not mean a rank can exit alone
Mathematically, nᵣ=0 with N positive gives a zero local loss-sum contribution. Distributed execution additionally requires compatible synchronization order. A rank cannot simply skip backward or a collective because it has no labels while others continue waiting.
A real implementation needs a zero-contribution path compatible with its model graph and DDP configuration, validated in the actual environment. A newly created constant zero unrelated to the model does not automatically trigger the required gradient synchronization. The empty list in this standard-library example expresses arithmetic only; it is not executable training guidance for that path.
When global N is zero, all ranks need a consistent, predetermined policy, such as jointly skipping the update unit. Define whether optimizer, scheduler and reporting step counts advance as well. Independent decisions based on local counts cannot establish a coordinated step. Data validity and communication order both matter.
Accumulation extends the denominator to the update unit
If several microbatches precede one optimizer.step, N normally covers all valid tokens in that update. Using a separate N for every microbatch can recreate an average of means inside the distributed calculation.
One approach determines update-unit counts beforehand and scales each loss sum against the common objective; another uses a verified equivalent gradient rescaling. The concrete design must agree with no_sync, AMP, clipping and the final synchronization boundary. A single-step algebraic identity should not be pasted untested into a complete loop.
The earlier gradient-accumulation reduction article tracks weights across microbatches. This article adds averaging across ranks. Follow both denominators to connect one optimizer update to a single stated objective.
Separate reporting scalars from backward scalars
Using W·Sᵣ/N for backward does not make each local scalar a global training metric. Reporting can separately aggregate detached loss sums and valid counts, then compute ΣSᵣ/N. Keeping per-partition counts also explains changes in update size.
Do not reconnect an arbitrary cross-rank logging aggregate into the training graph merely to show the same scalar everywhere. Differentiable collectives and ordinary metric aggregation have different autograd contracts that require actual API validation. Identical displayed numbers cannot prove correct backward behavior.
A final implementation test fixes data and initial parameters, compares against a single-process combined batch, and varies partition sizes. Allow appropriate floating-point reduction tolerances after excluding integer-factor and token-weighting errors. That multiprocess acceptance test was not run here, so the arithmetic checks do not certify a DDP implementation.
Use partition invariance to test the objective
Reassigning the same valid tokens to different ranks should not change an intended global token mean. Test that requirement directly before compensating with learning-rate changes. A learning rate can hide a scale discrepancy without restoring the relative weight of each token.
If the derivation and repartitioning test help locate a distributed normalization error, voluntary support through the section below is welcome. The reusable outcome is a precise explanation of whose gradients each update averages.
Sources and backfill note
Definitions were checked against versioned PyTorch 2.8 material available before the planned slot. The derivation, arithmetic code and two figures are original teaching material. The planned article date is October 18, 2026 at 23:30 Beijing time; this missed slot was backfilled on October 19, 2026, with the actual publication time reported separately. No PyTorch training, distributed process or GPU benchmark ran.
Support
If the explanations or examples helped you solve a problem or save time, you’re welcome to support this site with a voluntary tip.
Click a QR code to enlarge. More options: support page。


