ResNet-34's data-parallel forward at SYNCHRONISED BatchNorm — replica r IS shard r #

ResNet34FullB.lean says the typed batch-BN graph denotes resnet34ForwardBFull N w on one device (resnet34FwdGraphBFull_faithful). ResNet34RenderB renders the data-parallel step differently: at replicas > 1 every BatchNorm site is bnFwdSite's sync-BN composition — this replica's mean all-reduced, then Chan's σ²_r + (μ_r − μ)² all-reduced, packed, then bnSyncF — so the statistics each replica normalises by are the GLOBAL batch's. This file is that theorem's data-parallel twin: that forward graph, stated as a family over the R replicas, denotes on replica r exactly batchShard r of the single-device forward at the global batch R·N.

den (resnet34FwdGraphSyncFull R hR N epsStr w bf16 e r)
  = batchShard R N nCls (resnet34ForwardBFull (R * N) w X) r

given that each replica's input is its shard of one global batch X. The spec does not move: the right-hand side is the committed resnet34ForwardBFull, at N := R·N, and it does not move with the precision either: bf16 selects the conv kinds at every site (.convAt / .convStridedAt, as ResNet34FullB), bf16 := true being the graph of resnet34in_momdp64bf16, the run the book reports; each conv's shard step erases the switch through Bf16Erasure (Bf16Fold.denOp_convAt_id) and the rest of the proof is the f32 one.

How it is proved #

By induction on the chain, one block at a time, exactly as resnet34FwdGraphBFull_faithful is — with the shard hypothesis ∀ r, den (e r) = batchShard R N _ X r as the invariant carried from block to block:

• every conv, relu, pool, GAP and dense node is a per-example lift, and sharding commutes with it (batchShard_batchMap, and relu pointwise) — den_batchOp_shard, den_relu_shard;
• every BatchNorm site is bnSyncSiteLA, whose shard lemma den_bnSyncSiteLA is DataParallelSync.den_bnSyncF_allReduce (P1 on the graph) read at the network's index.

The index seam #

The conv/relu chain runs at the left-assoc index N·(c·h·w); bnSyncF and its statistics nodes at N·(c·(h·w)), the index bnBatchTensor4 is stated at. bnBatchF hides that seam inside its den (bnBatchLA is bnBatchTensor4 conjugated by the mul_assoc relabelling). The sync site cannot, because its statistics subgraph is shared across replicas, so it carries the relabelling as castIdx on the AST value — h ▸ e, the move ConvNeXtRenderB's reassocB already makes. It changes no emitted text: skel never sees an index, and the render writes the same SSA name at both types (bnFwdSite's zbn operand beside its zin one).

What is NOT claimed here #

The backward and the parameter collectives are ResNet34SyncStepTieB.lean's. That the R replicas' inputs ARE the shards of one batch is the driver's, as in DataParallel.Sync. The lowerer's all_reduce is trusted as every other op's lowering is.