loader: yield filtered MTP weights lazily to avoid OOM hang on multi-layer EAGLE (#25748)
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user