The sync-BN data-parallel kit — sharding, homogeneity and the collectives, per op kind #

Every net's sync-BN twin (T2 forward, T3 step) says that R replicas at batch N, each running the render's replica program over its shard, compute the single-device step at the global batch R·N. The per-op facts that argument is assembled from are stated here once:

what	names	namespace
the `c·h·w ↔ c·(h·w)` index cast and sharding through it; non-BN nodes commute with the batch cut; the BN sync site	`castIdx`, `laAssoc`, `batchShard_castIdx`, `den_batchOp_shard`, `den_relu_shard`, `den_addVB_shard`, `bnSyncSiteLA`	`StableHLO`
homogeneity — each cotangent step and gradient node is linear in its cotangent	`*_smul`	`ResNet34SyncTieB`, `MBConvSyncTieB`
sharding — each input-VJP is per example, so it commutes with the batch cut; sync-BN's backward is the shard of the global one	`*_shard`, `bnSyncInB_shard`	`ResNet34SyncTieB`, `MBConvSyncTieB`
P4 — the replica mean of a weight-gradient node is `1/R` of the global node	`den_allReduceMeanF_*_shard`	both (and `DataParallelSync` for conv W / BN β)
per-parameter DP ties at `R ×` the shards of a global cotangent	`Sync`, `Sync_of_scaled`	both
the divisor: a replica's loss cotangent is `R ×` its shard of the global one	`replicaLossCot_eq`	`ResNet34SyncTieB`

The namespaces are the nets that first needed each piece; the names are cited by every twin.

source

theorem Proofs.StableHLO.batchShard_castIdx {R N a b : ℕ} (hab : a = b) (X : Vec (R * N * a)) (r : Fin R) :

batchShard R N b (fun (i : Fin (R * N * b)) => X (Fin.cast ⋯ i)) r = fun (i : Fin (N * b)) => batchShard R N a X r (Fin.cast ⋯ i)

Sharding commutes with relabelling the per-example index. The batch axis is outside the per-example one, so relabelling within an example and cutting the batch do not interact.

source

theorem Proofs.StableHLO.den_batchOp_shard {R N a b : ℕ} (op : BatchableOp a b) (e : Fin R → SHlo (N * a)) (X : Vec (R * N * a)) (he : ∀ (r : Fin R), den (e r) = batchShard R N a X r) (r : Fin R) :

den (SHlo.batchOp op (e r)) = batchShard R N b (batchMap (R * N) (denOp op) X) r

A per-example node on every replica denotes the shard of the same node on the global batch.

source

theorem Proofs.StableHLO.den_relu_shard {R N n : ℕ} (e : Fin R → SHlo (N * n)) (X : Vec (R * N * n)) (he : ∀ (r : Fin R), den (e r) = batchShard R N n X r) (r : Fin R) :

den (SHlo.batchOp BatchableOp.relu (e r)) = batchShard R N n (relu (R * N * n) X) r

…stated for relu at the whole-batch relu, the form the committed forward is written in.

source

theorem Proofs.StableHLO.den_addVB_shard {R N n : ℕ} (a b : Fin R → SHlo (N * n)) (A B : Vec (R * N * n)) (ha : ∀ (r : Fin R), den (a r) = batchShard R N n A r) (hb : ∀ (r : Fin R), den (b r) = batchShard R N n B r) (r : Fin R) :

den ((a r).addVB (b r)) = batchShard R N n (fun (j : Fin (R * N * n)) => A j + B j) r

The residual fan-in on every replica is the shard of the global one.

source

def Proofs.StableHLO.bnSyncSiteLA (gN bN es t t' : String) (ds ds' : List ℕ) (R : ℕ) (hR : 0 < R) {N oc h w : ℕ} (ε : ℝ) (γ β : Vec oc) (x : Fin R → SHlo (N * (oc * h * w))) (r : Fin R) :

SHlo (N * (oc * h * w))

One sync-BN forward site, at the network index, on replica r — bnFwdSite's replicas > 1 branch: bnSyncF of this replica's operand, reading syncStats over all R replicas' operands (the mean collective t, then Chan's variance collective t'), with the mul_assoc relabelling on the way in and out.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

