MobileNetV2's data-parallel step at synchronised BatchNorm is the single-device step at R·N #

MobileNetV2StepTieB.lean threads the label-smoothed loss cotangent down the batch-BN backward chain on one device and ties every parameter gradient node to the certified gradient at that chain's cotangent. This is its data-parallel twin, for the render MobileNetV2RenderB emits at replicas > 1: R replicas at batch N, every one of the 52 BatchNorms synchronised, every parameter gradient all-reduced by its mean. The capstone mnv2_net_syncTiedB says that, for every parameter the render emits,

mean over the R replicas of replica r's gradient node, loss divided by B
  = the single-device gradient node at the global batch R·N, loss divided by R·B

— the gradient node mnv2_net_tiedB at N := R·N ties to the certified gradient. The right-hand side is the existing single-device chain at N := R·N, so the spec has not moved.

The same four steps as ResNet-34's twin #

ResNet34SyncStepTieB.lean is the template, and everything net-agnostic is imported from it rather than restated: the replica BN link bnSyncInB and its shard lemma, the ConvWSync / BnSync / DenseSync statements and their *_of_scaled closers, the divisor step replicaLossCot_eq, and the homogeneity of bnInB and cInB. The inverted-residual pieces it shares with EfficientNet-B0 — the depthwise, GAP and dense input-VJPs, the depthwise and XLA-SAME stem weight collectives — come from DataParallel.SyncKit.

1. Sharding — each replica's backward chain, handed its shard of a global cotangent, computes the shard of the global chain. The relu6 mask is pointwise, so it shards by rfl; the conv, depthwise, XLA-SAME strided depthwise, GAP and dense input-VJPs are per-example maps; every BN link is bnSyncInB.
2. The collectives — the mean over replicas of each replica's gradient node is 1/R of the global node at the global cotangent. MobileNetV2 adds three kinds the ResNet-34 kit does not have: the XLA-SAME strided conv weight (the stem), the depthwise weight and its XLA-SAME strided peer (§4).
3. Homogeneity — the single-device chain and its gradient nodes are linear in the loss cotangent (§1): R × the cotangent gives R × every node.
4. The divisor — replica r's loss cotangent is R × its shard of the global one, and that R cancels the collective's 1/R at every parameter.

What the DP render emits, and what is tied #

At the committed convBias := false, MobileNetV2RenderB emits 158 parameter gradients — stem 3 (sW, sg, sbt), b1 6, sixteen blocks × 9 (eW eg ebt dW dg dbt pW pg pbt), head 3 (hW, hg, hbt), dense 2 (Wd, bd) — and the capstone ties all 158. The 52 conv, depthwise and project BIAS nodes the single-device tie also states are not emitted at convBias := false and are not tied here (mnv2_net_tiedB keeps them for the flag).

What is NOT claimed #

The replicas' saved forward activations enter as the shards of the single-device forward's (batchShard r (mnv2PreB{k} (R*N) w X)); that the sync forward graph computes exactly those is StableHLO.mobilenetv2FwdGraphSyncFull_shard, the forward half. That the replicas' inputs are the shards of one batch is the driver's. The lowerer's all_reduce is trusted as every other op's lowering is.

Precision is a flag on the statement. The ImageNet run the book reports trains from mobilenetv2in_rmsdp64wxdols0eps0001bf16, whose conv and depthwise weight gradients are the bf16 kinds. Every conv and depthwise weight collective below is stated on the renderers' switch — ConvWSyncAt bf16 / ConvStridedXlaWSyncAt bf16 / DepthwiseWSyncAt bf16 / DepthwiseStridedXlaWSyncAt bf16 (Foundation.GradNodesBAt): the replicas' nodes and the batch-R·N node are the f32 kind at false and the bf16 kind at true — and the capstone takes bf16, so it covers the f32 DP artifacts (mobilenetv2_adamdp, mobilenetv2in_adamdp64, mobilenetv2in_rmsdp64, mobilenetv2in_rmsdp128) and the *bf16 ones' step read over ℝ at the identity rounding (Bf16Erasure), where each bf16 collective is 1/R of the batch-R·N bf16 node exactly as at f32. At a real rounding the bf16 weight-gradient collective differs from the single-device node by the per-replica rounding of each partial sum, which DataParallel.SyncBf16 states and this file does not model, as the f32 tie does not model f32 rounding.

Classifier dropout is an optional site on the statement, as in MobileNetV2StepTieB: the capstone takes cd : Option (Vec ((R * N) * 1280)), the GLOBAL mask, and replica r runs its chain at batchShard r of it — which is what the DP render's per-replica %do input is. none is the drop-free DP artifacts' statement exactly as it was; some M reaches mobilenetv2in_rmsdp64wxdols0eps0001bf16's head: the site is linear in the cotangent (dropoutOpt_smul) and commutes with the batch cut (dropoutOpt_shard, Foundation.DropSites), so the four steps above are unchanged, and the dense collective reads the dropped activation on every replica. The forward half with the site is StableHLO.mobilenetv2FwdGraphSyncFullDo_shard.

The block cotangents, each one line from the previous link's.