Documentation

LeanMlir.Proofs.Nets.EfficientNet.EfficientNetSyncB

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

EfficientNetFullB0.lean (T2) says the typed batch-BN graph denotes efficientnetForwardB_full 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 T2's 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 (efficientnetFwdGraphSync_full R hR N epsStr w e r)
  = batchShard R N 10 (efficientnetForwardB_full (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 efficientnetForwardB_full, at N := R·N.

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:

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 T3 half (EfficientNetSyncStepTieG.lean). ⚠ That the R replicas' inputs ARE the shards of one batch is the driver's. ⚠ The emitted artifact's stochastic-depth and classifier-dropout variants add per-example scalings T2 does not state. ⚠ The lowerer's all_reduce is trusted as every other op's lowering is.

theorem Proofs.StableHLO.den_swishF_shard {R N n : } (e : Fin RSHlo (N * n)) (X : Vec (R * N * n)) (he : ∀ (r : Fin R), den (e r) = batchShard R N n X r) (r : Fin R) :
den (e r).swishF = batchShard R N n (swish (R * N * n) X) r

Swish on every replica denotes the shard of the global swish — pointwise, like relu.

theorem Proofs.StableHLO.den_addV_shard {R N n : } (a b : Fin RSHlo (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).addV (b r)) = batchShard R N n (fun (j : Fin (R * N * n)) => A j + B j) r

The MBConv identity skip (addV) on every replica is the shard of the global one.

def Proofs.StableHLO.mbNoExpGraphSync (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc rd kh kw : } (q : MBWNoExp ic oc rd kh kw) (e : Fin RSHlo (N * (ic * h * w))) :
Fin RSHlo (N * (oc * h * w))

MBConv1 (b1, no expand) at sync-BN, over the replica family: depthwise → sync-BN → swish → SE → 1×1 project → sync-BN.

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    theorem Proofs.StableHLO.mbNoExpGraphSync_shard (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc rd kh kw : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (q : MBWNoExp ic oc rd kh kw) (e : Fin RSHlo (N * (ic * h * w))) (X : Vec (R * N * (ic * h * w))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (ic * h * w) X r) (r : Fin R) :
    den (mbNoExpGraphSync p epsStr R hR N h w q e r) = batchShard R N (oc * h * w) (mbNoExpW (R * N) h w q X) r
    def Proofs.StableHLO.mbBodyGraphSync (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc rd kh kw : } (q : MBW ic mid oc rd kh kw) (e : Fin RSHlo (N * (ic * h * w))) :
    Fin RSHlo (N * (oc * h * w))

    The MBConv6 body at sync-BN — expand 1×1 → sync-BN → swish → depthwise → sync-BN → swish → SE → project 1×1 → sync-BN — shared by the residual (b3, b5, …) and the no-skip widening (b9, b16) blocks.

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      theorem Proofs.StableHLO.mbBodyGraphSync_shard (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc rd kh kw : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (q : MBW ic mid oc rd kh kw) (e : Fin RSHlo (N * (ic * h * w))) (X : Vec (R * N * (ic * h * w))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (ic * h * w) X r) (r : Fin R) :
      den (mbBodyGraphSync p epsStr R hR N h w q e r) = batchShard R N (oc * h * w) (mbExpW (R * N) h w q X) r
      def Proofs.StableHLO.mbExpGraphSync (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc rd kh kw : } (q : MBW ic mid oc rd kh kw) (e : Fin RSHlo (N * (ic * h * w))) :
      Fin RSHlo (N * (oc * h * w))

      The no-skip widening block (b9, b16) is the body alone.

      Equations
      Instances For
        theorem Proofs.StableHLO.mbExpGraphSync_shard (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc rd kh kw : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (q : MBW ic mid oc rd kh kw) (e : Fin RSHlo (N * (ic * h * w))) (X : Vec (R * N * (ic * h * w))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (ic * h * w) X r) (r : Fin R) :
        den (mbExpGraphSync p epsStr R hR N h w q e r) = batchShard R N (oc * h * w) (mbExpW (R * N) h w q X) r
        def Proofs.StableHLO.mbResidGraphSync (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {c mid rd kh kw : } (q : MBW c mid c rd kh kw) (e : Fin RSHlo (N * (c * h * w))) :
        Fin RSHlo (N * (c * h * w))

        The residual MBConv6 block (b3, b5, b7, b8, b10, b11, b13–b15): the body plus the identity skip, addV body e, as mbResidGraphB orders it.

        Equations
        Instances For
          theorem Proofs.StableHLO.mbResidGraphSync_shard (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {c mid rd kh kw : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (q : MBW c mid c rd kh kw) (e : Fin RSHlo (N * (c * h * w))) (X : Vec (R * N * (c * h * w))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (c * h * w) X r) (r : Fin R) :
          den (mbResidGraphSync p epsStr R hR N h w q e r) = batchShard R N (c * h * w) (mbResidW (R * N) h w q X) r
          def Proofs.StableHLO.mbStridedGraphSync (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc rd kh kw : } (q : MBW ic mid oc rd kh kw) (e : Fin RSHlo (N * (ic * (2 * h) * (2 * w)))) :
          Fin RSHlo (N * (oc * h * w))

          The strided MBConv6 block (b2, b4, b6, b12): expand at the input grid 2h×2w, the strided depthwise down to h×w, SE and project there. Three sync sites, the expand one at 2h×2w.

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            theorem Proofs.StableHLO.mbStridedGraphSync_shard (p epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc rd kh kw : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (q : MBW ic mid oc rd kh kw) (e : Fin RSHlo (N * (ic * (2 * h) * (2 * w)))) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (ic * (2 * h) * (2 * w)) X r) (r : Fin R) :
            den (mbStridedGraphSync p epsStr R hR N h w q e r) = batchShard R N (oc * h * w) (mbStridedW (R * N) h w q X) r
            def Proofs.StableHLO.stemGraphSync (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc : } (Ws : Kernel4 oc ic 3 3) (bs : Vec oc) (εs : ) (γs βs : Vec oc) (e : Fin RSHlo (N * (ic * (2 * h) * (2 * w)))) :
            Fin RSHlo (N * (oc * h * w))

            Stem at sync-BN, over the replica family: 3×3/s2 conv (XLA-SAME) → sync-BN → swish.

            Equations
            • One or more equations did not get rendered due to their size.
            Instances For
              theorem Proofs.StableHLO.stemGraphSync_shard (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (Ws : Kernel4 oc ic 3 3) (bs : Vec oc) (εs : ) (γs βs : Vec oc) (e : Fin RSHlo (N * (ic * (2 * h) * (2 * w)))) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (ic * (2 * h) * (2 * w)) X r) (r : Fin R) :
              den (stemGraphSync epsStr R hR N h w Ws bs εs γs βs e r) = batchShard R N (oc * h * w) (stemB (R * N) Ws bs εs γs βs X) r
              def Proofs.StableHLO.headGraphSync (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {c oc nC : } (Wh : Kernel4 oc c 1 1) (bh : Vec oc) (εh : ) (γh βh : Vec oc) (Wfc : Mat oc nC) (bfc : Vec nC) (e : Fin RSHlo (N * (c * h * w))) :
              Fin RSHlo (N * nC)

              Head at sync-BN, over the replica family: 1×1 conv → sync-BN → swish → GAP → dense.

              Equations
              • One or more equations did not get rendered due to their size.
              Instances For
                theorem Proofs.StableHLO.headGraphSync_shard (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {c oc nC : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (Wh : Kernel4 oc c 1 1) (bh : Vec oc) (εh : ) (γh βh : Vec oc) (Wfc : Mat oc nC) (bfc : Vec nC) (e : Fin RSHlo (N * (c * h * w))) (X : Vec (R * N * (c * h * w))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (c * h * w) X r) (r : Fin R) :
                den (headGraphSync epsStr R hR N h w Wh bh εh γh βh Wfc bfc e r) = batchShard R N nC (headFwdB (R * N) Wh bh εh γh βh Wfc bfc X) r
                def Proofs.StableHLO.efficientnetFwdGraphSync_full (R : ) (hR : 0 < R) (N : ) (epsStr : String) (w : B0Weights) (e : Fin RSHlo (N * (3 * 224 * 224))) :
                Fin RSHlo (N * 10)

                The sync-BN data-parallel EfficientNet-B0 forward graph, over the replica family. T2's efficientnetFwdGraphB_full with every one of the 49 BatchNorms a bnSyncSiteLA over all R replicas, fed each replica's own input subgraph e r; block prefixes and collective tags are EfficientNetRender's.

                Equations
                • One or more equations did not get rendered due to their size.
                Instances For
                  theorem Proofs.StableHLO.efficientnetFwdGraphSync_full_shard (R : ) (hR : 0 < R) (N : ) (hN : 0 < N) (epsStr : String) (w : B0Weights) (e : Fin RSHlo (N * (3 * 224 * 224))) (X : Vec (R * N * (3 * 224 * 224))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (3 * 224 * 224) X r) (r : Fin R) :
                  den (efficientnetFwdGraphSync_full R hR N epsStr w e r) = batchShard R N 10 (efficientnetForwardB_full (R * N) w X) r

                  ⭐⭐ T2 at synchronised BatchNorm: replica r's forward IS shard r of the global-batch forward. Given that the replicas' inputs are the shards of one batch X of R·N images, the sync-BN graph on replica r denotes batchShard r of efficientnetForwardB_full (R * N) w X — the committed batch-BN forward, at the global batch. One block lemma per stage, the shard hypothesis threaded from each into the next.