I converted a german RoBERTa model into a Longformer with a capacity of 4096 token using this script 1. I use it for LEDModel(). But when I tokenize the data with max_length=4096, I get the following error when printing the loss or executing loss.backward():
RuntimeError: CUDA error: device-side assert triggered
I understand that this has something to do with invalid values outside of the expected index range of possible tokens (read that here). When I run the same code with max_length=512 it somehow works. Could it be that the convertion to a Longformer didn’t pick up the possible length of input sequences? But if so, why should that be a problem regarding the size of the vocabulary? There is a snippet of my code below.
Edit:
I think I know what’s wrong: LEDModel() wants an encoder decoder model and not just a transformer/longformer. So this class doesn’t chain two transformers together to create an encoder decoder.
So the question is: How can I use my Longformer for the Seq2Seq task so that I can also utilize gradient checkpointing?
MAX_LEN = 4096
batch_size = 1
epochs = 1
max_grad_norm = 0.06
weight_decay = 0.3
FULL_FINETUNING = True
learning_rate = 1e-5
adam_eps = 1e-10

torch.cuda.empty_cache()
model_path = sys.argv[2]
tokenizer = RobertaTokenizer.from_pretrained(model_path)
tokenizer.pad_token = 0

def train(train_dataloader, valid_dataloader, test_dataloader, voc_id_to_tok):
    model = LEDModel.from_pretrained(model_path)

    model.config.pad_token_id = tokenizer.pad_token_id
    model.config.vocab_size = len(voc_id_to_tok)
    print("vocab size", model.config.vocab_size)
    model.config.use_cache = False
    model.gradient_checkpointing_enable()
    model.resize_token_embeddings(len(tokenizer))

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    model.to(device)

    if FULL_FINETUNING:
        param_optimizer = list(model.named_parameters())
        no_decay = ['bias', 'gamma', 'beta']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
             'weight_decay_rate': weight_decay},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
             'weight_decay_rate': 0.0}
        ]
    else:
        param_optimizer = list(model.classifier.named_parameters())
        optimizer_grouped_parameters = [{"params": [p for n, p in param_optimizer]}]

    optimizer = AdamW(
        optimizer_grouped_parameters,
        lr=learning_rate,
        eps=adam_eps
    )

    scaler = GradScaler()

    # Total number of training steps is number of batches * 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,
        num_training_steps=total_steps
    )

    for epoch in trange(epochs, desc="Epoch"):
        # ========================================
        #               Training
        # ========================================
        # Perform one full pass over the training set.

        model.train()
        total_loss = 0

        # Training loop
        for step, batch in enumerate(train_dataloader):
            batch = tuple(t.to(device) for t in batch)

            b_input_ids, b_output_ids = batch  # b_input_mask, b_output_ids = batch

            model.zero_grad()
            with autocast():
                outputs = model(input_ids=b_input_ids, decoder_input_ids=b_output_ids)

                loss = outputs[0].mean()

            scaler.scale(loss).backward()
            total_loss += loss.item()
            clip_grad_norm_(parameters=model.parameters(), max_norm=max_grad_norm)

            scaler.step(optimizer)
            scaler.update()
            scheduler.step()

        avg_train_loss = total_loss / len(train_dataloader)
        print("Average train loss: {}".format(avg_train_loss))

        # Store the loss value for plotting the learning curve.
        loss_values.append(avg_train_loss)