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).
This commit is contained in:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user