Repository: pytorch
Issue: [RFC] Extend Autocast to CPU/CUDA with BF16 data type
Labels: feature,triaged,module: bfloat16,module: amp (automated mixed precision)

## 🚀 Feature
This RFC proposes to extend Pytorch Autocast to support datatype of BFloat16 on CPU device. Autocast provides easy-to-use methods for mixed precision which can use the low precision datatype in a convenient way. Currently, Autocast already support CUDA device with datatype of fp16.
To align with existing Autocast mechanism, below are the proposed changes at the high level:
### Frontend Changes
We propose to change current Autocast API from `torch.cuda.amp.autocast()` into `torch.autocast(enabled=True, dtype= torch.bfloat16, device='cpu')`:
* Move the API from CUDA to a more general namespace.
* Add new data type parameter in the API which is the target data type to be casted at Runtime.
* Add new device type parameter in the API since we extend Autocast to different devices.

### Backend Changes:
To extend the Autocast design with BFloat16 on CPU device and to adapt the frontend changes, the proposed backend changes are:
* Change the cast policy name of `fp16` to `user_defined_dtype`, since we propose to add the new parameter of `dtype` in the frontend API.
* Consolidate OP list under each cast policy: `user_defined_dtype`, fp32, fall through, fp32_set_opt_dtype, fp32_append_dtype, promote.

## Motivation
With Intel's next XEON server product, new instruction set of BF16 TMUL will be added. It's attractive to users getting their application's performance improvement with this new BF16 instruction. However, to get the performance boost with bf16 instruction, user still needs to convert the model to bfloat16 data type explicitly in current phase. It's not user-friendly. With the support of Autocast in CPU device, it will significantly improve the out-of-box user experience.

## Pitch
Extend the Pytorch Autocast design to support BFloat16 on the CPU device by generalizing the device and the data type supported by Autocast.

## Additional context
### Design Options: Share or Not Share Dispatch Key
Here the question is if we can share the dispatch key `Autocast` between different devices in the backend design. Under the current Pytorch design, sharing the dispatch key means sharing the operation's cast policy between different devices. However, the operation's belonging cast policy among devices are not fully aligned naturally.
Take CUDA FP16 and CPU BF16 as example:
| Proposed Cast Policy | Current CUDA OP List (FP16 AMP) | CPU OP List (BF16 AMP)
|  ----  | ----  | ----  
| user_defined_dtype | conv1D/2D/3D, matmul, linear, mm, bmm, `lstm_cell`, `gru_cell` etc. | ops support BF16: conv1D/2D/3D, matmul, linear, mm, bmm etc.
| fp32 | acos, asin, log, tan, gelu, `layer_norm`, group_norm, etc. | ops support FP32 only. (`lstm_cell`, `gru_cell` for CPU).
| fall through | ops support FP16. | ops support BF16. (`layer_norm` for CPU).
| fp32_set_opt_dtype | prod, softmax. int,  sum, log_softmax.int, etc. | keep this list.
| fp32_append_dtype | norm.scalar, etc. | keep this list.
| promote | cat, stack, tensordot etc. | keep this list.

There are two reasons which cause the difference:
* CUDA and CPU have different low precision op support status in current phase. They are developing continually (e.g., `lstm_cell`, `gru_cell`).  
* Some ops support bf16 but not fp16 inherently (e.g., `layer_norm`).

#### Design option1: Dedicated Dispatch Key `AutocastCPU`
Using a dedicated Dispatch Key `AutocastCPU` for CPU and `AutocastCUDA` for CUDA, CPU/CUDA OP list can be managed separately.
* Pros:
    * Much easier to enable and maintain for new devices.
    * No changes to CUDA Autocast mechanism.
    * No additional runtime dispatch cost.
* Cons:
    * More device-specific dispatch keys

#### Design option2: Share Dispatch Key `Autocast`
Sharing the Dispatch Key `Autocast` between different devices, we must consolidate CUDA OP list and CPU OP list. As we discussed previously, some operations (e.g., `lstm_cell`, `gru_cell`, `layer_norm`) have different belonging cast policy between CPU and CUDA. The solution to handle these operations without performance compromise is to move all these ops into new cast policy named as `runtime`. Then we decide the cast dtype of these operations at runtime. The details steps are:
1.	Classify all these operations into a new `runtime` cast policy.
2.	Generate the partial specialization wrap function for each of these ops under `runtime` cast policy.
3.	Inside the op's specific wrap function, we combine the information of runtime device for this op, dtype from frontend API and op specific capability under this runtime device to decide the final cast dtype.
* Pros:
    * Save device-specific dispatch keys.
* Cons:
    * More enabling and maintenance effort: Need to consider both CPU and CUDA, and when more devices added in.
    * Non-trivial changes to CUDA Autocast.
    * Added complexity for additional runtime device and data type dispatch in Autocast.