theorem Proofs.StableHLO.den_bnSyncSiteLA (gN bN es t t' : String) (ds ds' : List ℕ) (R : ℕ) (hR : 0 < R) {N oc h w : ℕ} (hm : N * (h * w) ≠ 0) (ε : ℝ) (γ β : Vec oc) (x : Fin R → SHlo (N * (oc * h * w))) (X : Vec (R * N * (oc * h * w))) (hx : ∀ (r : Fin R), den (x r) = batchShard R N (oc * h * w) X r) (r : Fin R) :

den (bnSyncSiteLA gN bN es t t' ds ds' R hR ε γ β x r) = batchShard R N (oc * h * w) (bnBatchLA (R * N) oc h w ε γ β X) r

⭐⭐ The sync-BN site on replica r is shard r of the global-batch BatchNorm. den_bnSyncF_allReduce (P1 on the graph), carried across the mul_assoc seam: the right-hand side is bnBatchLA — what bnBatchF denotes — at N := R·N.

source

theorem Proofs.StableHLO.nhw_ne_zero {N h w : ℕ} (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) :

N * (h * w) ≠ 0

The reduction width a BatchNorm site needs nonzero, from the three positive dimensions.

source

theorem Proofs.ResNet34SyncTieB.reluMaskB_smul (n : ℕ) (pre : Vec n) :

IsHomog (ResNet34TieB.reluMaskB n pre)

source

theorem Proofs.ResNet34SyncTieB.bn_grad_input_smul (n : ℕ) (ε γ : ℝ) (x : Vec n) :

IsHomog (bn_grad_input n ε γ x)

The three-term BatchNorm input-gradient is linear in dy — both its reductions are.

source

theorem Proofs.ResNet34SyncTieB.bnPerChannel_grad_input_smul (oc m : ℕ) (ε : ℝ) (γ : Vec oc) (x : Vec (oc * m)) :

IsHomog (bnPerChannel_grad_input oc m ε γ x)

source

theorem Proofs.ResNet34SyncTieB.bnBatchTensor4_grad_input_smul (N oc h w : ℕ) (ε : ℝ) (γ : Vec oc) (x : Vec (N * (oc * (h * w)))) :

IsHomog (bnBatchTensor4_grad_input N oc h w ε γ x)

source

theorem Proofs.ResNet34SyncTieB.bnInB_smul (N oc h w : ℕ) (ε : ℝ) (γ : Vec oc) (x : Vec (N * (oc * h * w))) :

IsHomog (ResNet34TieB.bnInB N oc h w ε γ x)

source

theorem Proofs.ResNet34SyncTieB.batchMap_smul {N a b : ℕ} (f : Vec a → Vec b) (hf : IsHomog f) :

IsHomog (StableHLO.batchMap N f)

source

theorem Proofs.ResNet34SyncTieB.batchMapAux_smul {N t a b : ℕ} (f : Vec t → Vec a → Vec b) (hf : ∀ (x : Vec t), IsHomog (f x)) (aux : Vec (N * t)) :

IsHomog (StableHLO.batchMapAux N f aux)

source

theorem Proofs.ResNet34SyncTieB.cInB_smul (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) :

IsHomog (EnetTiePoC.cInB N W b)

source

theorem Proofs.ResNet34SyncTieB.cStridedInB_smul (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) :

IsHomog (ResNet34TieB.cStridedInB N W b)

source

theorem Proofs.ResNet34SyncTieB.maxPool3s2BackFlat_smul (c h w : ℕ) (xv : Vec (c * (2 * h) * (2 * w))) :

IsHomog (StableHLO.maxPool3s2BackFlat c h w xv)

The 3×3/s2 pool's select_and_scatter is linear in the cotangent it scatters.

source

theorem Proofs.ResNet34SyncTieB.mpInB_smul (N c h w : ℕ) (x : Vec (N * (c * (2 * h) * (2 * w)))) :

IsHomog (ResNet34TieB.mpInB N c h w x)

source

theorem Proofs.ResNet34SyncTieB.batchSlice_smul {N a : ℕ} (X : Vec (N * a)) (s : ℝ) (n : Fin N) :

StableHLO.batchSlice N a (fun (i : Fin (N * a)) => s * X i) n = fun (i : Fin a) => s * StableHLO.batchSlice N a X n i

source

theorem Proofs.ResNet34SyncTieB.convWeightGradB_smul {N ic oc h w kH kW : ℕ} (xN cotN : String) (b : Vec oc) (x : Vec (N * (ic * h * w))) (W : Kernel4 oc ic kH kW) (cot : Vec (N * (oc * h * w))) (s : ℝ) (idx : Fin (oc * ic * kH * kW)) :

