Accelerator abstraction

[tjruwase]

This PR is to help reduce the burden of supporting Deep Learning accelerators in DeepSpeed. We expect at least two concrete benefits from:

1. Adding and maintaining accelerator logic will require minimal changes to DeepSpeed code.
2. DeepSpeed development will involve minimal changes to accelerator code.

This PR is heavily influenced by #2221. The main differences from #2221 are:

1. Introduction of an abstract DeepSpeedAccelerator class that encapsulates accelerator logic.
2. A single global DeepSpeedAcccelerator object that is lazily initialized based on available torch modules (e.g., torch.cuda, torch.xpu) and will be used throughout DeepSpeed code to access accelerator functionalities.
3. Concrete implementation of DeepSpeedAccelerator (e.g., for CUDA or XPU) will be external modules that will be imported by DeepSpeed during initialization.
4. Eliminate control flow related to accelerator selection (e.g., if ... else) outside of initialization code.

This PR includes a proposed implementation of a CUDA accelerator. Below is output of a basic test on V100x32GB:

[delock]

For device abstraction, one way is make literal_device() an interface of accelerator. During porting I found there is still need to use literal string to mark device in many places.

def literal_device(self, device_index=None):
    if device_index == None:
        return 'cuda'
    return 'cuda:{}'.format(device_index)

https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/activation_checkpointing/checkpointing.py#L557
-->
cuda_device = get_accelerator().literal_device(get_accelerator().current_device())

also it might make sense to have current_literal_device() for the sake of code readability.

We are testing accelerator abstraction interfaces in our environment. When we finish we will push to #2221. We are looking forward to graduately converging with #2320 .

[tjruwase]

For device abstraction, one way is make literal_device() an interface of accelerator. During porting I found there is still need to use literal string to mark device in many places.

In general, it is alright to extend the interface if you find compelling cases in your testing. However, I think the literal_device() scenario below could be handled differently as I will explain in a separate post.

We are testing accelerator abstraction interfaces in our environment. When we finish we will push to #2221. We are looking forward to graduately converging with #2320 .

Thanks, that would be great. The primary goal of this PR is to experiment and foster our alignment. So, it is great if your #2221 can subsume it since you are able to test the XPU accelerator.

[tjruwase]

def literal_device(self, device_index=None):
    if device_index == None:
        return 'cuda'
    return 'cuda:{}'.format(device_index)

https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/activation_checkpointing/checkpointing.py#L557 --> cuda_device = get_accelerator().literal_device(get_accelerator().current_device())

I think in this case, the deepspeed code should be changed so that Stream() is called without device argument so that it defaults to None. This way, each concrete accelerator can handle as appropriate. For CUDA, this will be current device, but perhaps some other accelerator may want to do things differently. What do you think?

also it might make sense to have current_literal_device() for the sake of code readability.

Can you please share a compelling case for this? I think that integer indices are sufficient to identify devices, perhaps I am missing some scenarios that you have encountered, so please correct me. My main goal is to minimize the API surface as much as possible to simplify development and testing. As an alternative, would it be sufficient to have an API that returns a human readable name for a device index. For example, the cuda implementation of such a method could be:

   def readable_device_name(self, device_index=None): 
         return f'cuda:{torch.cuda.current_device()}' if device_index is None else f'{cuda:{device_index}'

[delock]

The necessarty of a literal device string or readable device name lies behind that device name is used in pytorch tensor function for accelerators other than cuda.

For example, if we have a tensor T
T.to(1) would copy T to cuda device 1
If we want to copy it to an xpu device, we need to write
T.to('xpu:1')
We can also specify cuda device explicity by a literal device string
T.to('cuda:1')

pytorch also use device name 'cuda' to specify current cuda device, in that case, the device name is 'cuda' or 'xpu' instead of None.

In the following function that take cuda_device as parameter, tensor interface are used. Thats the reason we have to replace device_index by device name in places where device id are used in tensor interface.

DeepSpeed/deepspeed/runtime/utils.py

Line 122 in eed4032

def copy_to_device(item, device, criterion_func):

There is choice of when the conversion from device_index to device name. In #2221, we choose to pick device name as the default form, and pass device name instead of device index in deepspeed function parameters.

Alternatively, if device index are passed in all deepspeed function, then conversion to device name will be needed when pytorch tensor interface are called, or any other place where device name are needed.

[delock]

The necessarty of a literal device string or readable device name lies behind that device name is used in pytorch tensor function for accelerators other than cuda.

For example, if we have a tensor T T.to(1) would copy T to cuda device 1 If we want to copy it to an xpu device, we need to write T.to('xpu:1') We can also specify cuda device explicity by a literal device string T.to('cuda:1')

