wence-
commented
10 months ago torch.cuda.change_current_allocator(rmm_torch_allocator) pool = rmm.mr.PoolMemoryResource( rmm.mr.CudaMemoryResource(), initial_pool_size=2**30, ) device = (int(os.environ['LOCAL_RANK'])) rmm.mr.set_per_device_resource(device, pool)
This is almost correct. However, when creating a memory resource, you should do so with the target device active.
That is, you must call cudaSetDevice(device) before creating the pool and then calling set_per_device_resource.
This can be done like so:
import rmm
device = (int(os.environ['LOCAL_RANK']))
rmm._cuda.gpu.setDevice(device)
pool = rmm.mr.PoolMemoryResource(...)
rmm.mr.set_per_device_resource(device, pool)Since this is such a common pattern, the top-level rmm.reinitialize has some logic to handle this for you:
import rmm
device = int(os.environ["LOCAL_RANK"])
rmm.reinitialize(devices=device, pool_allocator=True, initial_pool_size=...)This doesn't have quite as much flexibility on the set up of the allocator, but if you just need a pool on top of a cuda memory resource then it works fine.
We could add an interface whereby you provide a zero-argument callback to construct the pool (and rmm.reinitialize would arrange to call it with the correct device active), but we haven't needed it so far.
li-yi-dong
harrism
Describe the bug I tried to use RMM with PyTorch. I launch my task with torchrun and set the rmm.mr for each device at the very beginning.
But each process occupies a chunk of memory on GPU0 like
Steps/Code to reproduce bug
Expected behavior I expected each process launched by torchrun only uses the memory on the GPU assigned by
LOCAL_RANKEnvironment details (please complete the following information): I'm using RMM v23.10.00 Here is the output of the print_env.sh
Additional context Add any other context about the problem here.