StableHLO.den (StableHLO.SHlo.convWeightGradB xN b x W (StableHLO.SHlo.operand cotN fun (i : Fin (N * (oc * h * w))) => s * cot i)) idx = s * StableHLO.den (StableHLO.SHlo.convWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx

source

theorem Proofs.ResNet34SyncTieB.convStridedWeightGradB_smul {N ic oc h w kH kW : ℕ} (xN cotN : String) (b : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cot : Vec (N * (oc * h * w))) (s : ℝ) (idx : Fin (oc * ic * kH * kW)) :

StableHLO.den (StableHLO.SHlo.convStridedWeightGradB xN b x W (StableHLO.SHlo.operand cotN fun (i : Fin (N * (oc * h * w))) => s * cot i)) idx = s * StableHLO.den (StableHLO.SHlo.convStridedWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx

source

theorem Proofs.ResNet34SyncTieB.bnGammaGradB_smul {N oc h w : ℕ} (vN epsStr cotN : String) (ε : ℝ) (v cot : Vec (N * (oc * (h * w)))) (s : ℝ) (k : Fin oc) :

StableHLO.den (StableHLO.SHlo.bnGammaGradB vN epsStr ε v (StableHLO.SHlo.operand cotN fun (i : Fin (N * (oc * (h * w)))) => s * cot i)) k = s * StableHLO.den (StableHLO.SHlo.bnGammaGradB vN epsStr ε v (StableHLO.SHlo.operand cotN cot)) k

source

theorem Proofs.ResNet34SyncTieB.bnBetaGradB_smul {N oc h w : ℕ} (cotN : String) (cot : Vec (N * (oc * (h * w)))) (s : ℝ) (k : Fin oc) :

StableHLO.den (StableHLO.SHlo.operand cotN fun (i : Fin (N * (oc * (h * w)))) => s * cot i).bnBetaGradB k = s * StableHLO.den (StableHLO.SHlo.operand cotN cot).bnBetaGradB k

source

theorem Proofs.ResNet34SyncTieB.denseWeightGradB_smul {N a c : ℕ} (xN cotN : String) (x : Vec (N * a)) (cot : Vec (N * c)) (s : ℝ) (idx : Fin (a * c)) :

StableHLO.den (StableHLO.SHlo.denseWeightGradB xN x (StableHLO.SHlo.operand cotN fun (i : Fin (N * c)) => s * cot i)) idx = s * StableHLO.den (StableHLO.SHlo.denseWeightGradB xN x (StableHLO.SHlo.operand cotN cot)) idx

source

theorem Proofs.ResNet34SyncTieB.denseBiasGradB_smul {N c : ℕ} (cotN : String) (cot : Vec (N * c)) (s : ℝ) (j : Fin c) :

StableHLO.den (StableHLO.SHlo.operand cotN fun (i : Fin (N * c)) => s * cot i).denseBiasGradB j = s * StableHLO.den (StableHLO.SHlo.operand cotN cot).denseBiasGradB j

source

theorem Proofs.ResNet34SyncTieB.cInB_shard {R N ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (DY : Vec (R * N * (oc * h * w))) (r : Fin R) :

EnetTiePoC.cInB N W b (batchShard R N (oc * h * w) DY r) = batchShard R N (ic * h * w) (EnetTiePoC.cInB (R * N) W b DY) r

source

theorem Proofs.ResNet34SyncTieB.cStridedInB_shard {R N ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (DY : Vec (R * N * (oc * h * w))) (r : Fin R) :

ResNet34TieB.cStridedInB N W b (batchShard R N (oc * h * w) DY r) = batchShard R N (ic * (2 * h) * (2 * w)) (ResNet34TieB.cStridedInB (R * N) W b DY) r

source

theorem Proofs.ResNet34SyncTieB.mpInB_shard {R N : ℕ} (c h w : ℕ) (X : Vec (R * N * (c * (2 * h) * (2 * w)))) (DY : Vec (R * N * (c * h * w))) (r : Fin R) :

ResNet34TieB.mpInB N c h w (batchShard R N (c * (2 * h) * (2 * w)) X r) (batchShard R N (c * h * w) DY r) = batchShard R N (c * (2 * h) * (2 * w)) (ResNet34TieB.mpInB (R * N) c h w X DY) r

source

noncomputable def Proofs.ResNet34SyncTieB.bnSyncInB (R : ℕ) (hR : 0 < R) (N oc h w : ℕ) (ε : ℝ) (γ : Vec oc) (xs dys : Fin R → Vec (N * (oc * h * w))) (r : Fin R) :

Vec (N * (oc * h * w))

Replica r's sync-BN input cotangent, as bnBackSite's replicas > 1 branch computes it, in the network layout: this replica's [μ ‖ σ² ‖ mean(γ·dy) ‖ mean(x̂·γ·dy)] (bnSyncDyStatsB, reading the forward's syncStats) all-reduced, then bnSyncBack. The replica peer of ResNet34TieB.bnInB, and like it written as the den of the emitted nodes over .operand leaves.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

theorem Proofs.ResNet34SyncTieB.reassocB_shard {R N oc h w : ℕ} (X : Vec (R * N * (oc * h * w))) (r : Fin R) :

EnetTiePoC.reassocB N oc h w (batchShard R N (oc * h * w) X r) = batchShard R N (oc * (h * w)) (EnetTiePoC.reassocB (R * N) oc h w X) r

reassocB of a shard is the shard of the reassocB.

source

theorem Proofs.ResNet34SyncTieB.bnSyncInB_shard (R : ℕ) (hR : 0 < R) (N oc h w : ℕ) (hm : N * (h * w) ≠ 0) (ε : ℝ) (γ : Vec oc) (xs dys : Fin R → Vec (N * (oc * h * w))) (X DY : Vec (R * N * (oc * h * w))) (hxs : ∀ (r : Fin R), xs r = batchShard R N (oc * h * w) X r) (hdys : ∀ (r : Fin R), dys r = batchShard R N (oc * h * w) DY r) (r : Fin R) :

bnSyncInB R hR N oc h w ε γ xs dys r = batchShard R N (oc * h * w) (ResNet34TieB.bnInB (R * N) oc h w ε γ X DY) r

⭐⭐ The sync-BN backward on replica r is shard r of the global-batch BN backward — den_bnSyncBack_allReduce (P2 on the graph) at the network index. The right-hand side is bnInB, the single-device chain's BN link, at N := R·N.

source

theorem Proofs.ResNet34SyncTieB.den_allReduceMeanF_convStridedWeightGradB_shard {N ic oc h w kH kW : ℕ} (R : ℕ) (hR : 0 < R) (t xN cotN : String) (ds : List ℕ) (b : Vec oc) (W : Kernel4 oc ic kH kW) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (DY : Vec (R * N * (oc * h * w))) (dy : Fin R → StableHLO.SHlo (N * (oc * h * w))) (hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N (oc * h * w) DY r) (idx : Fin (oc * ic * kH * kW)) :

StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) => StableHLO.SHlo.convStridedWeightGradB xN b (batchShard R N (ic * (2 * h) * (2 * w)) X r) W (dy r)) idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.convStridedWeightGradB xN b X W (StableHLO.SHlo.operand cotN DY)) idx

