[ROCm] Using a more precise memory profiling #12624
Conversation
…s padding (#394) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
vllm/platforms/rocm.py
Outdated
| @@ -164,4 +164,4 @@ def get_current_memory_usage(cls, | |||
| device: Optional[torch.types.Device] = None | |||
| ) -> float: | |||
| torch.cuda.reset_peak_memory_stats(device) | |||
| return torch.cuda.max_memory_allocated(device) | |||
| return torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0] | |||
There was a problem hiding this comment.
pass device to torch.cuda.mem_get_info() ?
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
And example is at https://github.com/ROCm/vllm/blob/main/vllm/model_executor/layers/quantization/fp8.py#L357 |
To fix another ROCm regression from #11809
Currently the profiling does not take into account padded memory (torch.nn.functional.pad), so if it is used (currently used in ROCm/vllm for fp8 and MoE), the amount of available cache memory would be computed incorrectly and OOM would happen.