Documentation

LeanMlir.Proofs.Foundation.DataParallel.Node

Data parallelism: the collective as an AST node #

DataParallel.Basic 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 per-replica batch-BN net. This file puts the collective in the SHlo AST so a tie can say something about the artifact.

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 allReduceMeanText, and the parser (StableHLOParse.parseStack) has its case. The batched renders print it through prettyAllReduceMean.

This file is what the node BUYS, stated once:

What is NOT claimed #

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 (Verified.Train, 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 (DataParallel.Sync relates them under sync-BN). The lowerer's all_reduce is trusted exactly as every other op's lowering is.

source
theorem Proofs.den_allReduceMeanF_eq_dpMean {n : ℕ} (R : ℕ) (hR : 0 < R) (t : String) (ds : List ℕ) (g : Fin R → StableHLO.SHlo n) :
StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds g) = dpMean fun (r : Fin R) => StableHLO.den (g r)

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

source
theorem Proofs.skel_allReduceMeanF_of_spmd {n : ℕ} (R : ℕ) (hR : 0 < R) (t : String) (ds : List ℕ) (g : Fin R → StableHLO.SHlo n) (hsp : ∀ (r : Fin R), StableHLO.skel (g r) = StableHLO.skel (g ⟨0, hR⟩)) (r : Fin R) :
StableHLO.skel (StableHLO.SHlo.allReduceMeanF R hR t ds g) = StableHLO.Raw.allReduceMean R t ds (StableHLO.skel (g 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.

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

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).

source
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 R → Vec (N * (ic * h * w))) (W : Kernel4 oc ic kH kW) (cot : Fin R → Vec (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

The conv weight-gradient 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).

source
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 R → StableHLO.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: 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).