From 23e88638aa75feb370164b091e9f2328da95f878 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 22:55:11 +0000 Subject: [PATCH] single_shot: memory-efficient MoE loading (CPU stacking, one-shot GPU transfer) Build stacked (E, N, K) tensors incrementally on CPU, then move to GPU in one shot. Avoids holding 384 individual expert weight+scale tensors on GPU simultaneously (~3x memory savings per layer). --- single_shot_inference.py | 110 ++++++++++++++++++++++++++------------- 1 file changed, 74 insertions(+), 36 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 1e056759..3446891a 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -862,49 +862,87 @@ def main(): # MoE weight loading helpers (stacked path for production GEMM) # ===================================================================== def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg): - """Load MoE expert weights into Nvfp4MoE via stacked path.""" + """Load MoE expert weights into Nvfp4MoE via stacked path. + + Memory-efficient: builds stacked tensors incrementally on CPU, + then moves to GPU in one shot. Avoids holding 384 individual + expert weight tensors on GPU simultaneously (~3× memory savings). + """ n_e = cfg["n_routed_experts"] moe_inter = cfg.get("moe_intermediate_size", 3072) H = cfg["hidden_size"] - l1_gate_fp4, l1_gate_sf, l1_gate_gs = [], [], [] - l1_up_fp4, l1_up_sf = [], [] - l2_fp4, l2_sf, l2_gs = [], [], [] - + # Build stacked tensors incrementally on CPU + # gate_proj and up_proj: (inter, K_packed) per expert → L1 stacked (E, 2*inter, K_packed) + # down_proj: (H, K_packed) per expert → L2 stacked (E, H, K_packed) + + # Get dimensions from first expert + w0 = all_w.get(f"{pfx}.experts.0.gate_proj.weight") + if w0 is None: + log.warning(f"L{li}: No expert weights found") + return + gate_N, gate_K = w0.shape # (inter, K_packed) + + l1_stacked = torch.zeros(n_e, 2 * gate_N, gate_K, dtype=w0.dtype) + l1_sf_stacked = None + l2_stacked = None + l2_sf_stacked = None + l1_gs = [] + l2_gs = [] + + # Determine L1 SF shape from first expert + ws0 = all_w.get(f"{pfx}.experts.0.gate_proj.weight_scale") + if ws0 is not None: + sf_N, sf_K = ws0.shape + l1_sf_stacked = torch.zeros(n_e, 2 * sf_N, sf_K, dtype=ws0.dtype) + + # Get L2 shape + dw0 = all_w.get(f"{pfx}.experts.0.down_proj.weight") + if dw0 is not None: + down_N, down_K = dw0.shape + l2_stacked = torch.zeros(n_e, down_N, down_K, dtype=dw0.dtype) + dws0 = all_w.get(f"{pfx}.experts.0.down_proj.weight_scale") + if dws0 is not None: + dsf_N, dsf_K = dws0.shape + l2_sf_stacked = torch.zeros(n_e, dsf_N, dsf_K, dtype=dws0.dtype) + + # Fill stacked tensors for eid in range(n_e): - for proj, fp4_l, sf_l, gs_l in [ - ('gate_proj', l1_gate_fp4, l1_gate_sf, l1_gate_gs), - ('up_proj', l1_up_fp4, l1_up_sf, None), - ('down_proj', l2_fp4, l2_sf, l2_gs), - ]: - w_k = f"{pfx}.experts.{eid}.{proj}.weight" - ws_k = f"{pfx}.experts.{eid}.{proj}.weight_scale" - isc_k = f"{pfx}.experts.{eid}.{proj}.input_scale" - w, ws, isc = all_w.get(w_k), all_w.get(ws_k), all_w.get(isc_k) - if w is not None and ws is not None: - fp4_l.append(w.to(dev)) - sf_l.append(ws.to(dev)) - if gs_l is not None: - gs_l.append(isc.float().item() if isc is not None else 1.0 / (6.0 * 448.0)) + # L1: gate + up + gw = all_w.get(f"{pfx}.experts.{eid}.gate_proj.weight") + gws = all_w.get(f"{pfx}.experts.{eid}.gate_proj.weight_scale") + gisc = all_w.get(f"{pfx}.experts.{eid}.gate_proj.input_scale") + uw = all_w.get(f"{pfx}.experts.{eid}.up_proj.weight") + uws = all_w.get(f"{pfx}.experts.{eid}.up_proj.weight_scale") - if l1_gate_fp4 and l1_up_fp4: - l1_stacked = torch.stack([torch.cat([g, u], dim=0) for g, u in zip(l1_gate_fp4, l1_up_fp4)]) - l1_sf_stacked = torch.stack([torch.cat([gs, us], dim=0) for gs, us in zip(l1_gate_sf, l1_up_sf)]) - l1_gs = l1_gate_gs - else: - l1_stacked = None; l1_sf_stacked = None; l1_gs = [1.0 / (6.0 * 448.0)] * n_e - if l2_fp4: - l2_stacked = torch.stack(l2_fp4) - l2_sf_stacked = torch.stack(l2_sf) - l2_gs = l2_gs - else: - l2_stacked = None; l2_sf_stacked = None; l2_gs = [1.0 / (6.0 * 448.0)] * n_e + if gw is not None and uw is not None: + l1_stacked[eid, :gate_N] = gw + l1_stacked[eid, gate_N:] = uw + if gws is not None and uws is not None and l1_sf_stacked is not None: + l1_sf_stacked[eid, :sf_N] = gws + l1_sf_stacked[eid, sf_N:] = uws + l1_gs.append(gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)) + + # L2: down + dw = all_w.get(f"{pfx}.experts.{eid}.down_proj.weight") + dws = all_w.get(f"{pfx}.experts.{eid}.down_proj.weight_scale") + disc = all_w.get(f"{pfx}.experts.{eid}.down_proj.input_scale") + if dw is not None: + l2_stacked[eid] = dw + if dws is not None and l2_sf_stacked is not None: + l2_sf_stacked[eid] = dws + l2_gs.append(disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)) + + # Move to GPU in one shot + l1_stacked = l1_stacked.to(dev) + l1_sf_stacked = l1_sf_stacked.to(dev) if l1_sf_stacked is not None else None + l2_stacked = l2_stacked.to(dev) if l2_stacked is not None else None + l2_sf_stacked = l2_sf_stacked.to(dev) if l2_sf_stacked is not None else None + l1_gs = l1_gs if l1_gs else [1.0 / (6.0 * 448.0)] * n_e + l2_gs = l2_gs if l2_gs else [1.0 / (6.0 * 448.0)] * n_e - if l1_stacked is not None: - moe.prepare_weights_from_stacked(l1_stacked, l1_sf_stacked, l1_gs, - l2_stacked, l2_sf_stacked, l2_gs) - else: - log.warning(f"L{li}: MoE weight stacking failed") + moe.prepare_weights_from_stacked(l1_stacked, l1_sf_stacked, l1_gs, + l2_stacked, l2_sf_stacked, l2_gs) def _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg):