Gradient Checkpointing for LongformerForSequenceClassification produces memory error

I am trying to include gradient checking for LongformerForSequenceClassification model.

I first tested this on 512 tokens and it seems to produce a cuda out of memory error (GPU is K80). If I don't include gradient_checkpoint=True, then there's no error, with the rest of the code staying the same. It's important for me to get this to work because I want to be able to work on at least 2000+ token lengths eventually. Without gradient checkpointing, I can't achieve this.

I have a multi-gpu setup but this breaks as well for single gpu. I'm running this on AWS sagemaker notebook. My batch size is 16.

```

def train_model(train_dataset, epochs=4) :
train_dataloader = create_data_loader(train_dataset)

model = LongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096',
gradient_checkpointing=True, # New to v3
num_labels=4)
model = torch.nn.DataParallel(model) # Parallelize for multi GPU
model.to(device) # Move to GPU

optimizer = AdamW(model.parameters(),
lr = 2e-5, # args.learning_rate - default is 5e-5, our notebook had 2e-5
eps = 1e-8 # args.adam_epsilon - default is 1e-8.
)

# Total number of training steps is [number of batches] x [number of epochs].
total_steps = len(train_dataloader) * epochs

# Create the learning rate scheduler.
scheduler = get_linear_schedule_with_warmup(optimizer,
num_warmup_steps = 0, # Default value in run_glue.py
num_training_steps = total_steps)

# Set the seed value all over the place to make this reproducible.
random.seed(seed_val)
np.random.seed(seed_val)
torch.manual_seed(seed_val)
torch.cuda.manual_seed_all(seed_val) #Re-added for GPU

# For each epoch...
for epoch_i in range(0, epochs):

# Measure time
t0 = time.time()

# Perform one full pass over the training set.
print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))

# Reset the total loss for this epoch.
total_train_loss = 0

# Put the model into training mode.
model.train()

# For each batch of training data...
for step, batch in enumerate(train_dataloader):
t1 = time.time()
print('.', end ="")
b_input_ids = batch[0].to(device) #Re-added for GPU
b_input_mask = batch[1].to(device)
b_labels = batch[2].to(device)

# Always clear any previously calculated gradients before performing a backward pass.
model.zero_grad()

# Perform a forward pass (evaluate the model on this training batch).
loss, logits = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)

# https://discuss.pytorch.org/t/how-to-fix-gathering-dim-0-warning-in-multi-gpu-dataparallel-setting/41733/2
loss = loss.mean()
total_train_loss += loss.item()

# Perform a backward pass to calculate the gradients.
loss.backward()

# Clip the norm of the gradients to 1.0. This is to help prevent the "exploding gradients" problem.
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

# Update parameters and take a step using the computed gradient.
optimizer.step()

# Update the learning rate.
scheduler.step()


# Calculate the average loss over all of the batches.
avg_train_loss = total_train_loss / len(train_dataloader)
print(" Average training loss: {0:.2f}".format(avg_train_loss))

training_time = format_time(time.time() - t0)
print(" Training epoch took: {:}".format(training_time))



print("Training complete!")
return(model)
```