{"cells":[{"metadata":{},"cell_type":"markdown","source":"# About this Notebook\n* In 2018 ,Google open sourced the device which was the backbone of their state of the art results ,the TPU's (more on them in a bit) . They made TPU's available through their cloud services for anyone to use at $2/hour and after some time TPU's were made available for free on Google colab .\n\n* As we all know TPU's are very fast and give state of the art results for neural networks , however initially it could only be used with TensorFlow and Keras , leaving the pytorch fans really frustated as they didn't want to shift to TF . This led to a chain of events for the development of way allowing TPU's to be used with Pytorch \n\n* Hence Pytorch-XLA module was developed which lets pytorch to run its graph on xla_devices like TPU's . \n\nIn the **Jigsaw competition** TPU's have been used in various different ways , on single cores, multi cores, using different chechpoints etc,etc . However when I tried to understand the Publically shared kernels to explore how to use TPU cores with pytorch, I found it was really difficult to comphrehend. Also for Jigsaw the best performing model is XLM-Roberta which is a fairly large model and throws an error if someone is not careful with memory management.\n\nThe complexity of TPU's and usage of XLM-Roberta with TPU's left me frustated and angry . I decided to take this TPU thing slowly and started with small models building on that with a lot of experimentations upto XLM-Roberta .\n\n**In this Notebook I share my experimentations with Pytorch-XLA and TPU's . I will start from basic TPU usage and build on that to show how to use TPUs on multiple cores and also with multithreading. I will also share some tips and tricks which would be useful when using TPU's with Pytorch XLA. If you want to learn to use TPU's the easy way, this might be a good place to start** \n\nAfter learning this You will be able to Decode and easily understand ALEX's and Abhishek's kernels","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Prerequisites:\n\n* This notebook assumes that you are familiar with Pytorch and have used it before to build models. If you have not , here are some useful links to learn pytorch from zero:
\n1) https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html
\n2) https://pytorch.org/tutorials/beginner/nn_tutorial.html
\n3) https://pytorch.org/tutorials/beginner/pytorch_with_examples.html
\n\n* Using Pytorch is fairly easy , it's just like using python only in a more advanced kind of way. Here is the link to the repository which contains Pytorch implementations on different architectures : https://github.com/tanulsingh/Pytorch-for-Everyone","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np \nimport pandas as pd\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n for filename in filenames:\n print(os.path.join(dirname, filename))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# What are TPU's? How they work? How are they different from a GPU?\n\nYou might be thinking why knowing how tpus work is important , well it's not a must but to exploit something fully we must know how it works right?\nTPUs are hardware accelerators specialized in deep learning tasks. For explanation of what TPU's are and how they work please go through the following videos :\n* [video1](https://www.youtube.com/watch?v=MXxN4fv01c8)\n* [video2](https://www.youtube.com/watch?v=kBjYK3K3P6M)
\nIts important to understand the underlying concepts of Pytorch XLA's . If you want to dig even deeper [here](https://codelabs.developers.google.com/codelabs/keras-flowers-data/#2) is a article by google explaining everything about TPU's\n\n\n# Key Takeaways\nFollowing are the key takeaways from the above videos and articles :-\n\n* Each TPU v3 board has 8 TPU cores and 64 GB's of memory\n* TPU's consist of two units, Matrix Multiply Unit (MXU) which runs matrix multiplications and a Vector Processing Unit (VPU) for all other tasks such as activations, softmax, etc.\n* TPU's v2/v3 use a new type of dtype called bfloat16 which combines the range of a 32-bit floating point number with just the storage space of only a 16-bit floating point number and this allows to do fit more matrices in the memory and thus more matrix multiplications. This increased speed comes at the cost of precision as bfloat16 is able to represent fewer decimal places as compared to 16-bit floating point integer but its ohk because neural networks can work at a reduced precision while maintaining their high accuracy\n* The ideal batch size for TPUs is 128 data items per TPU core but the hardware can already show good utilization from 8 data items per TPU core\n\n**Now we move onto the final question does TPU's directly run the Python code? Or is there something else working under the hood without credits**\n\n\n","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Under the Hood\n* We know that any deep learning framework first defines a computation graph which is then executed by any processing chip to train a neural network. Similarly, The TPU does not directly run Python code, it runs the computation graph defined by your program.However the computation graph is first converted into TPU machine code. Under the hood, a compiler called XLA (accelerated Linear Algebra compiler) transforms the graph of computation nodes into TPU machine code. This compiler also performs many advanced optimizations on your code and your memory layout. \n* In tensorflow the conversion from computation to TPU machine code automatically takes place as work is sent to the TPU, whereas there was no such support for Pytorch and thus XLA module was created to include XLA in our build chain explicitly.\n\n","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"Let's start with importing Pytorch-XLA and necessary Modules","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-output":true,"_kg_hide-input":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":"PS : Above code only works when TPU is on ๐
","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import os\nimport time\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom sklearn import metrics\nfrom sklearn.model_selection import train_test_split\nimport transformers\nfrom transformers import AdamW\n\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now if you have used Pytorch with GPU's before , you know that running your code and models on GPU's is this simple :\n* device = \"cuda\"\n* tensors.to(device)\n* model.to(device)
\nThus While building Pytorch-XLA developers had the same thing in mind and they wanted to create something similar, so that the end users have the same feel and structure while using TPU's as they had while using GPU's
\n**Pytorch-XLA treats each TPU core as an individual XLA device and thus using a TPU core is as easy as**:\n* device = xm.xla_device()\n* tensors.to(device)\n* model.to(device)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"t0 = torch.randn(2, 2, device=xm.xla_device()) #creating a tensor and sending it to the TPU\nt1 = torch.randn(2, 2, device=xm.xla_device()) #creating a tensor and sending it to the TPU\nprint(t0 + t1) # As both tensors are now on the same device i.e same TPU core we can perform any calculation on them like addition\nprint(t0.mm(t1)) # Matrix Multiplication","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Sending a model On TPU","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"l_in = torch.randn(10, device=xm.xla_device())\nlinear = torch.nn.Linear(10, 20).to(xm.xla_device())\nl_out = linear(l_in) #NOTE THAT THE TENSOR AND MODEL BOTH BE SENT TO THE DEVICE AS WE DID WITH GPUS , THEN ONLY we CAN PERFORM ANY OPERATION\nprint(l_out)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training on a Single TPU Core\nNow that we know the basics of Pytorch-XLA , let's see how to train models on TPU Single cores and what are the changes in code.\nI will be using a subset of Jigsaw Competition , and using BERT-BASE for simplicity , In the end of this notebook I will provide information about the experimentations I have done with XLM-Roberta and TPU's and give insights to efficiently train XLM-Roberta on TPU's
\n**First I will start with Functions that will remain the same irrespective of where and how we want to train the model**\n\n# Remaining same always","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class config:\n \n MAX_LEN = 224\n TRAIN_BATCH_SIZE = 32\n VALID_BATCH_SIZE = 8\n EPOCHS = 1\n MODEL_PATH = \"model.bin\"\n TRAINING_FILE = '/kaggle/input/jigsaw-multilingual-toxic-comment-classification/jigsaw-toxic-comment-train.csv'\n TOKENIZER = transformers.BertTokenizer.from_pretrained('bert-base-uncased',do_lower_case =True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BERTDataset(torch.utils.data.Dataset):\n def __init__(self,text,target):\n self.text = text\n self.target = target\n self.tokenizer = config.TOKENIZER\n self.max_len = config.MAX_LEN \n\n def __len__(self):\n return len(self.text)\n\n def __getitem__(self,idx):\n text = str(self.text[idx])\n text = \" \".join(text.split())\n\n inputs = self.tokenizer.encode_plus(\n text,\n None,\n add_special_tokens = True,\n max_length = self.max_len,\n pad_to_max_length = True\n )\n\n ids = inputs[\"input_ids\"]\n mask = inputs[\"attention_mask\"]\n\n return {\n 'ids': torch.tensor(ids,dtype=torch.long),\n 'mask': torch.tensor(mask,dtype=torch.long),\n 'targets': torch.tensor(self.target[idx],dtype=torch.long)\n }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BERTBaseUncased(nn.Module):\n def __init__(self):\n super(BERTBaseUncased,self).__init__()\n self.bert = transformers.BertModel.from_pretrained('bert-base-uncased')\n self.bert_drop = nn.Dropout(0.3)\n self.fc1 = nn.Linear(768,1)\n\n def forward(self,ids,mask):\n _,o2 = self.bert(\n ids,\n mask\n )\n bo = self.bert_drop(o2)\n out = self.fc1(bo)\n\n return out","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 eval_fn(data_loader, model, device):\n model.eval()\n fin_targets = []\n fin_outputs = []\n with torch.no_grad():\n for bi, d in enumerate(data_loader):\n ids = d[\"ids\"]\n mask = d[\"mask\"]\n targets = d[\"targets\"]\n\n ids = ids.to(device, dtype=torch.long)\n mask = mask.to(device, dtype=torch.long)\n targets = targets.to(device, dtype=torch.float)\n\n outputs = model(\n ids=ids,\n mask=mask\n )\n fin_targets.extend(targets.cpu().detach().numpy().tolist())\n fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())\n return fin_outputs, fin_targets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv(config.TRAINING_FILE).fillna(\"none\").sample(n=4000)\ntrain_df, valid_df, train_tar, valid_tar = train_test_split(train.comment_text, train.toxic, \n stratify=train.toxic.values, \n random_state=42, \n test_size=0.2, shuffle=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Functions that will change","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_fn(data_loader, model, optimizer, device, scheduler,epoch,num_steps):\n model.train()\n\n for bi, d in enumerate(data_loader):\n ids = d[\"ids\"]\n mask = d[\"mask\"]\n targets = d[\"targets\"]\n\n ids = ids.to(device, dtype=torch.long)\n mask = mask.to(device, dtype=torch.long)\n targets = targets.to(device, dtype=torch.float)\n \n optimizer.zero_grad()\n outputs = model(\n ids=ids,\n mask=mask\n )\n\n loss = loss_fn(outputs, targets)\n loss.backward()\n #--------------------------------#------------------------#----------------------------#--------------------------#\n ####################################### CHANGE HAPPENS HERE #######################################################\n xm.optimizer_step(optimizer,barrier=True)\n ###################################################################################################################\n #-------------------------------#------------------------#----------------------------#---------------------------#\n if scheduler is not None:\n scheduler.step()\n \n \n if (bi+1) % 10 == 0:\n print('Epoch [{}/{}], bi[{}/{}], Loss: {:.4f}' \n .format(epoch+1, 1, bi+1,num_steps, loss.item()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run():\n train_dataset = BERTDataset(\n text=train_df.values,\n target=train_tar.values\n )\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=config.TRAIN_BATCH_SIZE,\n sampler=train_sampler,\n drop_last=True,\n num_workers=4\n )\n\n valid_dataset = BERTDataset(\n text=valid_df.values,\n target=valid_tar.values\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=64,\n sampler=valid_sampler,\n drop_last=False,\n num_workers=4\n )\n \n #-----------------------------#---------------------#-----------------------------------#-----------------------------------#\n ##################################### Change occurs Here ####################################################################\n\n device = xm.xla_device()\n model = BERTBaseUncased()\n model.to(device)\n \n #############################################################################################################################\n #----------------------------#----------------------#------------------------------------#-----------------------------------#\n param_optimizer = list(model.named_parameters())\n no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n optimizer_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 \n lr = 3e-5 * xm.xrt_world_size() #You can or cannot make this change , it will work if not multiplied with xm.xrt_world_size()\n\n num_train_steps = int(len(train_dataset) / config.TRAIN_BATCH_SIZE / xm.xrt_world_size() * config.EPOCHS)\n optimizer = AdamW(optimizer_parameters, lr=lr)\n\n best_accuracy = 0\n for epoch in range(config.EPOCHS):\n train_fn(train_data_loader, model, optimizer, device, scheduler=None,epoch=epoch,num_steps=num_train_steps)\n \n outputs, targets = eval_fn(valid_data_loader, model, device)\n \n outputs = np.array(outputs) >= 0.5\n accuracy = metrics.roc_auc_score(targets, outputs)\n print(f\"AUC_SCORE = {accuracy}\")\n if accuracy > best_accuracy:\n xm.save(model.state_dict(), config.MODEL_PATH)\n best_accuracy = accuracy","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"run()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"What changes did we make ?\n* We declared device as xm.xla_device()\n* We used xm.optimizer_step(optimizer,barrier=True)
\nAnd these were all that was needed for running on a single TPU core.\n\n### Curious About What Barrier=True Does?\n\nTO understand We must know how XLA tensors work internally and how XLA creates graphs and runs operations. XLA tensors are Lazy and their internals differ from CPU and CUDA tensors. CPU and CUDA tensors launch operations immediately or eagerly. XLA tensors, on the other hand, are lazy. They record operations in a graph until the results are needed. Deferring execution like this lets XLA optimize it. Thus Calling xm.optimizer_step(optimizer, barrier=True) at the end of each training iteration causes XLA to execute its current graph and update the modelโs parameters.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Training on Multiple TPU cores\nLet's now move at the most interesting part of running Pytorch on MULTIPLE TPU CORES simultaneously. \n\nWorking with multiple Cloud TPU cores is different than training on a single Cloud TPU core. With a single Cloud TPU core we simply acquired the device and ran the operations using it directly. To use multiple Cloud TPU cores we must use other processes, one per Cloud TPU core. This indirection and multiplicity makes multicore training a little more complex than training on a single core, but it's necessary to maximize performance.\n\nTo understand we must understand four important things offered by XLA- Module\n* xla_multiprocessing\n* spawn() Function \n* ParallelLoader\n* XLA_USE_BF16\n\nWe Will first use these in our model and see how they work and then I will explain what these do","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!export XLA_USE_BF16=1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_fn(data_loader, model, optimizer, device, scheduler,epoch,num_steps):\n model.train()\n\n for bi, d in enumerate(data_loader):\n ids = d[\"ids\"]\n mask = d[\"mask\"]\n targets = d[\"targets\"]\n\n ids = ids.to(device, dtype=torch.long)\n mask = mask.to(device, dtype=torch.long)\n targets = targets.to(device, dtype=torch.float)\n \n optimizer.zero_grad()\n outputs = model(\n ids=ids,\n mask=mask\n )\n\n loss = loss_fn(outputs, targets)\n loss.backward()\n #--------------------------------#------------------------#----------------------------#--------------------------#\n ####################################### CHANGE HAPPENS HERE #######################################################\n xm.optimizer_step(optimizer)\n ###################################################################################################################\n #-------------------------------#------------------------#----------------------------#---------------------------#\n if scheduler is not None:\n scheduler.step()\n \n \n if (bi+1) % 10 == 0:\n print('Epoch [{}/{}], bi[{}/{}], Loss: {:.4f}' \n .format(epoch+1, 1, bi+1,num_steps, loss.item()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = BERTBaseUncased()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _run():\n \n train_dataset = BERTDataset(\n text=train_df.values,\n target=train_tar.values\n )\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=config.TRAIN_BATCH_SIZE,\n sampler=train_sampler,\n drop_last=True,\n num_workers=4\n )\n\n valid_dataset = BERTDataset(\n text=valid_df.values,\n target=valid_tar.values\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=64,\n sampler=valid_sampler,\n drop_last=False,\n num_workers=4\n )\n\n device = xm.xla_device()\n model.to(device)\n \n param_optimizer = list(model.named_parameters())\n no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n optimizer_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 \n lr = 0.4 * 1e-5 * xm.xrt_world_size()\n\n num_train_steps = int(len(train_dataset) / config.TRAIN_BATCH_SIZE / xm.xrt_world_size() * config.EPOCHS)\n optimizer = AdamW(optimizer_parameters, lr=lr)\n\n best_accuracy = 0\n #---------------------------------------#--------------------------------#----------------------------#-------------------------------#\n ########################################## Change occur In this Loop #################################################################\n for epoch in range(config.EPOCHS):\n para_loader = pl.ParallelLoader(train_data_loader, [device])\n train_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=None,epoch=epoch,num_steps=num_train_steps)\n \n para_loader = pl.ParallelLoader(valid_data_loader, [device])\n outputs, targets = eval_fn(para_loader.per_device_loader(device), model, device)\n \n ########################################################################################################################################\n #---------------------------------------#---------------------------------#------------------------------#------------------------------#\n \n outputs = np.array(outputs) >= 0.5\n accuracy = metrics.roc_auc_score(targets, outputs)\n print(f\"AUC_SCORE = {accuracy}\")\n if accuracy > best_accuracy:\n xm.save(model.state_dict(), config.MODEL_PATH)\n best_accuracy = accuracy","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n torch.set_default_tensor_type('torch.FloatTensor')\n a = _run()\n\nFLAGS={}\n#xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### NOTE :- I didn't know about this and got to know while I was preparing this notebook . It's always useful to experiment ๐ \nTo Run on mutliple TPU cores xlm.xla_device() should not be called before spawning i.e before the spawn function . In simple terms xm.xla_device should only be called through spawn function . To understand this point just uncomment this last line and you will get a Runtime error. But when you restart the kernel and don't run the blocks where xlm.xla_device() was called , the code works fine. Try doing this exercise for fun. Why this happens? More on that in a bit","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Things To notice:\n* xm.optimizer_step() does not take a barrier argument this time\n* Model was declared outside the run function and was sent to Xla Device in the run fucntion whereas when using single TPU's we did it simultaneously in one place\n* Something called Paraloader is wrapped around dataloader\n* USE of XLA_USE_BF16 Environment variable\n* And off course we now run the spawn function to execute the model training and eval
\nLet's now talk about each of these one by one . Let's start with:\n\n# XLA_USE_BF16 Environment variable\n\nPyTorch/XLA can use the bfloat16 datatype when running on TPUs. In fact, PyTorch/XLA handles float types (torch.float and torch.double) differently on TPUs. This behavior is controlled by the XLA_USE_BF16 environment variable:\n\n* By default both torch.float and torch.double are torch.float on TPUs.\n* If XLA_USE_BF16 is set, then torch.float and torch.double are both bfloat16 on TPUs.\n* If a PyTorch tensor has torch.bfloat16 data type, this will be directly mapped to the TPU bfloat16 (XLA BF16 primitive type).\n\n# Paraloader\n* ParallelLoader loads the training data onto each device i.e onto each TPU core\n* Wraps an existing PyTorch DataLoader with background data upload.\n\n# Barrier No longer needed\n* xm.optimizer_step(optimizer) no longer needs a barrier. ParallelLoader automatically creates an XLA barrier that evalutes the graph","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Spawn Function\nThis is the most important of all to know how to effectively use multi-processing and Multiple TPU cores. \n* What spawn function does is it creates multiple copies of the computation graphs to be fed to different cores or xla_devices . It also makes copies of the data on which the model is trained upon.\n* spawn() takes a function (the \"map function\"), a tuple of arguments (the placeholder flags dict), the number of processes to create, and whether to create these new processes by \"forking\" or \"spawning.\" \n* In the above code here, spawn() will create eight processes, one for each Cloud TPU core, and call _map_fn() -- the map function -- on each process. The inputs to _map_fn() are an index (zero through seven) and the placeholder flags. When the proccesses acquire their device they actually acquire their corresponding Cloud TPU core automatically.\n\n### Map_function\nLet's now talk about the map function . So it is the function which is called on the replicated n number of processes. Pytorch XLA makes nprocs copies as soon as the spawn function is called , one for each device , then the map function is called the first thing on each of these devices. Map function takes two arguments , one is process index (zero to n) and the placeholder flags which is a dictionary and can contain configuration of your model like max_len, epochs, num_workers,etc\n\n### Now back to why we cannot call xm.xla_device() before spawing?\n* This is because if we do this we aquire single TPU core as our XLA-Device and thus Pytorch-XLA will think that there is a single device and will not be able to use the other TPU cores .\n* While when spawing occrus the TPU cores are called as different devices running the same processes via the index input to the map_function .In pytorch XLA we can grab a specific Core using xm.xla_device(n=) and passing the index. That's how spawn is able to grab different TPU cores \n\n### How did each process in the above cell know to acquire its own Cloud TPU core?\nThe answer is context. Accelerators, like Cloud TPUs, manage their operations using an implicit stateful context. In the cell above, the `spawn()` function creates a multiprocessing context and gives it to each new, forked process, allowing them to coordinate.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Experimentations with TPUs for Jigsaw\n* Now since we know that spawn makes copies of data and processes , it is the reason for a common error of over memory usage when working with large models like XLM-Roberta and Roberta itself .\n* I have tried and tested the following:
\n> \nBERT- BASE UNCASED ------> Maxlength = 224 , bs= 128 epochs=3 works with multi-core without getting out of memory
\nBERT-BASE-MULTILINGUAL ------> Maxlength = 224 , bs=64 epcohs =3 ,works with multi-core without getting out of memory . bs of 128 does not work
\nXLM-Roberta ----> Gets out of memory even with bs =16 and max_len = 96 .
\nIdeal hyperparams for XLM-Roberta for this comp would be maxlebgth = 224 and bs =64 and it works fine on google colab because it gives 25gb of VMS
","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# End Notes\n\nI have tried to explain everything in great detail and I hope this resolves some of the uncertainities pertaining to TPU's . If you find something new or there is some doubt , Please mention it in the comment and we can discuss on it.\n\nI hope my efforts helped you and you have a better clarity regarding Pytorch-XLA and TPU's in general. Upvoting is gesture of appreciation which tells me that my notebook did help you , please be kind enough to upvote","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Further Reading\n* http://pytorch.org/xla/release/1.5/index.html\n* https://github.com/pytorch/xla","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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}