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

ResNet50StepTieB threads a loss cotangent g down the batch-BN bottleneck backward chain on ONE device and ties every parameter gradient node to the certified per-op gradient at the cotangent that chain delivers. This is its data-parallel twin, for the render ResNet50RenderB emits at replicas > 1: R replicas at batch N, every BatchNorm synchronised (bnFwdSite / bnBackSite / bnGammaSite), every parameter gradient all-reduced by its mean. The capstone r50_net_syncTiedB says that, for every one of the 161 parameters,

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, loss cotangent G

whenever each replica's loss cotangent is R × its shard of the global one (∀ r, gs r = batchShard r (R • G)) — the gradient node r50_net_tiedB at N := R·N ties to the certified per-op gradient. The right-hand side is the existing single-device chain at N := R·N, so the spec has not moved.

The loss cotangent is a binder, as it is in r50_net_tiedB. ResNet-50 ships both losses, so the capstone takes the scaled-shard relation between the replicas' cotangents and the global one as its hypothesis, and two corollaries discharge it, one per loss:

• r50_net_syncTiedB_smoothedCE — the label-smoothed chain, replicas dividing by B and the global step by R·B (SyncKit.replicaLossCot_eq);
• r50_net_syncTiedB_bce — BCE-with-logits at the COMMITTED divisors, N·K on a replica and (R·N)·K on the global step (replicaBceLossCot_eq, the BCE peer, proved here).

Four steps (ResNet-34's, with the bottleneck's links) #

1. Sharding — each replica's backward chain, handed its shard of a global cotangent, computes the shard of the global chain. Every non-BN link (relu mask, the 1×1 and 3×3 conv input-VJPs, the strided 3×3 and strided 1×1 input-VJPs, the stem pool's scatter, the head) is a per-example map; the BN link is bnSyncInB, whose shard lemma bnSyncInB_shard is P2 on the graph. In the strided block bn₁ runs at the INPUT grid 2h × 2w (v1.5: the stride is on the 3×3), so that site's statistics reduce over N·(2h)·(2w) per replica.
2. The collectives — the mean over replicas of each replica's gradient node is 1/R of the global node at the global cotangent (DataParallel.Sync's P4 lemmas and ResNet-34's strided and dense ones). The γ node is the sync one, bnSyncGammaGradB, reading the forward's all-reduced statistics.
3. Homogeneity — the single-device chain is linear in its cotangent (*_smul below for the three bottleneck forms, ResNet-34's for the links they are built from).
4. The divisor — the R the hypothesis carries cancels the collective's 1/R.

The index seam #

The chain runs at N·(c·h·w), the BN nodes at N·(c·(h·w)); the replica BN link bnSyncInB is the den of the emitted nodes over .operand leaves at reassocB (ResNet-34's), and reassocB_shard commutes the relabelling with sharding.

What is NOT claimed #

The replicas' saved forward activations enter as the shards of the single-device forward's (batchShard r (r50Pre_k (R*N) q w X)); that the sync forward graph computes exactly those is StableHLO.resnet50FwdGraphSyncFull_shard, the forward half. No stochastic depth (drop-path): the chain is the drop-free one, as r50_net_tiedB's is. The nodes are at either precision: bf16 selects the conv weight collectives' kind on both sides (ConvWSyncAt bf16 / ConvStridedWSyncAt bf16, Foundation.GradNodesBAt), true being the bf16 DP artifacts' (resnet50in160_lambaccdp4x128wxclipbcebf16, …), read at the identity rounding (Bf16Erasure, as ResNet34SyncStepTieB). The render has no conv-bias gradient ops, so there are none here. That the replicas' inputs are the shards of one batch is the driver's. The gradient accumulator and the optimizers run after the collective, so this is per micro-step. The lowerer's all_reduce is trusted as every other op's lowering is.