Comparing the Pros and Cons between these 2 designs, we prefer the Design Option1: Dedicated Dispatch Key `AutocastCPU`.

### Design Options: Add or Not Add `Device` Parameter
Here is to discuss whether we need to add the parameter `device` in the frontend API. Let us look at the case if we don’t add this `device` parameter. Then the question is when user enables Autocast, which dispatch Key should be included in the `local_dispatch_key_set`. It’s unknown since we have no information about the model or tensors when enabling Autocast under origin design.
There are two design options to this question:
#### Design option1: Add Device Specific Parameter
The design option1 is adding the device specific parameter in the frontend API. Then we can easily include the corresponding dispatch key into the `local_dispatch_key_set`.
* Pros:
    * Straight forward design without extra dispatch.
* Cons:
    * One more frontend Parameter in API.
#### Design option2: Add Extra Dispatch Layer
The design option2 is adding an extra dispatch layer with device independent dispatch key (such as Autocast) as first dispatch layer. Inside it, all operations will fall back into a general wrap function. Inside this wrap function, we will check each tensor’s backend, and include the device-specific dispatch key (AutocastX) in next dispatch layer. Here is the sample code of the first dispatch layer:
```
void generic_wrapper_fallback(const OperatorHandle& op, torch::jit::Stack* stack) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(DispatchKey::Autocast);
  auto num_arguments = op.schema().arguments().size();
  auto args = torch::jit::pop(*stack, num_arguments);
  bool use_cuda = false;
  for (size_t i = 0; i < num_arguments; i++) {
    if(args[i].isTensor()){
      Tensor impl = args[i].toTensor();
      torch::jit::push(*stack, std::move(impl));
      use_cuda = check_cuda(impl);
    }else if(args[i].isTensorList()){
      c10::List<at::Tensor> impl = args[i].toTensorList();
      torch::jit::push(*stack, std::move(impl));
      use_cuda = check_cuda(impl);
    }else{
      torch::jit::push(*stack, std::move(args[i]));
      use_cuda = check_cuda(args[i]);
    }
  }
  if(use_cuda){
     c10::impl::IncludeDispatchKeyGuard autocast_cuda(DispatchKey::AutocastCUDA);
     op.callBoxed(stack);
  }else{
    c10::impl::IncludeDispatchKeyGuard autocast_cpu(DispatchKey::AutocastCPU);
    op.callBoxed(stack);
  }
  return;
}
TORCH_LIBRARY_IMPL(_, Autocast, m) {
  m.fallback(torch::CppFunction::makeFromBoxedFunction<&generic_wrapper_fallback>());
}
```
* Pros:
    * Clean frontend API with less parameters.
* Cons:
    * Extra dispatch layer with boxed API which is not performance friendly.

Comparing the Pros and Cons between these 2 design options, we prefer the Design Option1: add the device parameter into the API.

#### Design option3: Open to discuss
Enhancing existing dispatch mechanism to avoid dispatch overhead from design option2, the detail implementation of this design option is still open to discuss.

### Details about the Use Cases
The new design supports all existing Autocast use cases. Here to elaborate some examples impacted by the new design.
#### Example1: Default Parameters
In the origin design, the default behavior of Autocast is to enable Autocast with FP16 dtype.
To follow this behavior, the new API default parameters are: `enabled=True`, `dtype=torch.float16`, `device='cuda'`, so if you are using Autocast with default parameter as:
```
for input, target in dataset:
  with torch.autocast():
      output = model(input)
      loss = loss_fn(output, target)
```
You are enabling Autocast with datatype of float16, and the op kernels registered under AutocastCUDA Dispatch Key will be dispatched.

#### Example2: Enable-Enable Nested Pattern with different parameters
We don’t support the nested Autocast with hybrid parameter use case such as:
```
for input, target in dataset:
  with torch.autocast(enabled=True, dtype=torch.bfloat16, device='cpu'):
      output1 = model1(input)
      model2 = model2.cuda()
      with torch.autocast(enabled=True, dtype=torch.float16, device='cuda'):
          output2 = model2(output1)
      output3 = model3(output2)
      loss = loss_fn(output, target)
```

#### Example3: Enable-Disable Nested Pattern
If you are using the nested Autocast with enable-disable pattern:
```
for input, target in dataset:
  with torch.autocast(enabled=True, dtype=torch.bfloat16, device='cpu'):
      output1 = model1(input)
      with torch.autocast(enabled=False):
          output2 = model2(output1)
     output3 = model3(output2)
      loss = loss_fn(output, target)
```
* The model1 will run under Autocast with dtype of torch.bfloat16 on CPU device.
* The model2 under the disable context means you want to disable Autocast for every device.
* The model3 is still under the first Autocast Context which will keep enabling Autocast with dtype of torch.bfloat16 on CPU device.


cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @mcarilli @ptrblck