loader: yield filtered MTP weights lazily to avoid OOM hang on multi-layer EAGLE (#25748)

This commit is contained in:
Xinyuan Tong
2026-05-20 12:33:54 +08:00
committed by GitHub
parent af22390af7
commit 0aedc5678b
+7 -6
View File
@@ -583,10 +583,12 @@ class DefaultModelLoader(BaseModelLoader):
@classmethod @classmethod
def _filter_mtp_weights( def _filter_mtp_weights(
cls, weights_iterator, prefix: str, draft_model_idx: int cls, weights_iterator, prefix: str, draft_model_idx: int
) -> Tuple[Tuple[str, torch.Tensor], ...]: ) -> Generator[Tuple[str, torch.Tensor], None, None]:
"""Filter MTP (Multi-Token Prediction) weights to keep only the """Filter MTP weights to keep only the specified draft model layer
specified draft model layer and remap it to layer 0.""" and remap it to layer 0. Yields lazily so the upstream buffered
filtered_weights = [] iterator's sliding window actually bounds CPU memory — eager
materialization caused page-reclaim hangs on large MoE checkpoints
with multi-layer EAGLE."""
for name, tensor in weights_iterator: for name, tensor in weights_iterator:
match = cls._MTP_PATTERN.match(name) match = cls._MTP_PATTERN.match(name)
if match is not None: if match is not None:
@@ -596,8 +598,7 @@ class DefaultModelLoader(BaseModelLoader):
new_name = name.replace(match.group(), "model.mtp.layers.0.") new_name = name.replace(match.group(), "model.mtp.layers.0.")
else: else:
new_name = name new_name = name
filtered_weights.append((prefix + new_name, tensor)) yield (prefix + new_name, tensor)
return tuple(filtered_weights)
def _get_all_weights( def _get_all_weights(
self, self,