P4 at the STRIDED conv weight — den_allReduceMeanF_convWeightGradB_shard's strided peer.

source

theorem Proofs.ResNet34SyncTieB.den_allReduceMeanF_denseWeightGradB_shard {N a c : ℕ} (R : ℕ) (hR : 0 < R) (t xN cotN : String) (ds : List ℕ) (A : Vec (R * N * a)) (DY : Vec (R * N * c)) (dy : Fin R → StableHLO.SHlo (N * c)) (hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N c DY r) (idx : Fin (a * c)) :

StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) => StableHLO.SHlo.denseWeightGradB xN (batchShard R N a A r) (dy r)) idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.denseWeightGradB xN A (StableHLO.SHlo.operand cotN DY)) idx

P4 at the dense weight — the head's Σ_n outer product, split by replica.

source

theorem Proofs.ResNet34SyncTieB.den_allReduceMeanF_denseBiasGradB_shard {N c : ℕ} (R : ℕ) (hR : 0 < R) (t cotN : String) (ds : List ℕ) (DY : Vec (R * N * c)) (dy : Fin R → StableHLO.SHlo (N * c)) (hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N c DY r) (j : Fin c) :

StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) => (dy r).denseBiasGradB) j = 1 / ↑R * StableHLO.den (StableHLO.SHlo.operand cotN DY).denseBiasGradB j

P4 at the dense bias — Σ_n cot, split by replica.

source

