support mtp

Author: hongchao

checkpoint_engine/worker.py

@@ -9,6 +9,7 @@
 
 from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid
 
+_WEIGHTS_TYPE = list[tuple[str, torch.Tensor]]
 
 def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor:
     func, args = handle
@@ -31,9 +32,9 @@ class FlattenedTensorMetadata(TypedDict):
 
 def _extract_weights(
     payload: list[FlattenedTensorMetadata], buffer: torch.Tensor
-) -> list[tuple[str, torch.Tensor]]:
+) -> _WEIGHTS_TYPE:
     assert buffer is not None
-    weights: list[tuple[str, torch.Tensor]] = []
+    weights: _WEIGHTS_TYPE = []
     for item in payload:
         shape = item["shape"]
         if isinstance(shape, list | tuple):
@@ -166,11 +167,18 @@ def update_weights_from_ipc(self, zmq_handles: dict[str, str]):
             self.device = torch.device(f"npu:{self.local_rank}")
         assert self.device is not None
 
+        def _load_weights(weights: _WEIGHTS_TYPE):
+            # Load main model weights
+            self.model_runner.model.load_weights(weights)
+            # Load drafter model weights if MTP/speculative decoding is enabled
+            if self.model_runner.get("use_spec_decode", False):
+                self.model_runner.drafter.model.load_weights(weights=weights)
+
         update_weights_from_ipc(
             self._zmq_ctx,
             zmq_handles[self._device_uuid],
             device_id=self.device.index,
-            run=self.model_runner.model.load_weights,
+            run=_load_weights,
             post_hook=lambda: process_weights_after_loading(
                 self.model_runner.model, self.model_config, self.device
             ),