Issue 1:
🐛 Bug

I want to use these two very useful functions at the same time: torch.cuda.amp and torch.utils.checkpoint.checkpoint, but I get this error:

Traceback (most recent call last):
File "/home/lart/Coding/MyProj/utils/some_tools/test_amp_ck.py", line 72, in <module>
scaler.scale(loss).backward()
File "/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/tensor.py", line 185, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/autograd/__init__.py", line 125, in backward
Variable._execution_engine.run_backward(
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same

Process finished with exit code 1

BUT IF I USE the APEX lib from NVIDIA, It works well...

To Reproduce

A demo:

The shared part:

# -*- coding: utf-8 -*-
# @Time : 2020/11/2
# @Author : Lart Pang
# @File : test_amp_ck.py
# @Project : MyProj
# @GitHub : https://github.com/lartpang
import apex
import torch
from torch.utils.checkpoint import checkpoint
from tqdm import tqdm

class MyNetWithCheckpoint(torch.nn.Module):
def __init__(self):
super(MyNetWithCheckpoint, self).__init__()
self.dummy_tensor = torch.ones(1, dtype=torch.float32,
requires_grad=True)
self.in_conv = torch.nn.Conv2d(3, 10, 3, 1, 1)
self.out_conv = torch.nn.Conv2d(10, 1, 1)

def forward(self, x):
x = checkpoint(self.ck_in_conv, x, self.dummy_tensor)
x = checkpoint(self.ck_out_conv, x)
return x

def ck_in_conv(self, x, dummy_arg=None):
assert dummy_arg is not None
x = self.in_conv(x)
return x

def ck_out_conv(self, x):
# assert dummy_arg is not None
x = self.out_conv(x)
return x

def cal_loss(pred, gt):
pred = pred.sigmoid()
loss = (pred - gt) ** 2
return loss.mean()

mynet = MyNetWithCheckpoint().cuda()

optimizer = torch.optim.SGD(mynet.parameters(), lr=0.001)
USE_AMP = True
ITER_NUM = 10

in_data = torch.randn(4, 3, 20, 20).cuda()
in_gt = torch.randint(low=0, high=2, size=(4, 3, 20, 20)).float().cuda()

When I use torch.cuda.amp:

# with torch.cuda.amp
scaler = torch.cuda.amp.GradScaler(USE_AMP)

for _ in tqdm(range(ITER_NUM), total=ITER_NUM):
with torch.cuda.amp.autocast(USE_AMP):
pred = mynet(in_data)
loss = cal_loss(pred=pred, gt=in_gt)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

It will output:

Traceback (most recent call last):
File "/home/lart/Coding/MyProj/utils/some_tools/test_amp_ck.py", line 72, in <module>
scaler.scale(loss).backward()
File "/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/tensor.py", line 185, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/autograd/__init__.py", line 125, in backward
Variable._execution_engine.run_backward(
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same

When I use APEX from NVIDIA:

# with apex.amp
mynet, optimizer = apex.amp.initialize(mynet, optimizer, opt_level="O1")

for _ in tqdm(range(ITER_NUM), total=ITER_NUM):
pred = mynet(in_data)
loss = cal_loss(pred=pred, gt=in_gt)
if USE_AMP:
with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
optimizer.zero_grad()

It will output:

Selected optimization level O1: Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled : True
opt_level : O1
cast_model_type : None
patch_torch_functions : True
keep_batchnorm_fp32 : None
master_weights : None
loss_scale : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled : True
opt_level : O1
cast_model_type : None
patch_torch_functions : True
keep_batchnorm_fp32 : None
master_weights : None
loss_scale : dynamic
100%|██████████| 10/10 [00:00<00:00, 60.96it/s]

Expected behavior

I hope these two features can be used at the same time.

Environment

➜ python collect_env.py
Collecting environment information...
PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration:
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti

Nvidia driver version: 440.100
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.0
[pip3] numpy==1.19.1
[pip3] segmentation-models-pytorch==0.1.0
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0
[conda] blas 1.0 mkl defaults
[conda] cudatoolkit 10.2.89 hfd86e86_1 defaults
[conda] efficientnet-pytorch 0.7.0 pypi_0 pypi
[conda] mkl 2020.2 256 defaults
[conda] mkl-service 2.3.0 py38he904b0f_0 defaults
[conda] mkl_fft 1.1.0 py38h23d657b_0 defaults
[conda] mkl_random 1.1.1 py38h0573a6f_0 defaults
[conda] numpy 1.19.1 py38hbc911f0_0 defaults
[conda] numpy-base 1.19.1 py38hfa32c7d_0 defaults
[conda] pytorch 1.6.0 py3.8_cuda10.2.89_cudnn7.6.5_0 pytorch
[conda] segmentation-models-pytorch 0.1.0 pypi_0 pypi
[conda] torchvision 0.7.0 py38_cu102 pytorch

Additional context

None. Thanks.

cc @mcarilli

Issue 2:
🐛 Bug

torch.utils.checkpoint.checkpoint fails when used with torch.cuda.amp on pytorch nightly. They both allow to reduce the memory usage so seem like a natural combination to use together.

To Reproduce

Here is an example program derived from https://pytorch.org/docs/master/notes/amp_examples.html#typical-mixed-precision-training

import torch
from torchvision.models import resnet18

amp_enabled = True
checkpoint_enabled = True
device = torch.device('cuda')

class ResNetWithCheckpoints(torch.nn.Module):
def __init__(self):
super().__init__()
self.base = resnet18()

def forward(self, x):
base = self.base
x = base.conv1(x)
x = base.bn1(x)
x = base.relu(x)
x = base.maxpool(x)

if checkpoint_enabled:
x = torch.utils.checkpoint.checkpoint(base.layer1, x)
x = torch.utils.checkpoint.checkpoint(base.layer2, x)
x = torch.utils.checkpoint.checkpoint(base.layer3, x)
x = torch.utils.checkpoint.checkpoint(base.layer4, x)
else:
x = base.layer1(x)
x = base.layer2(x)
x = base.layer3(x)
x = base.layer4(x)

x = base.avgpool(x)
x = torch.flatten(x, 1)
x = base.fc(x)

return x

model = ResNetWithCheckpoints()
model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()
scaler = torch.cuda.amp.GradScaler(enabled=amp_enabled)

xs = torch.randn(5, 3, 128, 128).to(device)
ys = torch.tensor(range(len(xs))).to(device)
for _ in range(20):
optimizer.zero_grad()
with torch.cuda.amp.autocast(enabled=amp_enabled):
ys_pred = model(xs)
loss = loss_fn(ys_pred, ys)
print(float(loss))
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

This script fails when both amp and checkpointing are enabled, and does not fail otherwise:

7.273253440856934
Traceback (most recent call last):
File "amp_checkpoint.py", line 51, in <module>
scaler.scale(loss).backward()
File "lib/python3.6/site-packages/torch/tensor.py", line 184, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "lib/python3.6/site-packages/torch/autograd/__init__.py", line 115, in backward
allow_unreachable=True) # allow_unreachable flag
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same

Expected behavior

I expect it to work with both amp and checkpoint enabled (I also expect memory savings from checkpointing but I'm not yet sure that my example should provide them).

Environment

Collecting environment information...
PyTorch version: 1.6.0.dev20200502+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: Could not collect

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration:
GPU 0: GeForce GTX 1050 Ti
GPU 1: GeForce RTX 2080 Ti (this is the one used here)

Nvidia driver version: 430.50
cuDNN version: /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7

Versions of relevant libraries:
[pip3] numpy==1.18.3
[pip3] torch==1.6.0.dev20200502+cu101
[pip3] torchvision==0.7.0.dev20200501+cu101
[conda] Could not collect