def Proofs.ResNet34SyncTieB.ConvWSync (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc kH kW : ℕ} (t xN cotN : String) (b : Vec oc) (X : Vec (R * N * (ic * h * w))) (W : Kernel4 oc ic kH kW) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) :

Prop

One conv weight, DP-tied: the mean over replicas of each replica's weight-gradient node (at its shard of the layer input and its own cotangent) IS the single-device node at the global batch. Tags are the render's: the collective is named for the parameter.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

def Proofs.ResNet34SyncTieB.ConvStridedWSync (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc kH kW : ℕ} (t xN cotN : String) (b : Vec oc) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) :

Prop

The strided conv weight, DP-tied.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

def Proofs.ResNet34SyncTieB.BnSync (R : ℕ) (hR : 0 < R) (N oc h w : ℕ) (tg tb vN epsStr cotN : String) (ε : ℝ) (V : Vec (R * N * (oc * h * w))) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) :

Prop

One BatchNorm's γ and β, DP-tied. Each replica's γ node is bnSyncGammaGradB, reading the forward's all-reduced [μ ‖ σ²] (syncStats over the replicas' pre-BN activations, tagged {tg}mu / {tg}var as bnFwdSite tags them), so its x̂ is the global batch's; the β node reads no statistic. The right-hand sides are the single-device bnGammaGradB / bnBetaGradB at N := R·N — BnPairTiedB's nodes.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

def Proofs.ResNet34SyncTieB.DenseSync (R : ℕ) (hR : 0 < R) (N : ℕ) {a c : ℕ} (tW tb xN cotN : String) (A : Vec (R * N * a)) (cots : Fin R → Vec (N * c)) (COT : Vec (R * N * c)) :

Prop

The dense weight and bias, DP-tied.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

theorem Proofs.ResNet34SyncTieB.inv_mul_R (R : ℕ) (hR : 0 < R) (v : ℝ) :

1 / ↑R * (↑R * v) = v

(1/R)·(R·v) = v — the collective's mean against the divisor's R.

source

theorem Proofs.ResNet34SyncTieB.convWSync_of_scaled (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc kH kW : ℕ} (t xN cotN : String) (b : Vec oc) (X : Vec (R * N * (ic * h * w))) (W : Kernel4 oc ic kH kW) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) (hc : ∀ (r : Fin R), cots r = batchShard R N (oc * h * w) (fun (i : Fin (R * N * (oc * h * w))) => ↑R * COT i) r) :

ConvWSync R hR N h w t xN cotN b X W cots COT

⭐ A replica family at R × the shards of COT gives the DP tie — the collective's 1/R (§4) cancels the R homogeneity (§1) carries. The same three lines for every kind below.

source

theorem Proofs.ResNet34SyncTieB.convStridedWSync_of_scaled (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc kH kW : ℕ} (t xN cotN : String) (b : Vec oc) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) (hc : ∀ (r : Fin R), cots r = batchShard R N (oc * h * w) (fun (i : Fin (R * N * (oc * h * w))) => ↑R * COT i) r) :

ConvStridedWSync R hR N h w t xN cotN b X W cots COT

source

theorem Proofs.ResNet34SyncTieB.bnSync_of_scaled (R : ℕ) (hR : 0 < R) (N oc h w : ℕ) (hm : N * (h * w) ≠ 0) (tg tb vN epsStr cotN : String) (ε : ℝ) (V : Vec (R * N * (oc * h * w))) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) (hc : ∀ (r : Fin R), cots r = batchShard R N (oc * h * w) (fun (i : Fin (R * N * (oc * h * w))) => ↑R * COT i) r) :

BnSync R hR N oc h w tg tb vN epsStr cotN ε V cots COT

source

theorem Proofs.ResNet34SyncTieB.denseSync_of_scaled (R : ℕ) (hR : 0 < R) (N : ℕ) {a c : ℕ} (tW tb xN cotN : String) (A : Vec (R * N * a)) (cots : Fin R → Vec (N * c)) (COT : Vec (R * N * c)) (hc : ∀ (r : Fin R), cots r = batchShard R N c (fun (i : Fin (R * N * c)) => ↑R * COT i) r) :

DenseSync R hR N tW tb xN cotN A cots COT

source

theorem Proofs.ResNet34SyncTieB.rowB_shard {R N K : ℕ} (Z : Vec (R * N * K)) (r : Fin R) :

ResNet34TieB.rowB N K (batchShard R N K Z r) = batchShard R N (1 * K) (ResNet34TieB.rowB (R * N) K Z) r

