EfficientNet-B0's data-parallel forward at SYNCHRONISED BatchNorm — replica r IS shard r #

EfficientNetFullB0.lean says the typed batch-BN graph denotes efficientnetForwardBFull N w on one device. At replicas > 1 the data-parallel render normalises every one of B0's 49 BatchNorms with the GLOBAL batch's statistics — the sync-BN composition bnFwdSite emits: this replica's mean all-reduced, then Chan's σ²_r + (μ_r − μ)² all-reduced, packed, then bnSyncF. This file is its 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 (efficientnetFwdGraphSyncFull R hR N epsStr w bf16 e r)
  = batchShard R N nCls (efficientnetForwardBFull (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 efficientnetForwardBFull, at N := R·N, and it does not move with the precision either: bf16 selects the conv and depthwise kinds at every site (.convAt / .convStridedXlaAt / .depthwiseAt / .depthwiseStridedXlaAt, as the single-device graph), bf16 := true being the kinds the *bf16 DP renders emit — efficientnetin_emarmsdp64dropdowxeps0001bf16, the run the book reports, among them (its drop sites are outside this graph, as they are at f32); 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 #

ResNet-34's recipe (ResNet34SyncB.lean), block by block, the shard hypothesis ∀ r, den (e r) = batchShard R N _ X r carried from each block into the next:

• every conv, depthwise, strided depthwise, squeeze-excite, GAP and dense node is a per-example lift and commutes with sharding (den_batchOp_shard) — the SE block included: its squeeze (GAP), both dense layers and the gate multiply all act on one example at a time, so it is one batchMap and needs nothing new;
• swish is pointwise (den_swishF_shard), and so is the residual addV (den_addV_shard);
• every BatchNorm is bnSyncSiteLA, whose shard lemma den_bnSyncSiteLA is P1 on the graph.

The collectives are tagged as bnFwdSite tags them: a BN site whose γ is %{p}eg gathers {p}egmu and {p}egvar, each over a [c] statistic. den ignores every string, so the tags fix which graph is stated, not what it denotes.

What is NOT claimed here #

The backward and the parameter collectives are the step half (EfficientNetSyncStepTieG.lean). That the R replicas' inputs are the shards of one batch is the driver's. The emitted artifacts' stochastic-depth and classifier-dropout variants are efficientnetFwdGraphSyncFullDrop (replica r's dropPathB at its own per-example scale on each of the nine skip-carrying blocks, its dropoutB at its own mask before the dense), and efficientnetFwdGraphSyncFullDrop_shard says it denotes shard r of efficientnetForwardBFullDrop at the global batch and the global masks the replicas' are the shards of. The lowerer's all_reduce is trusted as every other op's lowering is.