Calling backward with create_graph on the output of a DistributedDataParallel throws error

[carlosgmartin]

🐛 Bug

Calling backward with create_graph=True on the output of a DistributedDataParallel throws a RuntimeError.

To Reproduce

import torch
from torch import nn
from torch.distributed import init_process_group
from torch.nn.parallel import DistributedDataParallel as DDP

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.p = nn.Parameter(torch.tensor(1.))
    def forward(self):
        return self.p.pow(2)

model = Model()
init_process_group(
    'gloo',
    'tcp://localhost:12355',
    rank=0,
    world_size=1,
)
ddp_model = DDP(model)
ddp_model().backward(create_graph=True)

Traceback (most recent call last):
  File "test.py", line 21, in <module>
    ddp_model().backward(create_graph=True)
  File "/usr/local/lib/python3.8/site-packages/torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/usr/local/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
    Variable._execution_engine.run_backward(
RuntimeError: mul(): functions with out=... arguments don't support automatic differentiation, but one of the arguments requires grad.

Expected behavior

No error.

Environment

Collecting environment information...
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 11.5.2 (x86_64)
GCC version: Could not collect
Clang version: 12.0.5 (clang-1205.0.22.9)
CMake version: version 3.20.1
Libc version: N/A

Python version: 3.8.8 (default, Feb 27 2021, 02:19:17)  [Clang 12.0.0 (clang-1200.0.32.29)] (64-bit runtime)
Python platform: macOS-11.5.2-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.812
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] pytorch-lightning==1.2.10
[pip3] torch==1.9.0
[pip3] torchaudio==0.9.0.dev20210313
[pip3] torchmetrics==0.2.0
[pip3] torchvision==0.10.0
[conda] Could not collect

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23

[rohan-varma]

Verified I can repro the issue. The error is coming from at::mul_out here: https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L382. It doesn't occur if create_graph=False which is the default.

@albanD do you have any suggestions for what create_graph changes to result in this error in at::mul_out? My thought is that it is because grad ends up requiring grad itself because create_graph is for higher order derivatives? (Edit: yep, seems like this is the case)

[albanD]

Yes, the out variants are never differentiable. So if you try to use create_graph=True (which differentiates through the backward), that will fail.

[mrshenli]

closing this issue since this is expected behavior.

[carlosgmartin]

@mrshenli What do you mean "this is expected behavior"? What's the solution?

[mrshenli]

What do you mean "this is expected behavior"?

I was referring to DDP does not work with higher-order derivatives is an expected behavior today.

Re-opening this issue in case the intention was not calculating higher-order derivatives. Could you please share a bit more on the motivation for setting create_graph=True? Are you going to calculate grad on grads?

[carlosgmartin]

@mrshenli The motivation was indeed to calculate higher-order derivatives. If DDP doesn't support that, a more informative error message would be helpful.

What would it take for DDP to support higher-order derivatives?

[mrshenli]

What would it take for DDP to support higher-order derivatives?

Today's DDP expects interleaving forward-backward passes. For higher-order derivatives, I assume you need to calculate the loss based on the grads and call a 2nd backward? If so, this would mess up DDP's internal states. We probably need to disable DDP's autograd hooks in the first backward, and enable that in the 2nd backward invocation. DistributedDataParallel.no_sync() won't work here, as it requires another forward after exiting the context.

One hacky and dangerous solution might be exposing a private API to toggle this flag. Say, with ddp._ignore_hooks():.

pytorch/torch/csrc/distributed/c10d/reducer.cpp

Line 571 in 4a07761

if (!expect_autograd_hooks_) {

Hey @carlosgmartin, could you pease help us understand what are the applications that require higher-order derivatives? Thanks!