rowB of a shard is the shard of the rowB — a relabelling inside each example.

source

theorem Proofs.ResNet34SyncTieB.unrowB_shard {R N K : ℕ} (Z : Vec (R * N * (1 * K))) (r : Fin R) :

ResNet34TieB.unrowB N K (batchShard R N (1 * K) Z r) = batchShard R N K (ResNet34TieB.unrowB (R * N) K Z) r

source

theorem Proofs.ResNet34SyncTieB.replicaLossCot_eq (R N nCls : ℕ) (hR : 0 < R) (α B : ℝ) (aStr negAK bStr logN ohN : String) (Z : Vec (R * N * nCls)) (T : Vec (R * N * (1 * nCls))) (r : Fin R) :

ResNet34TieB.unrowB N nCls (StableHLO.den (smoothedLossCotGraph N nCls α B aStr negAK bStr logN ohN (ResNet34TieB.rowB N nCls (batchShard R N nCls Z r)) (batchShard R N (1 * nCls) T r))) = batchShard R N nCls (fun (i : Fin (R * N * nCls)) => ↑R * ResNet34TieB.unrowB (R * N) nCls (StableHLO.den (smoothedLossCotGraph (R * N) nCls α (↑R * B) aStr negAK bStr logN ohN (ResNet34TieB.rowB (R * N) nCls Z) T)) i) r

⭐ The divisor step. Replica r divides its smoothed-CE cotangent by B — the render's divConstB at the per-replica batch — and the single-device step at the global batch divides by R·B. At the replica's shard of the logits and targets, the replica's cotangent is R × its shard of the global one. Nothing else about the loss differs: softmax, the label-smoothing shift and the target are per example.

source

theorem Proofs.MBConvSyncTieB.dInB_smul (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) :

IsHomog (EnetTiePoC.dInB N W b)

source

theorem Proofs.MBConvSyncTieB.dStridedInB_smul (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) :

IsHomog (EnetTiePoC.dStridedInB N W b)

source

theorem Proofs.MBConvSyncTieB.gapInB_smul (N c h w : ℕ) :

IsHomog (EnetTiePoC.gapInB N c h w)

source

theorem Proofs.MBConvSyncTieB.rowDenseBackFlat_smul (N a c : ℕ) (W : Mat a c) :

IsHomog (StableHLO.rowDenseBackFlat N a c W)

The row-wise input-VJP dX = W·dy (the classifier's, and the SE excite dense's) is linear in dy.

source

theorem Proofs.MBConvSyncTieB.hasVJP3_backward_smul {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} (hf : HasVJP3 f) (x : Tensor3 c₁ h₁ w₁) (a : ℝ) (dy : Tensor3 c₂ h₂ w₂) :

(hf.backward x fun (i₁ : Fin c₂) (i₂ : Fin h₂) (i₃ : Fin w₂) => a * dy i₁ i₂ i₃) = fun (j₁ : Fin c₁) (j₂ : Fin h₁) (j₃ : Fin w₁) => a * hf.backward x dy j₁ j₂ j₃

A HasVJP3 backward is linear in its cotangent — HasVJP.backward_smul's three-axis peer, read off HasVJP3.correct. The stride-1 depthwise weight gradient is stated through one.

source

theorem Proofs.MBConvSyncTieB.depthwiseWeightGradB_smul {N c h w kH kW : ℕ} (xN cotN : String) (b : Vec c) (x : Vec (N * (c * h * w))) (W : DepthwiseKernel c kH kW) (cot : Vec (N * (c * h * w))) (s : ℝ) (idx : Fin (c * kH * kW)) :

