EfficientNet-B0's data-parallel step at SYNCHRONISED BatchNorm IS the single-device step at R·N #

EfficientNetStepTieG.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. This is its data-parallel twin, for the render EfficientNetRender.Basic emits at replicas > 1: R replicas at batch N, every one of the 49 BatchNorms synchronised (bnFwdSite / bnBackSite / bnGammaSite), every parameter gradient all-reduced by its mean. The capstone efficientnet_net_syncTiedG says that, for every parameter,

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 EnetTieG.efficientnet_net_tiedG at N := R·N ties to the certified gradient. The spec it is stated against has not moved: the right-hand side is the single-device tie's chain at N := R·N — its forward prefixes, loss cotangent and block .backwards verbatim, and its in-block cotangents as named definitions (xCotEc, tCotDn, …) that unfold to enetExpTiedG's lets, so each right-hand node is the single-device tie's node by rfl.

Five steps #

ResNet-34's four (ResNet34SyncStepTieB.lean), and one B0 needs because of how its single-device tie is written.

0. The block VJP is the chain (§ 0). The single-device tie threads each block's input cotangent as a certified VJP's .backward; a replica computes its own by the explicit chain, and only the explicit chain can be sharded. xCotIn_eq_vjp and its four peers say the two agree — the Batched.BackLinks stage graphs, read at an .operand leaf, with bnBatchLABack_faithful turning the one non-rfl link into bnBackB.
1. Sharding (§§ 2, 5). Every non-BN link is per-example — conv, depthwise and strided depthwise input-VJPs, swish and sigmoid masks, and the squeeze-excite backward: the gate cotangent gateCotB (seReduceB's den) and the fused input-VJP seInB (seBackBatched's) both read one example at a time, so they shard like ResNet-34's pool backward. The BN link is bnSyncInB, P2 on the graph, read through bnInB_eq_bnBackB onto the single-device tie's bnBackB.
2. The collectives. ResNet-34's conv, dense and BatchNorm ones, plus the depthwise, strided-depthwise and XLA-SAME stem conv weights from DataParallel.SyncKit, which also holds the depthwise, GAP and dense links of steps 1 and 3 that MobileNetV2 shares.
3. Homogeneity (§§ 1, 4). Every link is linear in its cotangent; most are a certified VJP's .backward, so HasVJP.backward_smul is the whole proof.
4. The divisor. Replica r's loss cotangent is R × its shard of the global one (replicaLossCot_eq), and steps 1–3 carry that R down to cancel each collective's 1/R.

What is covered, and what is NOT claimed #

The DP render runs convBias := false, so it emits 213 parameter collectives — stem 3, b1 10, fifteen MBConv6 blocks × 13, head 5 — and all 213 are tied. The single-device tie carries 49 more conjuncts, one per conv bias (a bnBetaGradB at the conv-output cotangent), for the convBias := true census; the DP artifacts do not emit them and they are not tied here.

Note: the replicas' saved forward activations enter as the shards of the single-device forward's; that the sync forward graph computes exactly those is EfficientNetSyncB's efficientnetFwdGraphSyncFull_shard, the forward half. Note: that the replicas' inputs are the shards of one batch is the driver's. Note: the lowerer's all_reduce is trusted as every other op's lowering is.

The training masks are optional sites on the statement, as in EfficientNetStepTieG: the capstone takes the GLOBAL masks sd : Option (Fin 9 → Vec (R * N)) (stochastic depth, nine blocks) and cd : Option (Vec ((R * N) * 1280)) (classifier dropout), and replica r runs its chain at shard r of each — exampleShard of a per-example scale, batchShard of the per-element mask — which is what the DP render's per-replica %dp<i> / %do inputs are. none is the drop-free DP artifacts' statement exactly as it was; some reaches efficientnetin_emarmsdp64dropdowxeps0001bf16's chain. Each site is linear in the cotangent and commutes with the batch cut (dropPathOpt_smul / _shard, dropoutOpt_smul / _shard, Foundation.DropSites), so the five steps are unchanged: a residual block's replica chain feeds its branch the family dropPathOptFam and the single-device one dropPathOpt … dy, the invariant passing the site by dropPathOptFam_scaled; the head's dense collective reads the dropped activation on every replica. The forward half with the sites is StableHLO.efficientnetFwdGraphSyncFullDrop_shard.

Precision is a flag on the statement. The ImageNet run the book reports trains from efficientnetin_emarmsdp64dropdowxeps0001bf16, 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 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.

The tail's forward activations, from the depthwise conv output dc — enetExpTiedG's dn, dr, s, e1, z, e2, se, pc, definitionally.

The tail's backward chain from the block-output cotangent dy — enetExpTiedG's cotPbn … cotDc, definitionally.

The stride-1 expand front (b3, b5, …, and the widenings b9 / b16).

The strided front (b2, b4, b6, b12): expand at the input grid 2h×2w.

The MBConv1 front (b1): the depthwise runs on the block input.

The stem (enetStemTiedG's chain) and the head (enetHeadTiedG's).

Each block's input cotangent IS its certified VJP #

The single-device tie threads the block-output cotangents by the block VJPs' .backward; the replicas compute theirs by the explicit chain. These say the two agree, so the sharding argument (which needs the explicit chain) lands on the single-device tie's own .backward terms. HasVJP.backward_unique swaps the bundle-level witness for the unfolded one, a vjpComp of the stage VJPs, which vjpComp_backward takes apart into the stage backwards.

The tail, on the replicas #

The fronts, on the replicas #

The scaled-shard invariant across each block #

Replicas at R × the shards of the single-device block-output cotangent DY hand the next block up R × the shards of the single-device tie's own .backward — sharding, then the explicit chain IS the VJP (§ 0), then HasVJP.backward_smul.