RuntimeError: Input tensor at index 3 has invalid shape [2, 2, 16, 128, 64] but expected [2, 4, 16, 128, 64]
Runtime error while finetuning a pretrained GPT2-medium model using Huggingface library in SageMaker - ml.p3.8xlarge instance.
The finetuning_gpt2_script.py contains the below,
Libraries:
from transformers import Trainer, TrainingArguments
from transformers import EarlyStoppingCallback
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from transformers import TextDataset,DataCollatorForLanguageModeling

Pretrained Models:
gpt2_model = GPT2LMHeadModel.from_pretrained(&quot;gpt2-medium&quot;)
gpt2_tokenizer = GPT2Tokenizer.from_pretrained(&quot;gpt2-medium&quot;)

Train and Test Data Construction:
train_dataset = TextDataset(
          tokenizer=gpt2_tokenizer,
          file_path=train_path,
          block_size=128)
    
test_dataset = TextDataset(
          tokenizer=gpt2_tokenizer,
          file_path=test_path,
          block_size=128)
    
data_collator = DataCollatorForLanguageModeling(
        tokenizer=gpt2_tokenizer, mlm=False,
    )

train_path &amp; test_path are unstructured text data file of size 1.45 Million and 200K lines of data
Training arguments:
training_args = TrainingArguments(
        output_dir=&quot;./gpt2-finetuned-models&quot;, #The output directory
        overwrite_output_dir=True, #overwrite the content of the output directory
        num_train_epochs=1, # number of training epochs
        per_device_train_batch_size=8, # batch size for training #32
        per_device_eval_batch_size=8,  # batch size for evaluation #64
        save_steps=100, # after # steps model is saved
        warmup_steps=500,# number of warmup steps for learning rate scheduler
        prediction_loss_only=True,
        metric_for_best_model = &quot;eval_loss&quot;,
        load_best_model_at_end = True,
        evaluation_strategy=&quot;epoch&quot;,
    )

training_args are the training arguments constructed to train the model.
Trainer:
trainer = Trainer(
        model=gpt2_model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        callbacks = [early_stop_callback],
    )
early_stop_callback = EarlyStoppingCallback(early_stopping_patience  = 3)

Training:
trainer.train()
trainer.save_model(model_path)

Here, the training is done for only 1 epoch in 4 GPUS using ml.p3.8xlarge instance.
The training is done by torch-distribution like below,
python -m torch.distributed.launch finetuning_gpt2_script.py

While training at the end of the epoch, observed the below error,
RuntimeError: Input tensor at index 3 has invalid shape [2, 2, 16, 128, 64] but expected [2, 4, 16, 128, 64]

Is the RuntimeError because of the way the train_dataset and test_datasetconstructed using TextData ?
Am I doing wrong in the torch-distribution ?


It could be related to a mismatch in the batch size (expecting a batch size of 4 but receiving a batch size of 2) as suggested here ? Solution provided is to set the parameter drop_last in your DataLoader like this:
tain_text = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True)

