reentrant torch.utils.checkpoint does not work with NamedTuple outputs

[rohan-varma]

🐛 Describe the bug

The reentrant (default) version of torch.utils.checkpoint does not work with NamedTuple outputs, it instead returns back a regular tuple so the outputs cannot be accessed via their original names. For a repro, see:

import torch
from torch.utils.checkpoint import checkpoint

from collections import namedtuple

Tup = namedtuple("Tup", "a b c")

tup = Tup(torch.ones(10, requires_grad=True), torch.ones(10, requires_grad=True), torch.ones(10, requires_grad=True))

def foo(tup):
	return Tup(tup.a + tup.b, tup.b, tup.a + tup.b + tup.c)

import pdb ; pdb.set_trace()
out = checkpoint(foo, tup, use_reentrant=False)

Versions

main

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501

[rohan-varma]

Upon further debugging, it seems that the type goes from NamedTuple -> Tuple after the CheckpointFunction.apply, i.e. in CheckpointFunction.forward the type that's returned is still correct.

[rohan-varma]

It seems that general torch.autograd.Function doesn't support NamedTuple:

import torch
from torch.utils.checkpoint import checkpoint

from collections import namedtuple

Tup = namedtuple("Tup", "a b c")

tup = Tup(torch.ones(10, requires_grad=True), torch.ones(10, requires_grad=True), torch.ones(10, requires_grad=True))

class MyOp(torch.autograd.Function):

	@staticmethod
	def forward(ctx, a, b, c):
		return Tup(a, b, c)

	@staticmethod
	def backward(ctx, a, b, c):
		return None

def foo(tup):
	print(f"input: got {tup}")
	return Tup(tup.a + tup.b, tup.b, tup.a + tup.b + tup.c)

# import pdb ; pdb.set_trace()
# out = checkpoint(foo, tup, use_reentrant=True)

op = MyOp.apply
a= torch.randn(5, 5, requires_grad=True)
out = op(a, a, a)
import pdb ; pdb.set_trace()
# import pdb ; pdb.set_trace()

in this case, out is not a namedtuple. @albanD @ezyang is this expected for custom autograd functions?

[ezyang]

It's not surprising to me. We could teach custom autograd about named tuple I suppose, it is on the easier side; but for arbitrary nesting you will need to manually flatten with pytrees.