Documentation

LeanMlir.Proofs.Nets.MobileNet.MobileNetV2SyncB

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

MobileNetV2FullB.lean (T2) says the typed batch-BN graph denotes mobilenetv2ForwardB_full N w on one device. MobileNetV2RenderB's data-parallel step normalises with the GLOBAL batch's statistics at replicas > 1: every one of the 52 BatchNorm sites is the sync-BN composition — 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 (mobilenetv2FwdGraphSync_full R hR N epsStr w e r)
  = batchShard R N nCls (mobilenetv2ForwardB_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 mobilenetv2ForwardB_full, at N := R·N.

How it is proved #

As ResNet34SyncB.lean proves ResNet-34's: by induction on the chain, one block at a time, with the shard hypothesis ∀ r, den (e r) = batchShard R N _ X r carried from block to block.

The graph reuses ResNet34SyncB's site verbatim — it is net-agnostic — so the only new lemma is the relu6 one; the rest of this file is MobileNetV2's six block shapes and their chain.

Names #

Parameter names are MobileNetV2FullB's (%b{k}{e,d,p}{W,g,bt}, %sW, %hW, …, the convBias := false zero-bias operands %zb{c}). A BatchNorm site with γ %b{k}dg gathers its statistics as %arsum / %armean of b{k}dgmu and b{k}dgvar, each over a [c] vector — the γ name without %, then mu / var.

What is NOT claimed here #

⚠ The backward and the parameter collectives are the T3 half (MobileNetV2SyncStepTieB.lean). ⚠ That the R replicas' inputs ARE the shards of one batch is the driver's, as in DataParallelSync.lean. ⚠ The lowerer's all_reduce is trusted as every other op's lowering is.

theorem Proofs.StableHLO.den_relu6_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 (SHlo.batchOp BatchableOp.relu6 (e r)) = batchShard R N n (relu6 (R * N * n) X) r

relu6 on every replica is the shard of the global relu6 — the clamp reads one cell, so cutting the batch before or after it is the same. The MobileNetV2 peer of den_relu_shard, stated at the whole-batch relu6 the committed forward is written in.

def Proofs.StableHLO.mnv2StemGraphSync (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc kH kW : } (Ws : Kernel4 oc ic kH kW) (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: 3x3/s2 XLA-SAME conv → sync-BN → relu6.

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    theorem Proofs.StableHLO.mnv2StemGraphSync_shard (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc kH kW : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (Ws : Kernel4 oc ic kH kW) (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 (mnv2StemGraphSync epsStr R hR N h w Ws bs εs γs βs e r) = batchShard R N (oc * h * w) (mnv2StemB (R * N) h w Ws bs εs γs βs X) r
    def Proofs.StableHLO.mnv2NoExpGraphSync (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc : } (p : IVWNoExp ic oc) (e : Fin RSHlo (N * (ic * h * w))) :
    Fin RSHlo (N * (oc * h * w))

    t = 1 bottleneck (b1) at sync-BN, over the replica family: depthwise → sync-BN → relu6 → project 1x1 → sync-BN. Two sync sites.

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      theorem Proofs.StableHLO.mnv2NoExpGraphSync_shard (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (p : IVWNoExp ic oc) (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 (mnv2NoExpGraphSync pfx epsStr R hR N h w p e r) = batchShard R N (oc * h * w) (mnv2NoExpB (R * N) h w p X) r
      def Proofs.StableHLO.mnv2ExpOnlyGraphSync (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc : } (p : IVW ic mid oc) (e : Fin RSHlo (N * (ic * h * w))) :
      Fin RSHlo (N * (oc * h * w))

      Stride-1 no-skip bottleneck (b11, b17) at sync-BN, over the replica family: expand → depthwise → project, a sync site after each, relu6 after the first two.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        theorem Proofs.StableHLO.mnv2ExpOnlyGraphSync_shard (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (p : IVW ic mid oc) (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 (mnv2ExpOnlyGraphSync pfx epsStr R hR N h w p e r) = batchShard R N (oc * h * w) (mnv2ExpOnlyB (R * N) h w p X) r
        def Proofs.StableHLO.mnv2ResidGraphSync (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {c mid : } (p : IVW c mid c) (e : Fin RSHlo (N * (c * h * w))) :
        Fin RSHlo (N * (c * h * w))

        Stride-1 skip bottleneck at sync-BN: the body plus the addVB identity skip, the block input shared between both arms on every replica.

        Equations
        Instances For
          theorem Proofs.StableHLO.mnv2ResidGraphSync_shard (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {c mid : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (p : IVW c mid c) (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 (mnv2ResidGraphSync pfx epsStr R hR N h w p e r) = batchShard R N (c * h * w) (mnv2ResidB (R * N) h w p X) r
          def Proofs.StableHLO.mnv2StridedGraphSync (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc : } (p : IVW ic mid oc) (e : Fin RSHlo (N * (ic * (2 * h) * (2 * w)))) :
          Fin RSHlo (N * (oc * h * w))

          Stride-2 downsampling bottleneck at sync-BN, over the replica family: expand at 2h x 2w, the XLA-SAME strided depthwise, project at h x w; three sync sites.

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            theorem Proofs.StableHLO.mnv2StridedGraphSync_shard (pfx epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic mid oc : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (p : IVW ic mid 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 (mnv2StridedGraphSync pfx epsStr R hR N h w p e r) = batchShard R N (oc * h * w) (mnv2StridedB (R * N) h w p X) r
            def Proofs.StableHLO.mnv2HeadGraphSync (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc nCls : } (Wh : Kernel4 oc ic 1 1) (bh : Vec oc) (εh : ) (γh βh : Vec oc) (Wd : Mat oc nCls) (bd : Vec nCls) (e : Fin RSHlo (N * (ic * h * w))) :
            Fin RSHlo (N * nCls)

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

            Equations
            • One or more equations did not get rendered due to their size.
            Instances For
              theorem Proofs.StableHLO.mnv2HeadGraphSync_shard (epsStr : String) (R : ) (hR : 0 < R) (N h w : ) {ic oc nCls : } (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (Wh : Kernel4 oc ic 1 1) (bh : Vec oc) (εh : ) (γh βh : Vec oc) (Wd : Mat oc nCls) (bd : Vec nCls) (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 (mnv2HeadGraphSync epsStr R hR N h w Wh bh εh γh βh Wd bd e r) = batchShard R N nCls (mnv2HeadB (R * N) h w Wh bh εh γh βh Wd bd X) r
              def Proofs.StableHLO.mobilenetv2FwdGraphSync_full (R : ) (hR : 0 < R) (N : ) (epsStr : String) {nCls : } (w : MNV2BWeights nCls) (e : Fin RSHlo (N * (3 * (2 * 112) * (2 * 112)))) :
              Fin RSHlo (N * nCls)

              The sync-BN data-parallel MobileNetV2 forward graph, over the replica family. T2's mobilenetv2FwdGraphB_full with every BatchNorm a bnSyncSiteLA over all R replicas; block prefixes, parameter names and collective tags are the render's.

              Equations
              • One or more equations did not get rendered due to their size.
              Instances For
                theorem Proofs.StableHLO.mobilenetv2FwdGraphSync_full_shard (R : ) (hR : 0 < R) (N : ) (hN : 0 < N) (epsStr : String) {nCls : } (w : MNV2BWeights nCls) (e : Fin RSHlo (N * (3 * (2 * 112) * (2 * 112)))) (X : Vec (R * N * (3 * (2 * 112) * (2 * 112)))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (3 * (2 * 112) * (2 * 112)) X r) (r : Fin R) :
                den (mobilenetv2FwdGraphSync_full R hR N epsStr w e r) = batchShard R N nCls (mobilenetv2ForwardB_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 examples, the sync-BN graph on replica r denotes batchShard r of mobilenetv2ForwardB_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.