pytorch also use device name 'cuda' to specify current cuda device, in that case, the device name is 'cuda' or 'xpu' instead of None.

In the following function that take cuda_device as parameter, tensor interface are used. Thats the reason we have to replace device_index by device name in places where device id are used in tensor interface.

DeepSpeed/deepspeed/runtime/utils.py

Line 122 in eed4032

def copy_to_device(item, device, criterion_func):

There is choice of when the conversion from device_index to device name. In #2221, we choose to pick device name as the default form, and pass device name instead of device index in deepspeed function parameters.

Alternatively, if device index are passed in all deepspeed function, then conversion to device name will be needed when pytorch tensor interface are called, or any other place where device name are needed.

In deepspeed code, https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/activation_checkpointing/checkpointing.py#L115 is the only place that current_device() can be used for other accelerators. In all other places where current_device is used, an alternative name like 'current_device_name()' should be used to return a device name.

[tjruwase]

The necessarty of a literal device string or readable device name lies behind that device name is used in pytorch tensor function for accelerators other than cuda.

Thanks for the explanation; I am convinced. I had incorrectly assumed that accelerator.device() could be used in the tensor interfaces in place of torch.device. But now, I find out that torch.cuda.device is not a subclass of torch.device.

>>> isinstance(torch.cuda.device(0), torch.device)
False

Hopefully in future the pytorch tensor interface could remove the 'cuda' hardcoding to simplify things. For now, I have added device_name() that returns the literal string. Will that work? Please feel free to modify as needed. Thanks!

[delock]

range_pop does not take msg as argument, only range_push does.
https://pytorch.org/docs/stable/generated/torch.cuda.nvtx.range_pop.html

[delock]

Our internal test shows this check 'is_available()' breaks unit test. Call to is_available would initialize cuda too early and cause initialization error.

specifically, this test is broken with this assertion.
pytest -k "test_ckpt_arg_none" test_activation_checkpointing.py

[tjruwase]

Can you please share a stack trace for this? Thanks!

[delock]

Here is the stack trace of this error in our CUDA environment

$ pytest -k "test_ckpt_arg_none" test_activation_checkpointing.py
......
Worker 0 exited with code 1
----------------------------- Captured stdout call -----------------------------
[2022-10-10 10:27:14,784] [INFO] [comm.py:639:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
------------------------------------------ Captured stderr call ------------------
Process Process-1:
Traceback (most recent call last):
  File "/home/gma/anaconda3/envs/ds/lib/python3.9/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/gma/anaconda3/envs/ds/lib/python3.9/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/gma/mingzhil/CI/frameworks.ai.benchmarking.other.deepspeed/tests/unit/common.py", line 250, in dist_init
    dist.barrier()
  File "/home/gma/mll/CI/frameworks.ai.benchmarking.other.deepspeed/deepspeed/comm/comm.py", line 127, in log_wrapper
    return func(*args, **kwargs)
  File "/home/gma/mll/CI/frameworks.ai.benchmarking.other.deepspeed/deepspeed/comm/comm.py", line 459, in barrier
    return cdb.barrier()
  File "/home/gma/mll/CI/frameworks.ai.benchmarking.other.deepspeed/deepspeed/comm/torch.py", line 153, in barrier
    return torch.distributed.barrier()
  File "/home/gma/anaconda3/envs/ds/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2776, in barrier
    work = default_pg.barrier(opts=opts)
RuntimeError: CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
......
===================== short test summary info ==========================
FAILED test_activation_checkpointing.py::test_ckpt_arg_none
=================== 1 failed, 21 deselected, 2 warnings in 1.06s ========================

[tjruwase]

@delock, since you have merged this into your PR, I feel this one has already achieved its goal of helping us align on core design issues. So, I propose to abandon this and focus on yours. What do you think?

[delock]

@tjruwase
Sure, let me merge your latest interface changes then we can do co-developing.

[delock]

@tjruwase I found there is one interface semantic difference between #2221 and #2320.

In #2320, the default behavior of device_name() is get current device name, and name() is use to get device name without index.
device_name() --> current device name such as 'cuda:0'
name()--> device name without index such as 'cuda'

In #2221, the default behavior of device_name() is get device name without index, and current_device_name() is to get current device name
device_name()--> device name without index such as 'cuda'
current_device_name()--> current device name such as 'cuda:0'

I think the variant in #2221 is better, because

1. the original code to use current device is torch.cuda.current_device(), use current_device_name() looks more natural modification.
2. 'cuda' and 'cuda:index' are both valid way to specify device name in pytorch, so let them share the same interface device_name() looks natural.

My suggestion is remove name() interface and keep interface device_name() and current_device_name().

[tjruwase]

My suggestion is remove name() interface and keep interface device_name() and current_device_name().

This sounds good with me.