Fine-Tuning GPT2 on Colab GPU... For Free! | by Joey S | Towards Data Science

Models these days are very big, and most of us don’t have the resources to train them from scratch. Luckily, HuggingFace has generously provided pretrained models in PyTorch, and Google Colab allows usage of their GPU (for a fixed time). Otherwise, even fine-tuning a dataset on my local machine without a NVIDIA GPU would take a significant amount of time. While the tutorial here is for GPT2, this can be done for any of the pretrained models given by HuggingFace, and for any size too.
Go to Google Colab and create a new notebook. It should look something like this.
Set to use GPU by clicking Runtime > Change runtime type
Then click Save.
We would run pip3 install transformers normally in Bash, but because this is in Colab, we have to run it with !
!pip3 install transformers
You can read more about WikiText data here. Overall, there’s WikiText-2 and WikiText-103. We’re going to use WikiText-2 because it’s smaller, and we have limits in terms of how long we can run on GPU, and how much data we can load into memory in Colab. To download and run, in a cell, run
%%bashwget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-raw-v1.zipunzip wikitext-2-raw-v1.zip
HuggingFace actually provides a script to help fine-tune models here. We can just download the script by running
!wget https://raw.githubusercontent.com/huggingface/transformers/master/examples/language-modeling/run_language_modeling.py
Now we are ready to fine-tune.
There are many parameters to the script, and you can understand them by reading the manual. I’m just going to go over the important ones for basic training.
output_dir is where the model will be output
model_type is what model you want to use. In our case, it’s gpt2. If you have more memory and time, you can select larger gpt2 sizes which are listed in HuggingFace pretrained models list.
model_name_or_path is the path to the model. If you want to train from scratch, you can leave this blank. In our case, it’s also gpt2
do_train tells it to train
train_data_file points to the training file
do_eval tells it to evaluate afterwards. Not always required, but good to have
eval_data_file points to the evaluation file
Some extra ones you MAY care about, but you can also skip this.
save_steps is when to save checkpoints. If you have limited memory, you can set this to -1 so it’ll skip saving until the end
per_gpu_train_batch_size is batch size for GPU. You can increase this if your GPU has enough memory. To be safe, you can start with 1 and ramp it up if you still have memory
num_train_epochs is the number of epochs to train. Since we’re fine-tuning, I’m going to set this to 2
overwrite_output_dir is used if the output_dir already has something in it, and you’re fine to overwrite the existing model
All in all, to train, run this in a cell
%%bashexport TRAIN_FILE=wikitext-2-raw/wiki.train.rawexport TEST_FILE=wikitext-2-raw/wiki.test.rawexport MODEL_NAME=gpt2export OUTPUT_DIR=outputpython run_language_modeling.py --output_dir=$OUTPUT_DIR \--model_type=$MODEL_NAME \--model_name_or_path=$MODEL_NAME \--do_train \--train_data_file=$TRAIN_FILE \--do_eval \--eval_data_file=$TEST_FILE \--per_gpu_train_batch_size=1 \--save_steps=-1 \--num_train_epochs=2
Note that if you want to fine-tune the model you just trained, you can change MODEL_NAME=gpt2 to MODEL_NAME=output/ so it’ll load the model we just trained
When you run this, if it takes some time without any output, you can hover over the RAM/Disk on the top right corner to see what’s happening.
The downside to Colab GPU is that it’s shared between Colab users. That means execution may not happen right away, as it’s being used by another user. When that happens, it’ll say something like
Waiting for 'Python 3 Google Compute Engine backend (GPU(' to finish its current execution.
There’s really nothing to do but wait it out.
After you’ve finished running the model, you can check that it exists in your output directory.
To use it, you can run something like
where = Toronto Raptors = is the equivalent of describing Toronto Raptors as the title of the article.
The result I got (and yours will differ) is
 = Toronto Raptors =  Toronto's first @-@ round draft pick in 2006 was selected by the Toronto Raptors with the seventh overall pick. He played in all 82 games, averaging 18 points per game and 6 assists. The Raptors won their third straight NBA championship in 2007, and won the 2009 NBA All @-@ Star Game. He played in a record 16 games for Toronto, averaging 19 points on 5 @.@ 6 rebounds and 6 assists in a season that saw Toronto win the Eastern Conference finals. He also played in the 2008 All @-@ Star Game and was named to the All @-@ Star Game MVP for the first time. He also was named to the All @-@ Star Game's all @-@ time career scoring list, and was the first player to accomplish the feat. He finished the season with an assist and an assist in eight games, and was the first player in NBA history to score in double figures. He was named to the All @-@ Star Game's All @-@ time scoring list in 2011, and was the first player to do this in consecutive seasons. = = Draft = = Toronto selected Jordan Matthews with the seventh overall pick in
From my example, I only generated the first 250 words, which is why it was cut off so abruptly. You can expand that if you want. Notice that this description of the Toronto Raptors is completely fake, as Jordan Matthews never played for the Raptors. The text coherency can be better as well, which can be adjusted by tuning with more epochs, or simply using a larger model. However, that requires more memory, so be careful with it.
In order for us to preserve this model, we should compress it and save it somewhere. This can be done easily with
! tar -czf gpt2-tuned.tar.gz output/
which creates a file called gpt2-tuned.tar.gz
To save it to your Google Drive from Colab, first you must have a Google account/Gmail account. In your Colab cell, you can run
from google.colab import drivedrive.mount('/content/drive')
without needing to install anything additional, since google.colab library comes with using Google Colab. When you run the code described above, you should see something like this
You have to click on the link, sign in and allow your Google Drive to access the your Colab. Then at the end, you’ll see something like
Copy that and paste it back into your Colab notebook.
Now you can copy your output model to your Google Drive by running
!cp gpt2-tuned.tar.gz /content/drive/My\ Drive/
And Voila! You’ve successfully tuned a pretrained model on a GPU, and saved it in your Google Drive. And you did it completely for free.
If you have any questions or improvements based on what we’ve worked on, please let me know in the comments.
You can run this Colab Notebook to reproduce everything shown above