[BugFix] Fix pipeline parallel (#24621)

Signed-off-by: Nick Hill <nhill@redhat.com>
This commit is contained in:
Nick Hill 2025-09-10 23:05:30 -07:00 committed by GitHub
parent 29799ddacc
commit e2d8c27f68
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 6 additions and 2 deletions

View File

@ -71,6 +71,10 @@ class UniProcExecutor(ExecutorBase):
self.shutdown()
return
def shutdown(self) -> None:
if worker := self.driver_worker:
worker.shutdown()
UniProcExecutorAsync = UniProcExecutor

View File

@ -2070,7 +2070,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
sampler_output = self._sample(logits, spec_decode_metadata)
with record_function_or_nullcontext("Bookkeep"):
assert isinstance(hidden_states, torch.Tensor)
(
num_nans_in_logits,
logprobs_lists,

View File

@ -45,7 +45,8 @@ class KVConnectorModelRunnerMixin:
@staticmethod
def ensure_kv_transfer_shutdown() -> None:
if has_kv_transfer_group():
# has_kv_transfer_group can be None during interpreter shutdown.
if has_kv_transfer_group and has_kv_transfer_group():
ensure_kv_transfer_shutdown()
@staticmethod