diff --git a/lmdeploy/pytorch/engine/logits_process.py b/lmdeploy/pytorch/engine/logits_process.py index e8da57d29d..3cdad21b47 100644 --- a/lmdeploy/pytorch/engine/logits_process.py +++ b/lmdeploy/pytorch/engine/logits_process.py @@ -285,6 +285,13 @@ def to_device(self, device: str, non_blocking: bool = False): return SamplingInputs(**out_dict) + def record_stream(self, stream: torch.cuda.Stream) -> None: + """Record forward-stream use of tensor fields.""" + for f in fields(self): + value = getattr(self, f.name) + if isinstance(value, torch.Tensor) and value.is_cuda: + value.record_stream(stream) + def get_delta(self) -> SamplingInputsDelta: """Get delta.""" delta = SamplingInputsDelta() diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py index 09f0a78237..dcf9a32434 100644 --- a/lmdeploy/pytorch/engine/model_agent/agent.py +++ b/lmdeploy/pytorch/engine/model_agent/agent.py @@ -2,6 +2,7 @@ import asyncio import time from collections import deque +from collections.abc import Mapping from contextlib import contextmanager from dataclasses import dataclass, field, fields from multiprocessing.reduction import ForkingPickler @@ -46,6 +47,7 @@ logger = get_logger('lmdeploy') _H2D_TRANSFER_KEY = '_h2d_transfer' +_H2D_INPUT_KEYS = ('inputs', 'delta', 'sampling_inputs', 'stopping_criteria', 'extra_inputs') @dataclass @@ -236,6 +238,23 @@ def _try_to_cuda(val, non_blocking: bool = False): raise RuntimeError(f'Can not cast {type(val)} to cuda.') +def _record_forward_input_stream(forward_inputs: Mapping[str, Any], stream: torch.cuda.Stream) -> None: + """Register the forward stream with each H2D payload owner.""" + for key in _H2D_INPUT_KEYS: + value = forward_inputs.get(key) + if value is None: + continue + if isinstance(value, torch.Tensor): + if value.is_cuda: + value.record_stream(stream) + continue + + record_stream = getattr(value, 'record_stream', None) + if not callable(record_stream): + raise TypeError(f'H2D input {key!r} must be a tensor or implement record_stream().') + record_stream(stream) + + SwapMap = dict[int, int] @@ -938,6 +957,7 @@ async def _async_loop_background(self, forward_event: asyncio.Event = None): if h2d_transfer is not None: self._keep_h2d_transfer(h2d_transfer) self.stream.wait_event(h2d_transfer.event) + _record_forward_input_stream(forward_inputs, self.stream) await self._async_step(**forward_inputs, ) if forward_event is not None: @@ -970,18 +990,20 @@ async def _async_loop_inputs_preprocess(self, forward_event: asyncio.Event = Non falling back to dummy inputs. """ non_blocking = True - keys = ['inputs', 'delta', 'sampling_inputs', 'stopping_criteria', 'extra_inputs'] while True: forward_inputs = await self._pre_in_que.get() forward_inputs_cuda = {} forward_inputs_cuda.update(forward_inputs) - h2d_refs = {k: forward_inputs_cuda[k] for k in keys if forward_inputs_cuda.get(k) is not None} + h2d_refs = { + key: forward_inputs_cuda[key] + for key in _H2D_INPUT_KEYS if forward_inputs_cuda.get(key) is not None + } logger.debug('preprocessing forward inputs.') with torch.cuda.stream(self.out_stream), torch.inference_mode(), record_function('inputs_H2D'): - for k in keys: - if k not in forward_inputs_cuda: + for key in _H2D_INPUT_KEYS: + if key not in forward_inputs_cuda: continue - forward_inputs_cuda[k] = _try_to_cuda(forward_inputs_cuda[k], non_blocking=non_blocking) + forward_inputs_cuda[key] = _try_to_cuda(forward_inputs_cuda[key], non_blocking=non_blocking) h2d_event = torch.cuda.Event() h2d_event.record() forward_inputs_cuda[_H2D_TRANSFER_KEY] = _H2DTransfer(h2d_event, h2d_refs) diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index 39d68dd854..9e5030a7fe 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/lmdeploy/pytorch/model_inputs.py @@ -101,6 +101,29 @@ def to_device(self, device: str, non_blocking: bool = False): return VisionModelInputs(**out_dict) + def record_stream(self, stream: torch.cuda.Stream) -> None: + """Record forward-stream use of vision tensor fields.""" + for f in fields(self): + key = f.name + value = getattr(self, key) + if isinstance(value, torch.Tensor): + if value.is_cuda: + value.record_stream(stream) + elif key == 'input_embedding_ranges': + for tensor in value or (): + if tensor.is_cuda: + tensor.record_stream(stream) + elif key == 'input_embeddings': + for tensors in value or (): + for tensor in tensors: + if tensor.is_cuda: + tensor.record_stream(stream) + elif key == 'input_multimodals': + for multimodals in value or (): + for items in multimodals.values(): + for item in items: + item.record_stream(stream) + def get_inputs(self, history_lengths: torch.Tensor, seq_lengths: torch.Tensor): """Get vision embedding inputs.""" input_embeddings = None @@ -172,6 +195,13 @@ def to_device(self, device: str, non_blocking: bool = False): return ModelInputsDelta(**out_dict) + def record_stream(self, stream: torch.cuda.Stream) -> None: + """Record forward-stream use of tensor fields.""" + for f in fields(self): + value = getattr(self, f.name) + if isinstance(value, torch.Tensor) and value.is_cuda: + value.record_stream(stream) + def log_info(self): """Get log info.""" ret = (f'num_tokens={self.indices.numel()}, batch_size={self.indices.numel()}' @@ -258,6 +288,16 @@ def to_device(self, device: str, non_blocking: bool = False): return ModelInputs(**out_dict) + def record_stream(self, stream: torch.cuda.Stream) -> None: + """Record forward-stream use of model tensor fields.""" + for f in fields(self): + value = getattr(self, f.name) + if isinstance(value, torch.Tensor): + if value.is_cuda: + value.record_stream(stream) + elif isinstance(value, VisionModelInputs): + value.record_stream(stream) + def build_dp_meta(self, num_tokens: list[int]): """Build dp meta.""" self.dp_meta = DPMeta.build(self.input_ids.numel(), num_tokens) diff --git a/lmdeploy/pytorch/multimodal/data_type.py b/lmdeploy/pytorch/multimodal/data_type.py index f778c2aeb5..833e8dff5a 100644 --- a/lmdeploy/pytorch/multimodal/data_type.py +++ b/lmdeploy/pytorch/multimodal/data_type.py @@ -99,6 +99,20 @@ def to_device(self, device: str, non_blocking: bool = False): out_dict['meta'] = new_meta return MultiModalData(**out_dict) + def record_stream(self, stream: torch.cuda.Stream) -> None: + """Record forward-stream use of multimodal tensor fields.""" + tensors = [self.data] if isinstance(self.data, Tensor) else self.data + for tensor in tensors: + if tensor.is_cuda: + tensor.record_stream(stream) + + for value in (self.meta or {}).values(): + if isinstance(value, Tensor): + if value.is_cuda: + value.record_stream(stream) + elif hasattr(value, 'record_stream'): + value.record_stream(stream) + MultiModalInputs = dict[str, list[MultiModalData]] diff --git a/lmdeploy/pytorch/strategies/base/model_agent.py b/lmdeploy/pytorch/strategies/base/model_agent.py index 4748f10597..b89b1379c6 100644 --- a/lmdeploy/pytorch/strategies/base/model_agent.py +++ b/lmdeploy/pytorch/strategies/base/model_agent.py @@ -35,6 +35,13 @@ def to_device(self, device: str, non_blocking: bool = False): """To device.""" return to_device(self, device, non_blocking) + def record_stream(self, stream: torch.cuda.Stream) -> None: + """Record forward-stream use of tensor fields.""" + for f in fields(self): + value = getattr(self, f.name) + if isinstance(value, torch.Tensor) and value.is_cuda: + value.record_stream(stream) + def broadcast(self, src: int, group, async_op=False): """Broadcast extra inputs.""" pass @@ -111,6 +118,13 @@ def to_device(self, device: str, non_blocking: bool = False): """To device.""" return to_device(self, device, non_blocking) + def record_stream(self, stream: torch.cuda.Stream) -> None: + """Record forward-stream use of tensor fields.""" + for f in fields(self): + value = getattr(self, f.name) + if isinstance(value, torch.Tensor) and value.is_cuda: + value.record_stream(stream) + class ModelAgentStrategy(ABC): """Base class for model agent strategies.""" diff --git a/tests/pytorch/engine/test_logits_process.py b/tests/pytorch/engine/test_logits_process.py index 01c6f9777c..ba2b2f142b 100644 --- a/tests/pytorch/engine/test_logits_process.py +++ b/tests/pytorch/engine/test_logits_process.py @@ -11,6 +11,35 @@ # yapf: enable +def test_sampling_inputs_record_stream_records_only_tensor_fields(): + from lmdeploy.pytorch.engine.logits_process import SamplingInputs + + recorded = [] + + class _CudaTensor(torch.Tensor): + + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.empty(1), False) + + @property + def is_cuda(self): + return True + + def record_stream(self, stream): + recorded.append((id(self), stream)) + + stream = object() + temperature = _CudaTensor() + nested_session_tensor = _CudaTensor() + inputs = SamplingInputs(temperature=temperature, + session_ctx=[{'persistent': nested_session_tensor}]) + + inputs.record_stream(stream) + + assert recorded == [(id(temperature), stream)] + + def test_process_temperature(): from lmdeploy.pytorch.engine.logits_process import _process_temperature_ diff --git a/tests/pytorch/engine/test_model_agent.py b/tests/pytorch/engine/test_model_agent.py index c074a47b5d..50b2f1b78c 100644 --- a/tests/pytorch/engine/test_model_agent.py +++ b/tests/pytorch/engine/test_model_agent.py @@ -1,7 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import asyncio from collections import deque -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from types import SimpleNamespace import pytest @@ -186,6 +186,178 @@ def test_spec_agent_reset_runtime_state_discards_chunk_carry(): assert agent._prev_chunk_last == {} +def test_record_forward_input_stream_uses_payload_protocol(): + from lmdeploy.pytorch.engine.model_agent.agent import _record_forward_input_stream + + recorded = [] + + class _CudaTensor(torch.Tensor): + + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.empty(1), False) + + @property + def is_cuda(self): + return True + + def record_stream(self, stream): + recorded.append(('tensor', id(self), stream)) + + class _Payload: + + def record_stream(self, stream): + recorded.append(('payload', id(self), stream)) + + stream = object() + tensor = _CudaTensor() + nested_tensor = _CudaTensor() + payload = _Payload() + forward_inputs = { + 'inputs': payload, + 'delta': tensor, + 'unowned_container': {'tensor': nested_tensor}, + } + + _record_forward_input_stream(forward_inputs, stream) + + assert recorded == [ + ('payload', id(payload), stream), + ('tensor', id(tensor), stream), + ] + + +def test_record_forward_input_stream_requires_payload_protocol(): + from lmdeploy.pytorch.engine.model_agent.agent import _record_forward_input_stream + + with pytest.raises(TypeError, match=r"H2D input 'sampling_inputs'.*record_stream"): + _record_forward_input_stream({'sampling_inputs': object()}, object()) + + +def test_strategy_inputs_record_stream_tensor_fields(): + from lmdeploy.pytorch.engine.model_agent.agent import _record_forward_input_stream + from lmdeploy.pytorch.strategies.ar.model_agent import ARStoppingCriteria + from lmdeploy.pytorch.strategies.ar_spec.model_agent import ARSpecExtraInputs + + recorded = [] + + class _CudaTensor(torch.Tensor): + + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.empty(1), False) + + @property + def is_cuda(self): + return True + + def record_stream(self, stream): + recorded.append((id(self), stream)) + + stream = object() + extra_tensor = _CudaTensor() + stopping_tensor = _CudaTensor() + extra_inputs = ARSpecExtraInputs(target_logits=extra_tensor) + stopping_criteria = ARStoppingCriteria(num_appendable_ids=stopping_tensor) + + _record_forward_input_stream( + { + 'extra_inputs': extra_inputs, + 'stopping_criteria': stopping_criteria, + }, stream) + + assert recorded == [ + (id(stopping_tensor), stream), + (id(extra_tensor), stream), + ] + + +def test_background_records_h2d_inputs_after_wait_and_before_forward(monkeypatch, event_loop): + from lmdeploy.pytorch.engine.model_agent import agent as agent_module + + events = [] + transfer = SimpleNamespace(event=object()) + forward_done = asyncio.Event() + + class _Stream: + + def wait_event(self, event): + assert event is transfer.event + events.append('wait') + + class _InputMaker: + + def __init__(self): + self.sent = False + + async def get(self): + if not self.sent: + self.sent = True + return { + agent_module._H2D_TRANSFER_KEY: transfer, + 'inputs': 'device-inputs', + } + await asyncio.Future() + + def step(self): + events.append('step') + + model_agent = _make_agent_with_queues() + model_agent.stream = _Stream() + model_agent.all_context = nullcontext + model_agent._keep_h2d_transfer = lambda item: events.append('keep') + + async def _async_step(**forward_inputs): + assert forward_inputs == {'inputs': 'device-inputs'} + events.append('forward') + forward_done.set() + + model_agent._async_step = _async_step + monkeypatch.setattr(agent_module, 'build_inputs_maker', lambda agent: _InputMaker()) + monkeypatch.setattr(agent_module.torch.cuda, 'stream', lambda stream: nullcontext()) + monkeypatch.setattr(agent_module, '_record_forward_input_stream', + lambda inputs, stream: events.append('record')) + + task = event_loop.create_task(model_agent._async_loop_background()) + event_loop.run_until_complete(asyncio.wait_for(forward_done.wait(), timeout=1)) + task.cancel() + event_loop.run_until_complete(asyncio.gather(task, return_exceptions=True)) + + assert events == ['keep', 'wait', 'record', 'forward', 'step'] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA allocator') +def test_record_forward_input_stream_defers_origin_stream_reuse(): + from lmdeploy.pytorch.engine.model_agent.agent import _record_forward_input_stream + + torch.cuda.synchronize() + torch.cuda.empty_cache() + h2d_stream = torch.cuda.Stream() + forward_stream = torch.cuda.Stream() + h2d_event = torch.cuda.Event() + numel = 1024 * 1024 + + with torch.cuda.stream(h2d_stream): + tensor = torch.full((numel, ), 7, dtype=torch.uint8, device='cuda') + h2d_event.record() + original_ptr = tensor.data_ptr() + + forward_stream.wait_event(h2d_event) + _record_forward_input_stream({'delta': tensor}, forward_stream) + with torch.cuda.stream(forward_stream): + torch.cuda._sleep(10_000_000) + observed = tensor.clone() + + del tensor + with torch.cuda.stream(h2d_stream): + replacement = torch.zeros((numel, ), dtype=torch.uint8, device='cuda') + + assert replacement.data_ptr() != original_ptr + forward_stream.synchronize() + h2d_stream.synchronize() + assert torch.all(observed == 7) + + class TestDrainQueues: def test_drain_empty_queues(self): diff --git a/tests/pytorch/test_model_inputs.py b/tests/pytorch/test_model_inputs.py index b947180eb6..0806d53a78 100644 --- a/tests/pytorch/test_model_inputs.py +++ b/tests/pytorch/test_model_inputs.py @@ -1,6 +1,7 @@ import torch -from lmdeploy.pytorch.model_inputs import DPMeta, ModelInputs, StepContext +from lmdeploy.pytorch.model_inputs import DPMeta, ModelInputs, ModelInputsDelta, StepContext, VisionModelInputs +from lmdeploy.pytorch.multimodal.data_type import MultiModalData def _make_model_inputs(is_decoding: bool, dp_is_decoding: bool | None = None): @@ -48,3 +49,83 @@ def test_step_context_global_is_decoding_uses_dp_global_state(): dp_meta=DPMeta(dp_is_decoding=False), ) assert not step_ctx.global_is_decoding() + + +def test_model_inputs_record_stream_matches_device_owned_fields(): + recorded = [] + + class _CudaTensor(torch.Tensor): + + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.empty(1), False) + + @property + def is_cuda(self): + return True + + def record_stream(self, stream): + recorded.append((id(self), stream)) + + stream = object() + input_ids = _CudaTensor() + embedding = _CudaTensor() + embedding_range = _CudaTensor() + multimodal_data = _CudaTensor() + multimodal_meta = _CudaTensor() + untouched_model_meta = _CudaTensor() + vision_inputs = VisionModelInputs( + input_embeddings=[[embedding]], + input_embedding_ranges=[embedding_range], + input_multimodals=[{ + 'image': [MultiModalData(data=multimodal_data, start=0, meta={'shape': multimodal_meta})] + }], + ) + inputs = _make_model_inputs(is_decoding=False) + inputs.input_ids = input_ids + inputs.vision_inputs = vision_inputs + inputs.model_metas = [{'persistent': untouched_model_meta}] + + inputs.record_stream(stream) + + assert set(recorded) == { + (id(input_ids), stream), + (id(embedding), stream), + (id(embedding_range), stream), + (id(multimodal_data), stream), + (id(multimodal_meta), stream), + } + + +def test_model_inputs_delta_record_stream_records_tensor_fields(): + recorded = [] + + class _CudaTensor(torch.Tensor): + + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.empty(1), False) + + @property + def is_cuda(self): + return True + + def record_stream(self, stream): + recorded.append((id(self), stream)) + + stream = object() + indices = _CudaTensor() + block_offsets = _CudaTensor() + delta = ModelInputsDelta(indices=indices, + block_offsets=block_offsets, + indice_cpu=None, + max_q_seqlen=1, + max_kv_seqlen=1, + sum_kv_seqlen=1) + + delta.record_stream(stream) + + assert recorded == [ + (id(indices), stream), + (id(block_offsets), stream), + ]