{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Training kernel for XLM-RoBERTa using PyTorch on the TPU\n\n## If you found this helpful, please give it an upvote!\n\nUsing the amazing [PyTorch XLA library](https://github.com/pytorch/xla)\n\nValidation will be added soon.\n\n- V1: Proof-of-concept. It runs!\n- V6: Add validation code\n\nFurther work:\n1. Tune LR further\n2. xhlulu dataset (prelim. code works)\n3. Further memory optimization tricks"},{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nIn this kernel, I will demonstrate how to properly train a larger model with [PyTorch XLA](https://pytorch.org/xla). PyTorch XLA allows one to train PyTorch models on [Google's tensor processing units (TPUs)](https://cloud.google.com/tpu). Kaggle provides 30 hours of free TPU compute.\n\nAbhishek has already made a great PyTorch XLA kernel [here](https://www.kaggle.com/abhishek/bert-multi-lingual-tpu-training-8-cores-w-valid) for training a multi-lingual BERT model on a modified dataset (training on a sample of 20,000 rows from the data, plus validation set). However, xhlulu shared a [kernel](https://www.kaggle.com/xhlulu/jigsaw-tpu-xlm-roberta) training an XLM-RoBERTa model using TPU and TensorFlow. xhlulu's uses a modified dataset, using the full train dataset downsampled.\n\nThis kernel is an attempt to train the XLM-RoBERTa model with PyTorch XLA. I use Abhishek's dataset to start out. I will attempt to use xhlulu's dataset in a follow-up kernel.\n\n\nThere are a couple challenges that make PyTorch XLA a little harder to use:\n1. The API is much more lower-level compared to TensorFlow. PyTorch XLA requires you to use XLA-specific dataloaders and XLA-specific optimizer stepping. Additionally, it requires the definition of a train/evaluation loop function that needs to be spawned using PyTorch XLA's multiprocessing functionality. PyTorch Lightning's [TPU support](https://pytorch-lightning.readthedocs.io/en/latest/tpu.html) may help alleviate some of these issues, but code may still need to be optimized further, at least when using Kaggle TPUs, due to the low amount of RAM available. TensorFlow simple requires the definition of a TPU distribution strategy, and you are set to go.\n2. PyTorch XLA works differently compared to TensorFlow, leading to higher host VM memory usage with PyTorch XLA compared to TensorFlow. Specifically, PyTorch XLA builds the XLA graphs, initializes the weights, runs input pipelines etc and then feeds them to the TPUs. On the other hand, TensorFlow hands the XLA graphs to the TPU and the TPU does most of the heavy lifting. This difference is amplified by the fact that the appropriate data and model needs to be replicated 8 times to utilize all eight cores of the TPU.\n\nIn this kernel, we will see some additional optimizations that will allow us to train the XLM-R model. There are _even more_ optimizations that could be done, which will be described in a future kernel."},{"metadata":{},"cell_type":"markdown","source":"# Imports\n\nThe code cell below will install PyTorch XLA."},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here are all of our imports. You will note that there are already an optimization applied in order for PyTorch XLA to train.\n\n`XLA_USE_BF16` is an environment variable that tells PyTorch XLA to automatically use [bfloat16](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus)."},{"metadata":{"_uuid":"993abe0b-2561-4abc-a6a2-b83d8dd4c846","_cell_guid":"16a3df23-8416-46b5-8661-c52345005b6d","trusted":true},"cell_type":"code","source":"import os\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'\n\nimport torch\nimport pandas as pd\nfrom scipy import stats\nimport numpy as np\n\nfrom tqdm import tqdm\nfrom collections import OrderedDict, namedtuple\nimport torch.nn as nn\nfrom torch.optim import lr_scheduler\nimport joblib\n\nimport logging\nimport transformers\nfrom transformers import AdamW, get_linear_schedule_with_warmup, get_constant_schedule, XLMRobertaTokenizer, XLMRobertaModel, XLMRobertaConfig\nimport sys\nfrom sklearn import metrics, model_selection","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here we import all the PyTorch XLA-specific modules."},{"metadata":{"_uuid":"993abe0b-2561-4abc-a6a2-b83d8dd4c846","_cell_guid":"16a3df23-8416-46b5-8661-c52345005b6d","trusted":true},"cell_type":"code","source":"import warnings\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"993abe0b-2561-4abc-a6a2-b83d8dd4c846","_cell_guid":"16a3df23-8416-46b5-8661-c52345005b6d","trusted":true},"cell_type":"code","source":"class AverageMeter:\n    \"\"\"\n    Computes and stores the average and current value\n    \"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset and Model"},{"metadata":{},"cell_type":"markdown","source":"Here is a simple class to create datasets from numpy arrays. I think this dataset class could likely be further optimized for decreased memory usage."},{"metadata":{"trusted":true},"cell_type":"code","source":"class ArrayDataset(torch.utils.data.Dataset):\n    def __init__(self,*arrays):\n        assert all(arrays[0].shape[0] == array.shape[0] for array in arrays)\n        self.arrays = arrays\n    \n    def __getitem__(self, index):\n        return tuple(torch.from_numpy(np.array(array[index])) for array in self.arrays)\n    \n    def __len__(self):\n        return self.arrays[0].shape[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here is the XLM-RoBERTa model definition inspired by the model used in xhlulu's kernel."},{"metadata":{"trusted":true},"cell_type":"code","source":"class CustomRoberta(nn.Module):\n    def __init__(self):\n        super(CustomRoberta, self).__init__()\n        self.num_labels = 1\n        self.roberta = transformers.XLMRobertaModel.from_pretrained(\"xlm-roberta-large\", output_hidden_states=False, num_labels=1)\n        self.dropout = nn.Dropout(p=0.2)\n        self.classifier = nn.Linear(1024, self.num_labels)\n\n    def forward(self,\n                input_ids=None,\n                attention_mask=None,\n                position_ids=None,\n                head_mask=None,\n                inputs_embeds=None):\n\n        _, o2 = self.roberta(input_ids,\n                               attention_mask=attention_mask,\n                               position_ids=position_ids,\n                               head_mask=head_mask,\n                               inputs_embeds=inputs_embeds)\n\n        logits = self.classifier(o2)       \n        outputs = logits\n        return outputs\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's define our model. If you check the PyTorch XLA documentation, you will note that the recommended setup is to define the model in the function that is run on each of the 8 cores. But doing so will lead to high VM memory usage. Therefore, this is the setup used for low-memory VMs."},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"mx = CustomRoberta();\nmx","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The trick also used in xhlulu's kernel is to pre-tokenize the dataset. This is done [over here](https://www.kaggle.com/tanlikesmath/xlm-r-large-tokenize-dataset). It uses the same `regular_encode` function defined in xhlulu's kernel. We now load it over here. \n\nAn additional trick used that may help memory usage is to load the dataset as a memory-mapped dataset (`mmap_mode='r'`). This way, the whole dataset isn't in the RAM at the same time."},{"metadata":{"trusted":true},"cell_type":"code","source":"tokenized_path = '../input/xlm-r-large-tokenize-dataset/'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x_train = np.load(tokenized_path+'x_train.npy',mmap_mode='r')\ntrain_toxic = np.load(tokenized_path+'df_train_toxic.npy',mmap_mode='r')\n\nx_valid = np.load(tokenized_path+'x_valid.npy',mmap_mode='r')\nvalid_toxic = np.load(tokenized_path+'df_valid_toxic.npy',mmap_mode='r')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x_train.shape, x_valid.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's create our dataset!"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = ArrayDataset(x_train, train_toxic)\nvalid_dataset = ArrayDataset(x_valid, valid_toxic)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Make sure to delete any unused variables:"},{"metadata":{"trusted":true},"cell_type":"code","source":"del x_train, x_valid\nimport gc;gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch_xla.version as xv\nprint('PYTORCH:', xv.__torch_gitrev__)\nprint('XLA:', xv.__xla_gitrev__)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!free -h","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def loss_fn(outputs, targets):\n    return nn.BCEWithLogitsLoss()(outputs, targets.view(-1, 1))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def reduce_fn(vals):\n    return sum(vals) / len(vals)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We now define our training loop and evaluation loop functions.\n\nTo get the loss of a batch, since the data is spread across the 8 cores, we have to _reduce_ the loss.\n\nPyTorch XLA requires that the optimizer be stepped using their own function `xm.optimizer_step(optimizer)`."},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_loop_fn(data_loader, model, optimizer, device, scheduler=None):\n    model.train()\n    for bi, d in enumerate(data_loader):\n\n        ids = d[0]\n        targets = d[1]\n\n        ids = ids.to(device, dtype=torch.long)\n        targets = targets.to(device, dtype=torch.float)\n\n        optimizer.zero_grad()\n        outputs = model(\n            input_ids=ids,\n        )\n        loss = loss_fn(outputs, targets)\n        if bi % 50 == 0:\n            loss_reduced = xm.mesh_reduce('loss_reduce',loss,reduce_fn)\n            xm.master_print(f'bi={bi}, loss={loss_reduced}')\n        loss.backward()\n        xm.optimizer_step(optimizer)\n        if scheduler is not None:\n            scheduler.step()\n            \n\n    model.eval()\n    \ndef eval_loop_fn(data_loader, model, device):\n    fin_targets = []\n    fin_outputs = []\n    for bi, d in enumerate(data_loader):\n        ids = d[0]\n        targets = d[1]\n\n        ids = ids.to(device, dtype=torch.long)\n        targets = targets.to(device, dtype=torch.float)\n\n        outputs = model(\n            input_ids=ids,\n        )\n\n        targets_np = targets.cpu().detach().numpy().tolist()\n        outputs_np = outputs.cpu().detach().numpy().tolist()\n        fin_targets.extend(targets_np)\n        fin_outputs.extend(outputs_np)    \n        del targets_np, outputs_np\n        gc.collect()\n    return fin_outputs, fin_targets","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We finally define our function that will be spawned by PyTorch XLA multiprocessing. This function will be run on each of the 8 cores. There are several things to note:\n1. We need to use a `DistributedSampler` that will appropriately distribute the dataset across the 8 cores.\n2. We are using `num_workers=0` as that decreases memory usage (only master process loading data).\n3. We put the model onto the TPU\n4. We use `ParallelLoader` which is a PyTorch XLA-specific DataLoader for loading data onto the TPU."},{"metadata":{"trusted":true},"cell_type":"code","source":"def _run():\n    MAX_LEN = 192\n    TRAIN_BATCH_SIZE = 16\n    EPOCHS = 2\n\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n          train_dataset,\n          num_replicas=xm.xrt_world_size(),\n          rank=xm.get_ordinal(),\n          shuffle=True)\n\n    train_data_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=TRAIN_BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=0,\n    )\n    \n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n          valid_dataset,\n          num_replicas=xm.xrt_world_size(),\n          rank=xm.get_ordinal(),\n          shuffle=False)\n\n    valid_data_loader = torch.utils.data.DataLoader(\n        valid_dataset,\n        batch_size=4,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=0\n    )\n\n    device = xm.xla_device()\n    model = mx.to(device)\n    xm.master_print('done loading model')\n\n    param_optimizer = list(model.named_parameters())\n    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n    optimizer_grouped_parameters = [\n        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]\n\n    lr = 0.5e-5 * xm.xrt_world_size()\n    num_train_steps = int(len(train_dataset) / TRAIN_BATCH_SIZE / xm.xrt_world_size() * EPOCHS)\n    \n    optimizer = AdamW(optimizer_grouped_parameters, lr=lr)\n    scheduler = get_linear_schedule_with_warmup(\n        optimizer,\n        num_warmup_steps=0,\n        num_training_steps=num_train_steps\n    )\n    xm.master_print(f'num_train_steps = {num_train_steps}, world_size={xm.xrt_world_size()}')\n\n\n    for epoch in range(EPOCHS):\n        gc.collect()\n        para_loader = pl.ParallelLoader(train_data_loader, [device])\n        xm.master_print('parallel loader created... training now')\n        gc.collect()\n        train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=scheduler)\n        del para_loader\n        para_loader = pl.ParallelLoader(valid_data_loader, [device])\n        gc.collect()\n        o, t = eval_loop_fn(para_loader.per_device_loader(device), model, device)\n        del para_loader\n        gc.collect()\n        auc = metrics.roc_auc_score(np.array(t) >= 0.5, o)\n        auc_reduced = xm.mesh_reduce('auc_reduce',auc,reduce_fn)\n        xm.master_print(f'AUC = {auc_reduced}')\n        gc.collect()\n    xm.save(model.state_dict(), \"xlm_roberta_model.bin\")\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Start training!\n\nLet's spawn the `_mp_fn` and start training!"},{"metadata":{"trusted":true},"cell_type":"code","source":"import time\n\n# Start training processes\ndef _mp_fn(rank, flags):\n    a = _run()\n\nFLAGS={}\nstart_time = time.time()\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Time taken: ',time.time()-start_time)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Acknowledgments:\n- Based on data from [Abhishek's code](https://www.kaggle.com/abhishek/bert-multi-lingual-tpu-training-8-cores-w-valid)\n- Model based on [xhlulu's code](https://www.kaggle.com/xhlulu/jigsaw-tpu-xlm-roberta)\n- Original attempt from [Aditya's code](https://www.kaggle.com/adityaecdrid/simple-xlmr-tpu-pytorch)\n- Discussion with Davide Libenzi and Daniel Sohn (PyTorch XLA team) - [code](https://www.kaggle.com/davidelibenzi/simple-xlmr-tpu-pytorch)\n- Fruitful discussions with Abhishek and Aditya\n\n\n# Fin\n\nIf you have any questions or suggestions, please drop a comment! :)"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}