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

ResNet50StepTieB (T3) threads a loss cotangent g down the batch-BN bottleneck backward chain on ONE device and ties every parameter gradient node to the certified gradient. 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 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 T3. 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 (ResNet34SyncTieB.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 (DataParallelSync'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.resnet50FwdGraphSync_full_shard, the forward half. ⚠ No stochastic depth (drop-path): the chain is the drop-free one, as T3's is. ⚠ The f32 nodes — the bf16 conv twins are not this statement. ⚠ 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.