Data parallelism at bf16 — every node shards exactly but the conv weight gradient #
DataParallel.Sync states P3/P4 at the f32 nodes. The bf16 data-parallel renders ResNet-34
and ResNet-50 train from (resnet34in_momdp64bf16, resnet50in_momdp64bf16,
resnet50in160_lambaccdp8x64wxclipbcebf16) swap six of those nodes for bf16 kinds, and each
has its own den: operands rounded going in, the result rounded once on the way out. This file
is what those six kinds do under sharding.
| kind | under sharding | lemma |
|---|---|---|
convBf16, convStridedBf16 (forward, .batchOp) | exact | den_convBf16_shard, den_convStridedBf16_shard |
convBackBatchedBf16, convStridedBackBatchedBf16 | exact | den_convBackBatchedBf16_shard, den_convStridedBackBatchedBf16_shard |
convWeightGradBBf16, convStridedWeightGradBBf16, all-reduced | rounded per replica | den_allReduceMeanF_convWeightGradBBf16_sub_global, den_allReduceMeanF_convStridedWeightGradBBf16_sub_global |
The first four round per element of a per-example map, so replica r's value is batchShard r
of the same node at batch R·N, exactly as at f32.
The weight gradients are not. The emitted weight-gradient convolution contracts the batch
in one op and stores bf16 once, so the node's den is rnd (Σ_n …) with the rounding outside
the batch sum (den_convWeightGradBBf16_eq_rnd). On R replicas each rounds its OWN partial
sum S_r before the f32 all-reduce; on one device the global sum Σ_r S_r is rounded once. So
the collective is (1/R)·Σ_r rnd S_r where the batch-R·N node is rnd (Σ_r S_r), and
den_allReduceMeanF_convWeightGradBBf16_sub_global states the difference exactly. At
rnd := id the two agree (…_shard_id), which is the f32 statement.
The divisor step at bf16 #
A DP render divides its loss by the per-replica batch, so its cotangents are R × the shard of
the global ones (DataParallel.Sync, "The 1/R"). At f32 linearity carries that factor
through every backward; at bf16 it has to pass through rnd as well. The *_smul lemmas below
take that as a hypothesis, ∀ x, rnd (s * x) = s * rnd x, and rndP_zpow_mul (in RndP.lean) proves it for
the repo's rounding model at every integer power of two — so for bf16 (rndP 7) at s = R and
s = 1/R whenever R is a power of two, the ImageNet runs' R = 4 among them.
What is NOT claimed #
Note: No whole-net statement at a real rounding. The whole-net sync ties are stated at either
precision at the identity rounding the renderers use (Bf16Erasure, GradNodesBAt); one at
rnd := rndP 7 would need a chain threaded with rnd, which no net has (planning/bf16_tie.md §4).
These are the per-node cases such a chain would walk with.
Note: rndP has an unbounded exponent (RndP.lean): bf16 overflow and subnormals are
outside it, where scaling by 4 can move a value across the format's boundary.
A per-example node on replica r is shard r of the same node at batch R·N. Stated
node to node, because at bf16 there is no ℝ-level forward to name on the right.
The bf16 forward conv shards exactly — every 1×1 and 3×3 in the ResNet bf16 renders.
…and so does the bf16 symmetric stride-2 conv — the 7×7 stem and the downsamples.
The bf16 conv input-VJP shards exactly. Its rounding is per element of a per-example map, so it never sees another replica's rows.
…and so does the bf16 strided conv input-VJP.
The bf16 weight-gradient node is the f32 node at rounded operands, rounded once.
…and the strided one.
Replica r's partial sum: the f32 conv weight-gradient node on shard r of the rounded
global operands. The collective and the batch-R·N node are both built from these; they
differ only in where rnd is applied.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The strided peer of convWGradShardSum.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The all-reduced bf16 weight gradient is the mean of the replicas' ROUNDED partial
sums. Each replica on its shard, at the shard-r block of the global cotangent.
The batch-R·N bf16 weight gradient rounds the SUM of the same partial sums, once.
The one difference, exactly. The all-reduced bf16 weight gradient minus 1/R of the
batch-R·N bf16 node is 1/R of (sum of the rounded partial sums − the rounded sum).
At rnd := id the difference vanishes — the collective is 1/R of the batch-R·N
node, which is the f32 statement den_allReduceMeanF_convWeightGradB_shard.
Strided: the all-reduced bf16 weight gradient is the mean of the rounded partial sums.
Strided: the batch-R·N node rounds the sum of the same partial sums, once.
Strided: the one difference, exactly.
rndP_zpow_mul and the Int.log shift lemmas behind it live in RndP.lean.
The bf16 conv input-VJP scales with its cotangent when rnd commutes with the scale.
…and the strided one.
The bf16 conv weight gradient scales with its cotangent when rnd commutes with it.
…and the strided one.