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

MobileNetV4StepTieB.lean threads a loss cotangent g 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 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.

The single-device tie 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 BnSync / DenseSync statements and their *_of_scaled closers (the conv and depthwise weight collectives are Foundation.GradNodesBAt's ConvWSyncAt bf16 / ConvStridedWSyncAt bf16 / DepthwiseWSyncAt bf16 / DepthwiseStridedWSyncAt bf16, on the renderers' precision switch), and the homogeneity of bnInB, cInB, cStridedInB and the head. The MBConv pieces come from DataParallel.SyncKit — the depthwise and SYMMETRIC strided-depthwise input-VJPs and weight collectives, the GAP backward and the row-wise dense input-VJP. The head's two 1×1 relabellings are per-example reindexes, so they commute with the batch cut (mnv4To11_shard, mnv4From11_shard) and with scaling.

1. Sharding — each replica's backward chain, handed its shard of a global cotangent, computes the shard of the global chain. The relu mask is 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 MobileNetV4StepTieB is, so widths stay variables; the capstone instantiates at the 21 concrete rows. MNv4's activation is relu everywhere (timm), so the masks are reluMaskB — MobileNetV2's relu6MaskB does not appear; the BN-only pre-DW has no mask at all.

The index seam #

ℝ-level, as MobileNetV4StepTieB 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 DataParallel.Sync's P2 (the sync-BN input-VJP shard identity) 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 the BnPairTiedB nodes MobileNetV4StepTieB uses 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 blocks × 12 (ten stride-1 rows and the three 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.

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.mnv4FwdGraphSyncFull_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. Every conv and depthwise weight collective below is stated on the renderers' switch (Foundation.GradNodesBAt): the replicas' nodes and the batch-R·N node are the f32 kind at bf16 := false and the bf16 kind at true — and the capstone takes bf16, so it covers the f32 DP artifacts and the *bf16 ones' step (mnv4in_emaaccdp8x128wxdowd005bf16, mnv4in_accdp8x128wxdropdowd01bf16) 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 stride-1 body (ExtraDW / ConvNeXt-like / FFN), each one line from the previous link's. The two ifs are MobileNetV4StepTieB's table dispatch; they split on both sides at once.

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

The stem, the fused stage and the head.