Documentation

LeanMlir.Proofs.Foundation.DataParallelNode

Data parallelism, piece 2: the collective as an AST node #

DataParallel.lean (piece 1) is the ℝ-level half: dpMean, the replica mean of gradients, IS the gradient of the mean of the per-replica losses, and that mean is the global-batch loss for a net with no batch coupling and provably is not for a batch-BN net. What it could not say was anything about the ARTIFACT, because the collective was ViTRender.emitGradAllReduce — emitted text outside the SHlo AST, a declared trusted carve-out that every train-step tie in the repo disclaimed in its own header.

Since 2026-09-07 the collective is SHlo.allReduceMeanF: R graphs of one skeleton (the same program on R replicas, each with its own values — SPMD), reduced by all_reduce(add) and divided by R. Its den is (1/R) Σ_r den (g r), its skel reads replica 0, its emit is the old text verbatim (allReduceMeanText), and the parser round-trip has its case (StableHLOParse.parseStack). Every batched render now calls prettyAllReduceMean where it called the text function, and every committed *dp* artifact re-rendered byte-identically.

This file is what the node BUYS, stated once:

What is NOT claimed #

⚠ Piece 3 is untouched: that the R graphs' values are the replica slices of ONE host batch, that every replica starts from the same parameters and that replica_groups names all R devices are the driver's (VerifiedTrain.lean, ffi/pjrt_ffi.c) and the *-dp-check gates'. ⚠ The per-replica gradient node's operand values differ per replica by construction; nothing here says what they are. ⚠ The lowerer's all_reduce is trusted exactly as every other op's lowering is.

theorem Proofs.den_allReduceMeanF_eq_dpMean {n : } (R : ) (hR : 0 < R) (t : String) (ds : List ) (g : Fin RStableHLO.SHlo n) :

The node denotes dpMean of its operands. Definitional: den's arm is piece 1's formula.

theorem Proofs.skel_allReduceMeanF_of_spmd {n : } (R : ) (hR : 0 < R) (t : String) (ds : List ) (g : Fin RStableHLO.SHlo n) (hsp : ∀ (r : Fin R), StableHLO.skel (g r) = StableHLO.skel (g 0, hR)) (r : Fin R) :

SPMD: the R graphs of one skeleton. Under the hypothesis that every replica's operand graph has replica 0's skeleton — free in every render, whose operand family is .operand grad at every r, and the content of "the same program on R replicas" — the node's skeleton is each replica's, so pretty prints ONE program and den sums R of them.

theorem Proofs.den_allReduceMeanF_eq_lossGrad_meanLoss {P : } (R : ) (hR : 0 < R) (t : String) (ds : List ) (g : Fin RStableHLO.SHlo P) (L : Fin RVec P) (θ : Vec P) (hdiff : ∀ (r : Fin R), LossDifferentiableAt (L r) θ) (hg : ∀ (r : Fin R), StableHLO.den (g r) = lossGrad (L r) θ) :

Piece 1 composed with the node. If each replica's gradient node denotes the gradient of that replica's loss, the all-reduced node denotes the gradient of the MEAN of the per-replica losses — the function a data-parallel run minimises (dpMeanGrad_eq_grad_meanLoss).

theorem Proofs.den_allReduceMeanF_convWeightGradB {N ic oc h w kH kW : } (R : ) (hR : 0 < R) (t : String) (ds : List ) (xN cotN : String) (b : Vec oc) (x : Fin RVec (N * (ic * h * w))) (W : Kernel4 oc ic kH kW) (cot : Fin RVec (N * (oc * h * w))) (idx : Fin (oc * ic * kH * kW)) :
StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) => StableHLO.SHlo.convWeightGradB xN b (x r) W (StableHLO.SHlo.operand cotN (cot r))) idx = 1 / R * r : Fin R, n : Fin N, j : Fin (oc * h * w), pdiv (fun (v' : Vec (oc * ic * kH * kW)) => (conv2d (Kernel4.unflatten v') b (Tensor3.unflatten (StableHLO.batchSlice N (ic * h * w) (x r) n))).flatten) W.flatten idx j * StableHLO.batchSlice N (oc * h * w) (cot r) n j

4b's fold composed with the node. The all-reduced conv weight-gradient node denotes the replica mean of the certified batched Σ_n gradients, each at its own replica's activations and cotangent. One op kind; every other *GradB composes identically (Finset.sum_congr over the replicas, then its own fold lemma).

theorem Proofs.adamW_at_allReduceMeanF {n : } (θN mN vN b1N ob1N b2N ob2N bc1N bc2N lrN epsN wdN : String) (ds : List ) (β₁ β₂ ε lr wd bc₁ bc₂ : ) (θ m v : Vec n) (R : ) (hR : 0 < R) (t : String) (ds' : List ) (g : Fin RStableHLO.SHlo n) :
(StableHLO.den (StableHLO.SHlo.adamWParamF θN mN vN b1N ob1N b2N ob2N bc1N bc2N lrN epsN wdN ds β₁ β₂ ε lr wd bc₁ bc₂ θ m v (StableHLO.SHlo.allReduceMeanF R hR t ds' g)), StableHLO.den (StableHLO.SHlo.adamMNextF mN b1N ob1N ds β₁ m (StableHLO.SHlo.allReduceMeanF R hR t ds' g)), StableHLO.den (StableHLO.SHlo.adamVNextF vN b2N ob2N ds β₂ v (StableHLO.SHlo.allReduceMeanF R hR t ds' g))) = adamWStep β₁ β₂ ε lr wd bc₁ bc₂ θ m v (dpMean fun (r : Fin R) => StableHLO.den (g r))

⭐⭐ The optimizer tail at the all-reduced node — §4d piece 2's target statement: den (tail (allReduceMeanF R g)) is adamWStep at dpMean of the per-replica gradient nodes. adamW_triple_faithful is ∀ e, so this is that theorem at e := allReduceMeanF … and the node's den; the same one line closes it for every other certified tail (mom_pair_faithful, rmsProp_triple_faithful, lamb_triple_faithful).