deepseek-ai / deepseek-ai/DeepEP
After upgrading PyTorch to 2.9.1, DeepEP’s buffer usage on GPU memory increased.
- Dominant language
- Cuda
- Stars
- 10.1k
- Forks
- 1.4k
- Avg merge
- 4d 1h
- Merged PRs (30d)
- 2
Description
Hi DeepSeek team,
I added some code in `tests/test_low_latency.py` to measure the memory consumed by Buffer:
```
import logging
logger = logging.getLogger(__name__)
gpu_id = torch.cuda.current_device()
free_mem, total_mem = torch.cuda.mem_get_info(gpu_id)
logger.warning(f'----------- Free memory of GPU {gpu_id} before allocation: {free_mem / 1024 ** 3:.2f} GB, Total memory: {total_mem / 1024 ** 3:.2f} GB')
buffer = deep_ep.Buffer(group,
num_rdma_bytes=num_rdma_bytes,
low_latency_mode=True,
num_qps_per_rank=num_experts // num_ranks,
allow_nvlink_for_low_latency_mode=not args.disable_nvlink,
explicitly_destroy=True,
allow_mnnvl=args.allow_mnnvl,
enable_shrink=args.shrink_test)
free_mem, total_mem = torch.cuda.mem_get_info(gpu_id)
logger.warning(f"----------- Free memory of GPU {gpu_id} after allocation: {free_mem / 1024 ** 3:.2f} GB, Total memory: {total_mem / 1024 ** 3:.2f} GB")
```
With PyTorch 2.8.1, the buffer consumes about 3.1 GB per GPU:
```
----------- Free memory of GPU 7 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 3 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 0 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 6 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 2 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 1 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 5 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 4 before allocation: 77.84 GB, Total memory: 79.10 GB
...
----------- Free memory of GPU 7 after allocation: 74.72 GB, Total memory: 79.10 GB
----------- Free memory of GPU 0 after allocation: 74.72 GB, Total memory: 79.10 GB
----------- Free memory of GPU 5 after allocation: 74.72 GB, Total memory: 79.10 GB
----------- Free memory of GPU 2 after allocation: 74.72 GB, Total memory: 79.10 GB
----------- Free memory of GPU 3 after allocation: 74.72 GB, Total memory: 79.10 GB
----------- Free memory of GPU 4 after allocation: 74.72 GB, Total memory: 79.10 GB
----------- Free memory of GPU 1 after allocation: 74.72 GB, Total memory: 79.10 GB
----------- Free memory of GPU 6 after allocation: 74.72 GB, Total memory: 79.10 GB
```
However, with PyTorch 2.9.1, the buffer consumption increases to about 4.8 GB per GPU:
```
----------- Free memory of GPU 6 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 1 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 3 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 7 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 0 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 2 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 4 before allocation: 77.84 GB, Total memory: 79.10 GB
----------- Free memory of GPU 5 before allocation: 77.84 GB, Total memory: 79.10 GB
...
----------- Free memory of GPU 6 after allocation: 73.00 GB, Total memory: 79.10 GB
----------- Free memory of GPU 1 after allocation: 73.00 GB, Total memory: 79.10 GB
----------- Free memory of GPU 4 after allocation: 73.00 GB, Total memory: 79.10 GB
----------- Free memory of GPU 5 after allocation: 73.00 GB, Total memory: 79.10 GB
----------- Free memory of GPU 0 after allocation: 73.00 GB, Total memory: 79.10 GB
----------- Free memory of GPU 7 after allocation: 73.00 GB, Total memory: 79.10 GB
----------- Free memory of GPU 3 after allocation: 73.00 GB, Total memory: 79.10 GB
----------- Free memory of GPU 2 after allocation: 73.00 GB, Total memory: 79.10 GB
```
This behavior is also observed in SGLang. Do you have any insights into what might be causing it?
Thanks!
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.