MobileNetV4-Conv-M's data-parallel step at SYNCHRONISED BatchNorm IS the single-device step at R·N #

MobileNetV4StepTieB.lean (T3) threads a loss cotangent g 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 MobileNetV4RenderB emits at replicas > 1: R replicas at batch N, every one of the 77 BatchNorms synchronised, every parameter gradient all-reduced by its mean. The capstone mnv4_net_syncTiedB says that, for every parameter the render emits,

mean over the R replicas of replica r's gradient node, replica cotangent gs r
  = the single-device gradient node at the global batch R·N, global cotangent G

whenever each replica's loss cotangent is R × its shard of the global one — the gradient node mnv4_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.

⚠ T3 binds the loss cotangent, so this twin does too #

mnv4_net_tiedB takes the logits' cotangent g as a binder (mnv4_lossCot_is_smoothedCE_grad instantiates it). The twin therefore takes the global cotangent G and the replica family gs as binders, with the scaled-shard hypothesis ∀ r, gs r = batchShard R N nCls (fun i => R * G i) r — exactly the invariant ResNet-34's and MobileNetV2's twins carry down their chains, and there discharge from replicaLossCot_eq. mnv4_net_syncTiedB_smoothedCE discharges it the same way for the label-smoothed softmax chain the artifacts emit: replicas dividing their loss by B, the single-device step by R·B.

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

ResNet34SyncStepTieB.lean is the template, and everything net-agnostic is imported from it: the replica BN link bnSyncInB and its shard lemma, the ConvWSync / ConvStridedWSync / BnSync / DenseSync statements and their *_of_scaled closers, and the homogeneity of bnInB, cInB, cStridedInB and the head. The MBConv pieces come from MBConvSyncTieB.lean — the depthwise and SYMMETRIC strided-depthwise input-VJPs and weight collectives, and the XLA-SAME stem collective. Swish's backward is a certified VJP's .backward, so its homogeneity is HasVJP.backward_smul and its sharding is definitional.

1. Sharding — each replica's backward chain, handed its shard of a global cotangent, computes the shard of the global chain. The relu mask and swish's backward are pointwise; the conv, depthwise, strided conv and strided depthwise input-VJPs are per-example maps; every BN link is bnSyncInB. The table's k = 0 dispatch (if s.postDWk = 0, if s.preDWk = 0) is the same if on both sides, so it splits once.
2. The collectives — the mean over replicas of each replica's gradient node is 1/R of the global node at the global cotangent. MobileNetV4 needs no kind the kit does not have.
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 — the hypothesis on gs; its R cancels the collective's 1/R at every parameter.

⭐ Everything is GENERIC IN THE ROW (s : UibSpec), as T3 is, so widths stay variables; the capstone instantiates at the 21 concrete rows. MNv4's activation is relu (the fused stage's is swish), so the masks are reluMaskB — MobileNetV2's relu6MaskB does not appear.

The index seam #

ℝ-level, as T3 is: the replica BN link bnSyncInB is the den of the emitted nodes (bnSyncDyStatsB → all-reduce → bnSyncBack) over .operand leaves at reassocB, and bnSyncInB_shard carries P2 across the N·(c·h·w) / N·(c·(h·w)) seam; BnSync's γ and β collectives read reassocB of the pre-BN activation and of the cotangent, exactly as T3's BnPairTiedB nodes do.

What the DP render emits, and what is tied #

MobileNetV4RenderB emits 233 parameter gradients — stem 3 (sW, sg, sbt), fused 6 (f0cW f0cg f0cbt f0pW f0pg f0pbt), thirteen ExtraDW-profile blocks × 12 (ten stride-1 rows and the three pre-strided rows 1, 3, 11: u{p}{q,e,d,p}{W,g,bt}), four ConvNeXt-like × 9 (no d), four FFN × 6 (no q, no d), head 8 (h1W h1g h1bt hW hg hbt Wd bd) — and the capstone ties all 233. ⚠ There are no conv-bias gradients to exclude: the render has no convBias flag, binds every bias slot to %zb{c}, and emits none. ⚠ Conv-M has no post-strided row, so the render's post-strided backward is never emitted for this table and T3 has no chain for it; neither does this file.

What is NOT claimed #

⚠ The replicas' saved forward activations enter as the shards of the single-device forward's (batchShard r (mnv4Blk{k} (R*N) w X)); that the sync forward graph computes exactly those is StableHLO.mnv4FwdGraphSync_full_shard, the forward half. ⚠ That the replicas' inputs are the shards of one batch is the driver's. ⚠ The statement is at the f32 nodes: the *bf16 artifact's bf16 conv twins are outside it, as for every other net. ⚠ The lowerer's all_reduce is trusted as every other op's lowering is.

The stride-1 body (ExtraDW / ConvNeXt-like / FFN), each one line from the previous link's. The two ifs are T3's table dispatch; they split on both sides at once.

The pre-strided block (rows 1, 3, 11).

The stem, the fused stage (swish, no mask) and the two-conv head.