StableHLO.den (StableHLO.SHlo.depthwiseWeightGradB xN b x W (StableHLO.SHlo.operand cotN fun (i : Fin (N * (c * h * w))) => s * cot i)) idx = s * StableHLO.den (StableHLO.SHlo.depthwiseWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx

source

theorem Proofs.MBConvSyncTieB.depthwiseStridedWeightGradB_smul {N c h w kH kW : ℕ} (xN cotN : String) (b : Vec c) (x : Vec (N * (c * (2 * h) * (2 * w)))) (W : DepthwiseKernel c kH kW) (cot : Vec (N * (c * h * w))) (s : ℝ) (idx : Fin (c * kH * kW)) :

StableHLO.den (StableHLO.SHlo.depthwiseStridedWeightGradB xN b x W (StableHLO.SHlo.operand cotN fun (i : Fin (N * (c * h * w))) => s * cot i)) idx = s * StableHLO.den (StableHLO.SHlo.depthwiseStridedWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx

source

theorem Proofs.MBConvSyncTieB.convStridedXlaWeightGradB_smul {N ic oc h w kH kW : ℕ} (xN cotN : String) (b : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cot : Vec (N * (oc * h * w))) (s : ℝ) (idx : Fin (oc * ic * kH * kW)) :

StableHLO.den (StableHLO.SHlo.convStridedXlaWeightGradB xN b x W (StableHLO.SHlo.operand cotN fun (i : Fin (N * (oc * h * w))) => s * cot i)) idx = s * StableHLO.den (StableHLO.SHlo.convStridedXlaWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx

source

theorem Proofs.MBConvSyncTieB.dInB_shard {R N c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (DY : Vec (R * N * (c * h * w))) (r : Fin R) :

EnetTiePoC.dInB N W b (batchShard R N (c * h * w) DY r) = batchShard R N (c * h * w) (EnetTiePoC.dInB (R * N) W b DY) r

source

theorem Proofs.MBConvSyncTieB.dStridedInB_shard {R N c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (DY : Vec (R * N * (c * h * w))) (r : Fin R) :

EnetTiePoC.dStridedInB N W b (batchShard R N (c * h * w) DY r) = batchShard R N (c * (2 * h) * (2 * w)) (EnetTiePoC.dStridedInB (R * N) W b DY) r

source

theorem Proofs.MBConvSyncTieB.gapInB_shard {R N : ℕ} (c h w : ℕ) (DY : Vec (R * N * c)) (r : Fin R) :

EnetTiePoC.gapInB N c h w (batchShard R N c DY r) = batchShard R N (c * h * w) (EnetTiePoC.gapInB (R * N) c h w DY) r

source

theorem Proofs.MBConvSyncTieB.rowDenseBackFlat_shard {R N : ℕ} (a c : ℕ) (W : Mat a c) (DY : Vec (R * N * c)) (r : Fin R) :

StableHLO.rowDenseBackFlat N a c W (batchShard R N c DY r) = batchShard R N a (StableHLO.rowDenseBackFlat (R * N) a c W DY) r

The row-wise input-VJP is batchMap of W·, so it shards like every per-example lift.

source

theorem Proofs.MBConvSyncTieB.den_allReduceMeanF_depthwiseWeightGradB_shard {N c h w kH kW : ℕ} (R : ℕ) (hR : 0 < R) (t xN cotN : String) (ds : List ℕ) (b : Vec c) (W : DepthwiseKernel c kH kW) (X DY : Vec (R * N * (c * h * w))) (dy : Fin R → StableHLO.SHlo (N * (c * h * w))) (hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N (c * h * w) DY r) (idx : Fin (c * kH * kW)) :

StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) => StableHLO.SHlo.depthwiseWeightGradB xN b (batchShard R N (c * h * w) X r) W (dy r)) idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.depthwiseWeightGradB xN b X W (StableHLO.SHlo.operand cotN DY)) idx

P4 at the depthwise weight — each replica's Σ_n over its own examples, averaged, is 1/R of the global batch's Σ_n.

source

theorem Proofs.MBConvSyncTieB.den_allReduceMeanF_depthwiseStridedWeightGradB_shard {N c h w kH kW : ℕ} (R : ℕ) (hR : 0 < R) (t xN cotN : String) (ds : List ℕ) (b : Vec c) (W : DepthwiseKernel c kH kW) (X : Vec (R * N * (c * (2 * h) * (2 * w)))) (DY : Vec (R * N * (c * h * w))) (dy : Fin R → StableHLO.SHlo (N * (c * h * w))) (hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N (c * h * w) DY r) (idx : Fin (c * kH * kW)) :

StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) => StableHLO.SHlo.depthwiseStridedWeightGradB xN b (batchShard R N (c * (2 * h) * (2 * w)) X r) W (dy r)) idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.depthwiseStridedWeightGradB xN b X W (StableHLO.SHlo.operand cotN DY)) idx

P4 at the strided depthwise weight.

source

