[Python] RuntimeError: Input tensor at index 3 has invalid shape [2, 2, 16, 128, 64] but...

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("gpt2-medium")
gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")


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 & test_path are unstructured text data file of size 1.45 Million and 200K lines of data

Training arguments:

training_args = TrainingArguments(
output_dir="./gpt2-finetuned-models", #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 = "eval_loss",
load_best_model_at_end = True,
evaluation_strategy="epoch",
)


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]

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

Continue reading...