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:
@@ -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)."""
|
||||
|
||||
Reference in New Issue
Block a user