theorem Proofs.MBConvSyncTieB.den_allReduceMeanF_convStridedXlaWeightGradB_shard {N ic oc h w kH kW : ℕ} (R : ℕ) (hR : 0 < R) (t xN cotN : String) (ds : List ℕ) (b : Vec oc) (W : Kernel4 oc ic kH kW) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (DY : Vec (R * N * (oc * h * w))) (dy : Fin R → StableHLO.SHlo (N * (oc * h * w))) (hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N (oc * h * w) DY r) (idx : Fin (oc * ic * kH * kW)) :

StableHLO.den (StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) => StableHLO.SHlo.convStridedXlaWeightGradB xN b (batchShard R N (ic * (2 * h) * (2 * w)) X r) W (dy r)) idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.convStridedXlaWeightGradB xN b X W (StableHLO.SHlo.operand cotN DY)) idx

P4 at the XLA-SAME strided conv weight (the stem) — den_allReduceMeanF_convWeightGradB_shard's peer. Only the certificate differs from the symmetric strided one; the batch split is the same.

source

def Proofs.MBConvSyncTieB.DepthwiseWSync (R : ℕ) (hR : 0 < R) (N h w : ℕ) {c kH kW : ℕ} (t xN cotN : String) (b : Vec c) (X : Vec (R * N * (c * h * w))) (W : DepthwiseKernel c kH kW) (cots : Fin R → Vec (N * (c * h * w))) (COT : Vec (R * N * (c * h * w))) :

Prop

A depthwise weight, DP-tied — the collective over the [c, 1, kH, kW] kernel the render all-reduces, against the single-device node at the global batch.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

def Proofs.MBConvSyncTieB.DepthwiseStridedWSync (R : ℕ) (hR : 0 < R) (N h w : ℕ) {c kH kW : ℕ} (t xN cotN : String) (b : Vec c) (X : Vec (R * N * (c * (2 * h) * (2 * w)))) (W : DepthwiseKernel c kH kW) (cots : Fin R → Vec (N * (c * h * w))) (COT : Vec (R * N * (c * h * w))) :

Prop

The strided depthwise weight, DP-tied.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

def Proofs.MBConvSyncTieB.ConvStridedXlaWSync (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc kH kW : ℕ} (t xN cotN : String) (b : Vec oc) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) :

Prop

The stem's XLA-SAME strided conv weight, DP-tied. Tags are the render's: the collective is named for the parameter, over its [oc, ic, kH, kW] shape.

Equations

One or more equations did not get rendered due to their size.

Instances For

source

theorem Proofs.MBConvSyncTieB.depthwiseWSync_of_scaled (R : ℕ) (hR : 0 < R) (N h w : ℕ) {c kH kW : ℕ} (t xN cotN : String) (b : Vec c) (X : Vec (R * N * (c * h * w))) (W : DepthwiseKernel c kH kW) (cots : Fin R → Vec (N * (c * h * w))) (COT : Vec (R * N * (c * h * w))) (hc : ∀ (r : Fin R), cots r = batchShard R N (c * h * w) (fun (i : Fin (R * N * (c * h * w))) => ↑R * COT i) r) :

DepthwiseWSync R hR N h w t xN cotN b X W cots COT

source

theorem Proofs.MBConvSyncTieB.depthwiseStridedWSync_of_scaled (R : ℕ) (hR : 0 < R) (N h w : ℕ) {c kH kW : ℕ} (t xN cotN : String) (b : Vec c) (X : Vec (R * N * (c * (2 * h) * (2 * w)))) (W : DepthwiseKernel c kH kW) (cots : Fin R → Vec (N * (c * h * w))) (COT : Vec (R * N * (c * h * w))) (hc : ∀ (r : Fin R), cots r = batchShard R N (c * h * w) (fun (i : Fin (R * N * (c * h * w))) => ↑R * COT i) r) :

DepthwiseStridedWSync R hR N h w t xN cotN b X W cots COT

source

theorem Proofs.MBConvSyncTieB.convStridedXlaWSync_of_scaled (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc kH kW : ℕ} (t xN cotN : String) (b : Vec oc) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cots : Fin R → Vec (N * (oc * h * w))) (COT : Vec (R * N * (oc * h * w))) (hc : ∀ (r : Fin R), cots r = batchShard R N (oc * h * w) (fun (i : Fin (R * N * (oc * h * w))) => ↑R * COT i) r) :

ConvStridedXlaWSync R hR N h w t xN cotN b X W cots COT

Documentation

LeanMlir.Proofs.Foundation.DataParallelSyncKit

The sync-BN data-parallel kit — sharding, homogeneity and the collectives, per op kind #