torch.func.jvp fails with BERT training
Labels: module: autograd, triaged, actionable, module: forward ad, module: functorch

As the title suggests just as a fun experiment I am trying to training a huggingface BERT model using forward gradient. I am getting this error in jvp computation.

> `RuntimeError: wrapper->level().value() <= current_level INTERNAL ASSERT FAILED at "../aten/src/ATen/functorch/ADInterpreters.cpp":39, please report a bug to PyTorch. escaped?`

Here is the minimal implementation

```
import torch
from torch.nn.utils import parameters_to_vector, vector_to_parameters
from transformers import BertTokenizer, BertForMaskedLM, BatchEncoding
from torch.utils.data import DataLoader
from datasets import load_dataset
import torch.optim as optim
import gc
from torch.func import jvp
from functools import partial


class ForwardADTraining:
    def __init__(self, model, dataloader, criterion, device):
        self.model = model.to(device)
        self.dataloader = dataloader
        self.criterion = criterion
        self.device = device

    def train_one_epoch(self):
        print("Training started")
        with open('log.txt', 'w') as f:
            for i, data in enumerate(self.dataloader, 0):
                inputs = {k: v.to(self.device) for k, v in data.items()}
                labels = data['labels'].to(self.device)

                # Get the model parameters
                model_params = list(self.model.parameters())
                original_params = tuple([p.detach().clone() for p in model_params])

                # Generate a random tangent vector
                tangents_vector = tuple(
                    [torch.normal(mean=0, std=1, size=p.size(), device=self.device) for p in original_params])

                # Define the loss function as a callable function for jvp
                def compute_loss(*params):
                    for param, p in zip(model_params, params):
                        param.data = p.data
                    input_ids = inputs["input_ids"]
                    attention_mask = inputs.get("attention_mask", None)
                    if attention_mask is not None:
                        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                    else:
                        outputs = self.model(input_ids=input_ids)
                    loss = self.criterion(outputs.logits.view(-1, outputs.logits.shape[-1]), labels.view(-1))
                    return loss

                # Compute the forward gradient (Jacobian-vector product)
                loss_value, grad_vector = jvp(compute_loss, original_params, tangents_vector)

                # Update the model parameters manually
                for param, grad in zip(model_params, grad_vector):
                    param.data -= 0.001 * grad  # Here, 0.001 is the learning rate

                # print statistics
                if i % 100 == 99:
                    print(f'Batch: {i + 1}, loss: {loss_value.item()}')
                    f.write(f'Batch: {i + 1}, loss: {loss_value.item()}\n')

                # Cleanup memory
                del inputs, labels, loss_value, grad_vector
                torch.cuda.empty_cache()
                gc.collect()
        print("Training finished")


# Load the tokenizer and model
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForMaskedLM.from_pretrained('bert-base-uncased')

# Load the dataset
dataset = load_dataset('wikitext', 'wikitext-2-raw-v1')


# Preprocessing function to tokenize the texts and prepare labels
def preprocess(examples):
    tokenized_inputs = tokenizer(examples['text'], truncation=True, padding='max_length', max_length=128,
                                 return_tensors='pt')
    labels = tokenized_inputs.input_ids.clone()
    labels[:-1] = tokenized_inputs.input_ids[1:]
    return {"input_ids": tokenized_inputs.input_ids.squeeze(), "labels": labels.squeeze()}


# Preprocess the dataset
encoded_dataset = dataset.map(preprocess)
encoded_dataset.set_format(type='torch', columns=['input_ids', 'labels'])

# Create a DataLoader
dataloader = DataLoader(encoded_dataset['train'], batch_size=16)

# Define the criterion
criterion = torch.nn.CrossEntropyLoss()

# Define device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Create the training object and train for one epoch
trainer = ForwardADTraining(model, dataloader, criterion, device)
trainer.train_one_epoch()
`
```

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @Chillee @samdow @kshitij12345 @janeyx99