Fix for warning as default stream was used in enqueueV3 #3191
Conversation
github-actions
bot
added
component: core
|
@keehyuna there is some code related to cudagraphs, can you check how it handles streams and perhaps write some docs on how the runtime is suppose to do streams in all cases? |
|
There was a problem hiding this comment.
This is when problem happend. capture stream is changed to default and non default stream.
This is when torch.cuda.set_stream() is used. Non default stream is used for cuda graph/enqueueV3(). But stream is not restored after Forward()
.
This is proposed fix to keep side stream to cuda graph or enqueueV3()
| @@ -333,7 +331,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr | |||
| if (need_cudagraphs_record) { | |||
| // If cudagraphs needs to record a graph, capture the enqueueV3 call in a graph | |||
| c10::cuda::CUDAStream recording_stream = compiled_engine->engine_stream; | |||
| compiled_engine->cudagraph.capture_begin(compiled_engine->cudagraph_mempool_id); | |||
There was a problem hiding this comment.
This was reverted to fix below assert() from torch. We don't share memory across captures, I think we can use internally created pool.
https://pytorch.org/docs/stable/notes/cuda.html#graph-memory-management
File "/root/trt/TensorRT/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py", line 274, in forward
outputs: List[torch.Tensor] = torch.ops.tensorrt.execute_engine(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_ops.py", line 1113, in __call__
return self._op(*args, **(kwargs or {}))
RuntimeError: it->second->use_count > 0 INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":2056, please report a bug to PyTorch
| @@ -464,7 +461,8 @@ def cudagraphs_validate_shapes(self, inputs: Sequence[torch.Tensor]) -> bool: | |||
| if new_shape_key != self.shape_key: | |||
| logger.debug(f"Resetting Cudagraph on new shape key {new_shape_key}") | |||
| self.shape_key = new_shape_key | |||
| self.cudagraph.reset() # type: ignore | |||
| if self.cudagraph: | |||
| self.cudagraph.reset() | |||
There was a problem hiding this comment.
self.cudagraph can be None when torch_compile backend is used.
self.cudagraph is initialized when cudagraphs mode is enabled. But this init was called at compile()
https://github.com/pytorch/TensorRT/blob/main/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py#L144-L145
|
@lanluo-nvidia please cherry-pick |
Description
torch.cuda.current_stream()/c10::cuda::getCurrentCUDAStream() always returns default stream and it leads running enqueueV3() with default stream.
torch.cuda.set_stream/c10::cuda::setCurrentCUDAStream is required to set current stream when new stream is acquired from pool
Fixes #3190
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: