Allow to disable polling for CUDA synchronization

[f0k]

🚀 Feature

When using CUDA, by default, synchronization of the host and GPU works through a tight loop that polls a particular memory location, running a CPU core at 100%. CUDA provides two alternative modes: (1) a tight loop that yields to other threads in between, and (2) synchronization via an interrupt (translated into an event by the driver). PyTorch should provide some way to select either of the three modes.

Motivation

Burning a CPU core at 100% is not very green. It doesn't save a lot of time either, especially when the GPU has to process large workloads in between synchronizations. And it makes it harder to see how much the CPU is actually utilized.

Pitch

On the backend, the way to select the synchronization mode is to run either of:

• cudaSetDeviceFlags(cudaDeviceScheduleAuto)
• cudaSetDeviceFlags(cudaDeviceScheduleSpin)
• cudaSetDeviceFlags(cudaDeviceScheduleYield)
• cudaSetDeviceFlags(cudaDeviceScheduleBlockingSync)

This can be done right after the first cudaSetDevice() call, but before the device is used (otherwise it's a cudaErrorSetOnActiveProcess error) -- this may be hard to guarantee. However, it can also be done before any cudaSetDevice() call, in that case it will set the default for any future activated devices.

So I'd suggest to add a function such as torch.cuda.set_sync_mode() that takes a string "auto", "spin", "yield", or "block" and directly issues the corresponding cudaSetDeviceFlags() call. It's up to the user to run this early enough then.

Disclaimer: I haven't looked into what would happen with multiprocessing.

Alternatives

The main alternative for implementing this would be figuring out where to safely place a cudaSetDeviceFlags() call in the device managing code, and have it read some global variables at that time. We could have a torch.cuda.sync_mode variable then, or still guard it by a setter function.

It's also possible to solve this in user code doing something like:

import ctypes
ctypes.CDLL('libcudart.so').cudaSetDeviceFlags(4)  # = cudaDeviceScheduleBlockingSync

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @ngimel @VitalyFedyunin @mruberry

[f0k]

However, it can also be done before any cudaSetDevice() call, in that case it will set the default for any future activated devices.

Or not. The following does not set the device flags as intended:

import torch
import ctypes
c = ctypes.CDLL('libcudart.so')

c.cudaSetDeviceFlags(4)

torch.cuda.set_device(0)
torch.zeros(1).cuda()

i = ctypes.c_uint()
c.cudaGetDeviceFlags(ctypes.byref(i))
print(i)  # prints 0

The following does work (sort of, it adds bit 8 although I didn't specify it):

import torch
import ctypes
c = ctypes.CDLL('libcudart.so')

torch.cuda.set_device(0)

c.cudaSetDeviceFlags(4)

torch.zeros(1).cuda()

i = ctypes.c_uint()
c.cudaGetDeviceFlags(ctypes.byref(i))
print(i)  # prints 12

So something in torch.cuda.set_device() seems to reset the device flags. Assuming it's hard to figure out what that is, an alternative solution would be to add the synchronization mode as an optional keyword argument to torch.cuda.set_device(), since it seems it needs to be called anyway to be able to provide device flags.

[smessmer]

This sounds reasonable to me. cc @ezyang for confirmation.

[mruberry]

@csarofeen as an fyi

[soumith]

@mruberry @csarofeen @ngimel please approve or deny, as this has more to do with CUDA and you are the owners of that.

[mruberry]

@f0k

This is a cool feature idea and the proposal is clear and well-written.

If we did have this feature, what would you recommend the default mode be and why?

Further, what guidance would you give to users about which mode they should set? How much of a difference would they experience?

[ezyang]

I quickly checked and there are no occurrences of cudaSetDeviceFlags in the codebase today. So I'm not sure why it's getting reset. @mruberry I'd assume the default mode would be what we do today in PyTorch.

If there's a way to get the first code sample to work, then we can just expose a very thin wrapper on the CUDA API, similar to some existing wrappers we have. But if it's necessary to set this after we change devices, we might have to work a bit harder to give a good UX.

[f0k]

If we did have this feature, what would you recommend the default mode be and why?

As @ezyang said, the default should be not to call cudaSetDeviceFlags, or to call it with 0 (automatic scheduling), to not change existing behavior.

Further, what guidance would you give to users about which mode they should set? How much of a difference would they experience?

I'd personally set it to 4 (blocking sync). It can be minimally slower, and it can reduce CPU use significantly, but both depends on the experiment. In my current case, it's reducing CPU usage by 30% and runs ~0.5% slower (that's 18 seconds per hour), which I'm willing to accept to reduce energy consumption.

I quickly checked and there are no occurrences of cudaSetDeviceFlags in the codebase today. So I'm not sure why it's getting reset.

Yes, I had checked as well and am a bit clueless. For the full picture, I originally added some more gets and sets:

import ctypes
import torch

c = ctypes.CDLL('libcudart.so')

def set_flags(f):
    return c.cudaSetDeviceFlags(f)

def get_flags():
    i = ctypes.c_uint()
    c.cudaGetDeviceFlags(ctypes.byref(i))
    return i.value

print('Before anything:', get_flags())
print('Setting to 12:', set_flags(12))
print('Retrieving:', get_flags())

torch.cuda.set_device(0)

print('After set_device()', get_flags())
print('Setting to 2:', set_flags(2))
print('Retrieving:', get_flags())

torch.zeros(10).cuda()
print('After using device()', get_flags())
print('Setting to 1:', set_flags(1))
print('Retrieving:', get_flags())

This prints:

Before anything: 8
Setting to 12: 0
Retrieving: 12
After set_device() 0
Setting to 2: 0
Retrieving: 0
After using device() 10
Setting to 1: 36
Retrieving: 10

So before doing anything, we can set and get the flags as we want, but they will be ignored. After calling cuda.set_device(), we can set the flags, but not retrieve them. After using the device, the flags are what we set after the set_device() call, with 8 added to them, and we cannot change them any more. The initial flags (before doing anything) are always ignored, also when commenting out the second cudaSetDeviceFlags() call.

I then went back to the first code example and changed it to not use torch at all:

import ctypes
c = ctypes.CDLL('libcudart.so')

c.cudaSetDeviceFlags(4)
c.cudaSetDevice(0)
c.cudaMalloc(ctypes.byref(ctypes.c_void_p()), 1024)

i = ctypes.c_uint()
c.cudaGetDeviceFlags(ctypes.byref(i))
print(i)  # prints 12

This works. This means torch.cuda.set_device(0) must be resetting the flags somewhere. Note that you can also change the flags on a per-context basis using the driver API instead of the runtime API, but I didn't find a cuCtxCreate or cuDevicePrimaryCtxSetFlags in the codebase either. I don't know how PyTorch manages its contexts, and what other ways there are to change the flags.

But if it's necessary to set this after we change devices, we might have to work a bit harder to give a good UX.

Calling torch.cuda._lazy_init() also resets the flags, but still allows me to change them, similarly to after the torch.cuda.set_device() call. So I guess we could have torch.cuda.set_sync_mode() record the desired synchronization mode somewhere, and then have torch.cuda._lazy_init() read and apply it if it's not the default. We might be able to just use the torch.cuda._lazy_call() machinery for this.
/edit: But the _lazy_call() would not be rerun for forks. So it should probably set a global variable somewhere instead.
/edit /edit: Doing this at the end of _lazy_init() will not set the flags for all GPUs. Solving this correctly will need some better knowledge of the internals.

[ezyang]

@f0k I'm as clueless as you. Can we maybe breakpoint the relevant driver/runtime api functions in gdb and see which ones are getting hit?