MobileNetV2's data-parallel forward at SYNCHRONISED BatchNorm — replica r IS shard r #
MobileNetV2FullB.lean (T2) says the typed batch-BN graph denotes mobilenetv2ForwardB_full N w
on one device. MobileNetV2RenderB's data-parallel step normalises with the GLOBAL batch's
statistics at replicas > 1: every one of the 52 BatchNorm sites is the sync-BN composition —
this replica's mean all-reduced, then Chan's σ²_r + (μ_r − μ)² all-reduced, packed, then
bnSyncF. This file is T2's data-parallel twin: that forward graph, stated as a family over the
R replicas, denotes on replica r exactly batchShard r of the single-device forward at the
global batch R·N.
den (mobilenetv2FwdGraphSync_full R hR N epsStr w e r)
= batchShard R N nCls (mobilenetv2ForwardB_full (R * N) w X) r
given that each replica's input is its shard of one global batch X. ⭐ The spec does not
move: the right-hand side is the committed mobilenetv2ForwardB_full, at N := R·N.
How it is proved #
As ResNet34SyncB.lean proves ResNet-34's: by induction on the chain, one block at a time, with
the shard hypothesis ∀ r, den (e r) = batchShard R N _ X r carried from block to block.
- every conv, depthwise, GAP and dense node is a per-example lift (
den_batchOp_shard), and so is the XLA-SAMEstrided conv and depthwise — the padding lives inside the per-example map; - relu6 is pointwise, so it commutes with sharding (
den_relu6_shard, the peer ofden_relu_shard); the identity skip isden_addVB_shard; - every BatchNorm site is
bnSyncSiteLA, whose shard lemmaden_bnSyncSiteLAis P1 on the graph read at the network index.
The graph reuses ResNet34SyncB's site verbatim — it is net-agnostic — so the only new lemma is
the relu6 one; the rest of this file is MobileNetV2's six block shapes and their chain.
Names #
Parameter names are MobileNetV2FullB's (%b{k}{e,d,p}{W,g,bt}, %sW, %hW, …, the
convBias := false zero-bias operands %zb{c}). A BatchNorm site with γ %b{k}dg gathers its
statistics as %arsum / %armean of b{k}dgmu and b{k}dgvar, each over a [c] vector — the
γ name without %, then mu / var.
What is NOT claimed here #
⚠ The backward and the parameter collectives are the T3 half (MobileNetV2SyncStepTieB.lean).
⚠ That the R replicas' inputs ARE the shards of one batch is the driver's, as in
DataParallelSync.lean. ⚠ The lowerer's all_reduce is trusted as every other op's lowering is.
relu6 on every replica is the shard of the global relu6 — the clamp reads one cell, so
cutting the batch before or after it is the same. The MobileNetV2 peer of den_relu_shard,
stated at the whole-batch relu6 the committed forward is written in.
Stem at sync-BN, over the replica family: 3x3/s2 XLA-SAME conv → sync-BN → relu6.
Equations
- One or more equations did not get rendered due to their size.
Instances For
t = 1 bottleneck (b1) at sync-BN, over the replica family: depthwise → sync-BN → relu6 →
project 1x1 → sync-BN. Two sync sites.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-1 no-skip bottleneck (b11, b17) at sync-BN, over the replica family: expand → depthwise → project, a sync site after each, relu6 after the first two.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-1 skip bottleneck at sync-BN: the body plus the addVB identity skip, the block input
shared between both arms on every replica.
Equations
- Proofs.StableHLO.mnv2ResidGraphSync pfx epsStr R hR N h w p e r = (Proofs.StableHLO.mnv2ExpOnlyGraphSync pfx epsStr R hR N h w p e r).addVB (e r)
Instances For
Stride-2 downsampling bottleneck at sync-BN, over the replica family: expand at 2h x 2w, the
XLA-SAME strided depthwise, project at h x w; three sync sites.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Head at sync-BN, over the replica family: 1x1 conv → sync-BN → relu6 → GAP → dense.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The sync-BN data-parallel MobileNetV2 forward graph, over the replica family. T2's
mobilenetv2FwdGraphB_full with every BatchNorm a bnSyncSiteLA over all R replicas; block
prefixes, parameter names and collective tags are the render's.
Equations
- One or more equations did not get rendered due to their size.
Instances For
⭐⭐ T2 at synchronised BatchNorm: replica r's forward IS shard r of the global-batch
forward. Given that the replicas' inputs are the shards of one batch X of R·N examples,
the sync-BN graph on replica r denotes batchShard r of mobilenetv2ForwardB_full (R * N) w X — the committed batch-BN forward, at the global batch. One block lemma per stage, the
shard hypothesis threaded from each into the next.