ViT-Tiny data-parallel: R replicas of N are one device at R·N, through the optimizer and the EMA #
ViTParamGrad proves, at one replica, that each of ViT-Tiny's 200 gradient nodes is the
derivative of the smoothed loss of the whole net in its parameter slot
(vit_net_lossGrad_smoothedCE). The vitin_*dp* artifacts, the book's
vitin_emadp256x4wxdropeps0000001smxerfbf16 among them, run that program on R replicas: each
gradient node feeds allReduceMeanF, the all-reduced gradient feeds the AdamW tail, and the
updated weight feeds the EMA shadow. This file composes those three with the one-replica theorem.
ViT has no BatchNorm, so no op couples examples and the composition needs no synchronised statistic: the loss is a sum over examples, a shard's forward is the shard of the forward, and the mean of the replicas' losses is the single-device loss at the global batch.
vitNetB_shard— replicar's forward at its shard of the images and of the stochastic-depth masks is its shard of the batch-R·Nforward.smoothedBatchLossDiv_shard— the batch-R·Nloss at divisorR·Bis the mean over replicas of the batch-Nlosses at divisorB, each at its shard of the logits and targets.HasGradAt.at_allReduceMeanF— the collective, at any slot: if each replica's node is the gradient of that replica's loss, the all-reduced node is the gradient of the mean loss.vit_dp_node_lossGrad— the three composed, for every parameter slot: whenever each replica's gradient node is the gradient of its own shard's loss in the slot (whichvit_net_lossGrad_smoothedCEproves for each of the 200 nodes), the all-reduced node is the gradient, in that slot, of the smoothed loss of the single-device net at the global batchR·N.vit_dp_node_eq_globalreads it as an equality with the batch-R·Nnode.vit_dp_step— the optimizer and the shadow: the emitted AdamW triple at the all-reduced node isadamWStepat that global-batch gradient, and the EMA node over the updated weight isemaNext d e θ', withθ'that step's parameter.
A slot is any slot : Vec P → ViTTieWeights nC, the weights with one tensor replaced by a
vector: fun θ => { w with bcls := θ }, fun θ => { w with b3 := { w.b3 with Wq := Mat.unflatten θ } },
and so on. ViTNetLossTiedGB's 200 conjuncts are exactly hypotheses hg of this shape, one per
slot; vit_dp_bcls instantiates one end to end, the classifier bias, with no hypothesis left but
0 < ε and targets that sum to one.
What is not claimed #
That every replica holds the same weights, that the shards are cut from one host batch and that
replica_groups names all R devices are the driver's (DataParallel.Basic,
Verified.Train) and the *-dp-check gates'. The learning-rate, bias-correction and decay
scalars are runtime inputs, binders here. The theorems are over ℝ.
The collective at a slot. If each replica's node is the gradient at θ of that replica's
loss G r, the all-reduced node is the gradient at θ of the mean of the R losses.
dpMeanGrad_eq_grad_meanLoss in HasGradAt's vocabulary, at the SHlo node.
The smoothed batch loss shards. The loss is a sum over examples divided by a constant, so
the batch-R·N loss at divisor R·B is the mean over the R replicas of the batch-N losses
at divisor B, each at its shard of the targets and logits.
Replica r's shard of the stochastic-depth masks: example n of replica r reads the
global batch's entries at (r, n), at both sites of every block. none stays none.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Replica r's forward is its shard of the batch-R·N forward, at its shard of the images
and of the masks. No ViT op couples examples, so every stage commutes with the cut
(batchShard_batchMap, batchShard_batchMapIdx).
Replica r's loss, as a function of the weights: the smoothed loss at divisor B of the
batch-N net on replica r's shard of the images, masks and targets.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The single-device loss at the global batch, as a function of the weights: the smoothed loss
at divisor R·B of the batch-R·N net.
Equations
- Proofs.ViTTieGB.vitGlobalLoss gf R N ε α B sd img t w = Proofs.smoothedBatchLossDiv (R * N) nC α (↑R * B) t (Proofs.ViTTieGB.vitNetB gf (R * N) ε w sd img)
Instances For
The mean of the replicas' losses is the single-device loss at the global batch, as
functions of the weights: smoothedBatchLossDiv_shard at vitNetB_shard.
Every all-reduced ViT-Tiny gradient node is the single-device gradient at the global batch
R·N. For any parameter slot slot and any family g of per-replica gradient nodes: if
each replica's node is the gradient, in the slot, of that replica's loss on its own shard
(vit_net_lossGrad_smoothedCE at batch N, for each of the 200 nodes), then the node the
render all-reduces them into is the gradient, in the slot, of the smoothed loss of the
batch-R·N net at divisor R·B.
…and so equals the batch-R·N node. Whatever single-device node gG is the gradient of
the global-batch loss in the slot (vit_net_lossGrad_smoothedCE at batch R·N), the
all-reduced node denotes the same vector.
One slot end to end: the classifier bias. The all-reduced biasGradB node of the
data-parallel render, each replica's at the smoothed-loss cotangent it emits on its own
shard, is the gradient in bcls of the single-device loss at the global batch R·N. The
hypothesis of vit_dp_node_lossGrad is discharged by vit_net_lossGrad_smoothedCE at each
replica; the other 199 slots discharge the same way from their conjuncts.
The data-parallel ViT-Tiny update, through the optimizer and the EMA shadow. Under
vit_dp_node_lossGrad's hypothesis, for any slot:
- the AdamW triple the render emits at the all-reduced node is
adamWStepat∇, the gradient in the slot of the single-device loss at the global batchR·N(adamW_at_allReduceMeanF;wd = 0is thewxrenders' excluded slots); - the EMA node over the updated weight is
emaNext d e θ', i.e.d·e + (1 − d)·θ', withθ'that step's parameter (emaShadow_faithful).
The rate, the bias corrections and the decay are runtime scalars, binders here.