Documentation

LeanMlir.Proofs.Codegen.SyncBnSites

The three BatchNorm emit sites every batch-BN renderer shares — batch BN or SYNCHRONISED BN #

Used by ResNet34RenderB, ResNet50RenderB (through ResNet34RenderB), MobileNetV2RenderB, MobileNetV4RenderB and EfficientNetRender.Basic, so every batch-BN renderer emits the same sync-BN composition. Each site is one BatchNorm node at sync = false and at sync = true the sync-BN subgraph whose den Foundation/DataParallel/Sync.lean states (the forward's two collectives, the backward's one, the γ gradient reading the forward's statistics). sync alone picks the branch; replicas only sizes the collectives. Callers pass sync := replicas > 1 (ResNet34RenderB adds forceSync, sync-BN on one device), so sync = false at several replicas would be per-replica batch BN.

Tags. The forward site's tag is its γ parameter's name without %, and its collectives are named from it ({tag}mu, {tag}var). The backward site's tag is the collective's full name, which callers build as {γ name}dst. Tags must be unique per site and must not collide with any parameter's name.

source
def Proofs.StableHLO.bnFwdSite (B oc hh ww : ℕ) (sync : Bool) (replicas : ℕ) (epsStr gN btN tag xIn : String) :
StateM EmitS (String × String × String)

One BatchNorm FORWARD site. At sync = false the batch-BN node bnBatchF. At sync = true the sync-BN composition, in TWO rounds (Chan's parallel variance): this replica's μ_r (bnBatchMeanB) all-reduced to the global μ; then σ²_r + (μ_r − μ)² (bnBatchVarAtB) all-reduced to the global σ² — each by prettyAllReduceMean, the SAME collective node the parameter gradients ride; packed (bnPackB); then bnSyncF, normalising with the global statistics. Returns (code, y, st), st the packed [μ ‖ σ²] SSA name ("" exactly when sync = false): the backward, the γ gradient and the handed-back running stats all read it, which is what makes the four agree on ONE x̂.

tag names the collectives' SSA values (%arsum{tag}mu / %armean{tag}var …), so it must be unique per site and disjoint from every parameter's ({p}g1 is a parameter; {p}g1mu is this).

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    def Proofs.StableHLO.bnBackSite (B oc hh ww : ℕ) (sync : Bool) (replicas : ℕ) (epsStr gN xN tag dyIn st : String) :
    StateM EmitS (String × String)

    One BatchNorm BACKWARD site — the input cotangent. At sync = false bnBatchBack; at sync = true this replica's [μ ‖ σ² ‖ mean(γ·dy) ‖ mean(x̂·γ·dy)] (bnSyncDyStatsB, reading the forward's packed st, so its x̂ is the forward's), all-reduced, then bnSyncBack. One collective per BN layer in the backward (three per layer per step with the forward's two).

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      source
      def Proofs.StableHLO.bnGammaSite (B oc hh ww : ℕ) (sync : Bool) (epsStr xN dyIn st : String) :
      StateM EmitS (String × String)

      One BatchNorm γ-GRADIENT site. bnGammaGradB rebuilds x̂ from ITS OWN batch, which under sync-BN is not the x̂ the forward used — so at sync = true it is bnSyncGammaGradB, reading the forward's packed global statistics. β's gradient is Σ dy, reads no statistic, and stays bnBetaGradB either way.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For