Data parallelism at bf16 — every node shards exactly but the conv weight gradient #
DataParallelSync.lean 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 and its strided peer |
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 (DataParallelSync.lean, "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_two_pow_mul proves it for
the repo's rounding model at every power of two — so for bf16 (rndP 7) at R = 4
(rndP_mul_four), the replica count of every ImageNet run.
What is NOT claimed #
⚠ No whole-net statement: there is no single-device bf16 chain for either net to tie a twin to,
here or in the f32 tier. These are the per-node cases a bf16 twin would walk its chain with.
⚠ rndP has an unbounded exponent (Binary32Instance.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.
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.