Environment
Python 3.9, transformer package in huggingface, and datasets package in huggingface.
And also install: https://github.com/chengxuz/pt_framework
Install the current repo using pip install . or pip install -e ..
Where to put data
First, define the environment variable BABYLM_ROOT_DIR to be where your models and data will live.
The downloaded data should be put at ${BABYLM_ROOT_DIR}/datasets/ so that this folder contains the following four subfolders: babylm_100M, babylm_10M, babylm_dev, and babylm_test. Note that the T5 training script expects .txt file inputs, so we create a single dev file by running this command in the ${BABYLM_ROOT_DIR}/datasets/babylm_dev/ folder: cat *.dev > babylm_dev.txt.
The trained models will be put at ${BABYLM_ROOT_DIR}/models/ and the records will be put at ${BABYLM_ROOT_DIR}/model_recs/.
Training Command
OPT-125M
Run the following command under the scripts folder.
python -m torch.distributed.launch --nproc_per_node=1 --master_port=29123 general_train.py --setting "BabyLM/exp_strict.py:opt125m_s1"
This command will load a training setting specified by function opt125m_s1 at src/babylm_baseline_train/configs/BabyLM/exp_strict.py.
RoBERTa-Base
Run the following command under the scripts folder.
python -m torch.distributed.launch --nproc_per_node=1 --master_port=29123 general_train.py --setting "BabyLM/exp_strict_mask.py:roberta_s1"
T5-Base
Run the following command under the scripts folder.
./train_t5_babylm.sh
Note that this training script uses a different backend than the OPT and RoBERTa models. This script is a slightly modified version of the flax T5 pre-training script from huggingface; the original lives here.
Where important parameters are defined
Learning rate schedule is defined at function get_learning_rate_params in script basic_param_setter.py under src/babylm_baseline_train folder.
Optimizer is in the scripts/general_train.py script inside the get_key_params funciton.
How to load the pretrained models
See the functions in src/babylm_baseline_train/models/ckpt_loader.py.
Questions?
Feel free to open issues here. Or just contact us through Slack/emails.