fix: free per-expert weight lists after stacking in CuTeDSL runner

_ensure_stacked() creates stacked copies of all weights but never freed
the per-expert lists. For 256 experts on a 175GB model, this doubles
weight memory to ~350GB, causing OOM.

Now the per-expert lists (l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs)
are set to None after stacking, keeping only the single stacked copy.
This commit is contained in:
2026-05-16 08:54:52 +00:00
parent 99c11c218d
commit cdd813cf7e

View File

@@ -45,7 +45,11 @@ class CuTeDSLMoERunner:
self._l2_gsb = None
def _ensure_stacked(self):
"""Lazily stack weight tensors into the format the kernel expects."""
"""Lazily stack weight tensors into the format the kernel expects.
After stacking, the per-expert lists are freed to avoid holding
two copies of ~175GB of weight data in GPU memory.
"""
if self._l1_mat_b is not None:
return
self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4))
@@ -54,6 +58,13 @@ class CuTeDSLMoERunner:
self._l2_scale_b = assemble_scales_3d_side(self.l2_sf)
self._l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=self.device)
self._l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=self.device)
# Free per-expert lists — stacked tensors are the only copy now
self.l1_fp4 = None
self.l1_sf = None
self.l1_gs = None
self.l2_fp4 = None
self.l2_sf = None
self.l2_gs = None
def prepare_weights_direct(self, l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs):
"""Set weights directly from checkpoint (no dequant→requant)."""