deepseek-ai / deepseek-ai/DeepEP

After upgrading PyTorch to 2.9.1, DeepEP’s buffer usage on GPU memory increased.

Open
#510 2 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.