[Python] Applying `vmap` to a model with `torch.utils.checkpoint.checkpoint`

I need to use vmap to calculate gradient with respect to a model that uses torch.utils.checkpoint.checkpoint.

I got the following error

RuntimeError: You tried to vmap over _NoopSaveInputs, but it does not have vmap support. Please override and implement the vmap staticmethod or set generate_vmap_rule=True. For more details, please see https://pytorch.org/docs/master/notes/extending.func.html
​

After following the note given in the error, I tried adding generate_vmap_rule = True under the _NoopSaveInputs class, and I got the following error instead

RuntimeError: torch.func transforms don't yet support saved tensor hooks. Please open an issue with your use case.
​

It seems I need to define the vmap staticmethod manually, but I am not sure how to do it.

I found a github issue asking the same question. But there is still no update after a year later.

How to make vmap to work with torch.utils.checkpoint.checkpoint?

Here's the minimal code to reproduce my problem

import torch
import torch.nn as nn
import torch.optim as optim
import random
from torch.func import vmap
from torch.utils.checkpoint import checkpoint
# from torch.autograd.function import once_differentiable

# Set a fixed seed for reproducibility
seed = 42
torch.manual_seed(seed)
random.seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # For multi-GPU setups
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# Define a simple Transformer model
class SimpleTransformer(nn.Module):
def __init__(self, input_dim, model_dim, num_heads, num_layers, output_dim):
super(SimpleTransformer, self).__init__()
self.embedding = nn.Embedding(input_dim, model_dim)
self.transformer = nn.Transformer(d_model=model_dim, nhead=num_heads, num_encoder_layers=num_layers)
self.fc_out = nn.Linear(model_dim, output_dim)

def forward(self, src, tgt):
src = self.embedding(src)
tgt = self.embedding(tgt)
# output = self.transformer(src, tgt)
output = checkpoint(self.transformer, src, tgt, use_reentrant=False)
return self.fc_out(output)

def compute_loss(model,
weights,
buffers,
src,
tgt,
):

# Forward pass
output = torch.func.functional_call(
model, (weights, buffers), args=(src.unsqueeze(1), tgt[:-1].unsqueeze(1))
)

# Reshape output and target for loss calculation
output = output.view(-1, output_dim)
tgt = tgt[1:].view(-1) # Shift target sequence by one
loss = criterion_mean(output, tgt)
print(loss.shape)
return loss

# Hyperparameters
input_dim = 1000 # Vocabulary size
model_dim = 512 # Embedding dimension
num_heads = 8 # Number of attention heads
num_layers = 6 # Number of transformer layers
output_dim = 1000 # Output vocabulary size
seq_length = 10 # Sequence length

# Initialize model, loss function, and optimizer
model = SimpleTransformer(input_dim, model_dim, num_heads, num_layers, output_dim)
criterion_mean = nn.CrossEntropyLoss()
criterion_raw = nn.CrossEntropyLoss(reduction='none')
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Dummy data (batch_size = 2, seq_length = 10)
src = torch.randint(0, input_dim, (seq_length, 2)) # Source sequence
tgt_raw = torch.randint(0, output_dim, (seq_length, 2)) # Target sequence

# Print src and tgt to verify they are the same each time
print("Source sequence (src):")
print("Target sequence (tgt):")

# Forward pass
output = model(src, tgt_raw[:-1, :]) # Exclude the last token in the target sequence

# Reshape output and target for loss calculation
output = output.view(-1, output_dim)
tgt = tgt_raw[1:, :].view(-1) # Shift target sequence by one

weights = dict(model.named_parameters())
buffers = dict(model.named_buffers())

grads_fn = torch.func.grad(compute_loss, has_aux=False, argnums=1)
gs = vmap(grads_fn,
in_dims=(None, None, None, 1, 1),
randomness='different')(
model,
weights,
buffers,
src,
tgt_raw
)

Continue reading...