semiAIfoundry

Executable software · Atlas v1.3.0

Uneven batches and the global gradient

Preserve the intended training objective when ranks contribute different numbers of valid examples or tokens.

Explore in AtlasLogical values and distributed ownership; no physical scale or vendor timing

Learning objectiveDerive and test a globally normalized gradient from unequal local batch means, then distinguish objective equivalence from replica synchronization.

The mechanism

A synchronized gradient can be consistent across replicas and still represent the wrong objective. Suppose the intended loss is the mean over all valid training items in one optimizer update. A rank that averages two items and a rank that averages six have already divided by different denominators. Averaging their two gradient means assigns the smaller batch more influence per item. The global item mean instead weights each local mean by its valid-item count. PyTorch’s pinned DistributedDataParallel contract documents averaging gradients between workers, so the loss normalization and collective normalization must be considered together. The count can mean examples, unmasked target tokens or a sum of explicit sample weights; choosing it is part of the objective, not a hardware detail. This lesson derives one scalar update from a concrete squared-error loss. It assumes the same current parameter at both ranks and no cross-sample operation such as batch-dependent normalization. A rank exhausting its entire input stream is a separate synchronization problem: matching collective participation does not by itself fix per-item weighting.

  1. Define which items count in the loss and record their global count N for this optimizer update. Exclude padding only when the objective says to exclude it.
  2. Compute local loss sums or local means and retain the valid count associated with each reduction.
  3. For a rank-mean reducer over P active participants, scale a local mean gradient by P nᵣ/N before the averaging collective. Equivalently reduce local sums and divide once by N.
  4. During gradient accumulation, use the count over the entire update window rather than silently averaging unequal microbatch means.
  5. Apply the same synchronized gradient and optimizer state transition at every replica, then compare the result with the unpartitioned reference objective.
  6. Handle empty local contributions and exhausted ranks with explicit participation rules; never divide an empty local mean by zero or simply omit a required collective.

Science and assumptions

Global item-mean gradient

N=∑r=0P−1nr,g=1N∑r=0P−1nrgˉrN=\sum_{r=0}^{P-1}n_r,\qquad g=\frac{1}{N}\sum_{r=0}^{P-1}n_r\bar g_r

Sum local valid-item counts to N. Weight each local mean gradient by its count and divide their sum by N.

Symbols and units
N
Total valid equally weighted items in the update

count

nr
Valid items at rank r

count

P
Participating rank count

count

ḡr
Mean local gradient at shared parameter state

loss/parameter

g
Global item-mean gradient

loss/parameter

Each local ḡᵣ is the mean of gradients for nᵣ valid equally weighted items. Multiply by nᵣ to recover the local sum.

Assumptions N>0, same parameters and differentiable separable per-item loss. No cross-sample dependence; sample weights replace counts when the objective uses unequal weights.

Compensate a mean-across-ranks reducer

g=1P∑r=0P−1(PnrNgˉr),wnew=w−ηgg=\frac{1}{P}\sum_{r=0}^{P-1}\left(\frac{P n_r}{N}\bar g_r\right),\qquad w_{\mathrm{new}}=w-\eta g

Before a mean-across-ranks reduction, scale local mean gradients by P times nrn_{r} over N. Scalar SGD subtracts learning rate times global gradient.

Symbols and units
P
Fixed participants in the mean reducer

count

nr, N
Local and global valid-item counts

count

ḡr, g
Local and global mean gradients

loss/parameter

w, wnew
Parameter before and after the update

parameter

η
Learning rate for the declared objective scaling

parameter²/loss

The factor P nᵣ/N ensures the subsequent division by P leaves the intended global item mean.

Assumptions P fixed participating ranks for this update; no custom reducer scaling. The scalar SGD update has no momentum, clipping, weight decay or loss scaling.

Engineering example

Take ℓ(w;y)=½(w−y)² at w=10. Rank 0 has targets [9,9], producing gradients [1,1], count 2 and mean 1. Rank 1 has [7,7,7,7,7,7], producing six gradients of 3, count 6 and mean 3. The unpartitioned objective has gradient (2×1+6×3)/8=2.5. Averaging rank means gives (1+3)/2=2 instead. With learning rate 0.1, correct SGD gives wnew=9.75; the unweighted rank mean gives 9.8. For a reducer that divides by P=2, scale the local means by 2×2/8=0.5 and 2×6/8=1.5. It then averages 0.5 and 4.5 to obtain 2.5. If the ranks each held four valid items, equal weighting would be appropriate. Token-level training needs the valid target-token count when that is the declared loss denominator; sequence counts need not be proportional.

Distinctions and tradeoffs

A mean across ranks is a mean across training items only when local denominators and weights make it so.

Keeping exact counts and a reference objective adds bookkeeping and sometimes a small count reduction. It prevents an apparent scaling improvement from silently changing sample influence.

The number of completed updates is useful only relative to the training objective those updates optimize. Ragged data, masking and uneven accumulation can change that objective through an unnoticed denominator.

Specifications and interfaces

Normalization unit
Valid examples or tokens

Choose the intended loss denominator explicitly.

Replica consistency
Equal synchronized gradients

Does not establish agreement with the desired global objective.

Runtime scope
Default mean reduction

Custom hooks, joined ranks or summed losses can change the scale.

Interfaces

Per-rank loss contributions Input
Local sums or means, valid counts, shared parameter state
Global optimizer update Output
One correctly normalized gradient applied consistently

Check your understanding

Two ranks have counts 2 and 6 and mean gradients 1 and 3. What is the global item-mean gradient?

  1. 2
  2. 2.5
  3. 4
Answer and reasoning

2.5

The weighted numerator is 2×1+6×3=20 across eight items, so the result is 2.5.

A mean-across-two-ranks reducer receives local means. Which scale factors preserve the eight-item objective?

  1. 1 and 1
  2. 2 and 6
  3. 0.5 and 1.5
Answer and reasoning

0.5 and 1.5

Scaling by P nᵣ/N gives 0.5 and 1.5; the reducer then averages 0.5 and 4.5.

Does equal synchronized state prove the chosen sample weights were correct?

  1. Yes
  2. No
  3. Only if the network is fast
Answer and reasoning

No

All replicas can agree on the same incorrectly normalized gradient.

Evidence and scope

The PyTorch source is pinned to v2.8.0 and supports the default averaging convention, not a claim about every backend or custom hook. The squared-error data, gradient weighting derivation and SGD step are authored examples. Join semantics after input exhaustion are outside this fixed-participant update.

The ideal ring all-reduce lab varies communication payload, ranks, startup and bandwidth. It does not execute these matrix shards or alter the optimizer; the exact computation is worked through in this lesson.

Sources

Published lesson data · Atlas releases and licensing

Continue exploring

Foundations

Related lessons