![\"Dropout](\"https://production-media.paperswithcode.com/methods/Screen_Shot_2020-05-23_at_6.19.24_PM.png\")
\n Visualization of how Dropout works\n
\nThis article is the continuation of Optimization approaches for Transformers, where the authors described some optimization approaches, which can significantly reduce memory utilization and time for training or inference of the Deep Learning models, in particular language models (LMs) such as Transformers in the Natural Language Understanding (NLP) tasks.
\nIn the article, we propose to study less known and less used approaches for optimization of language models, but the further described approaches can potentially reduce memory footprint and time by a lot too. Worth noting, that some of the proposed methods can be also successfully used not only with language models, and some others even don't directly belong to the model's optimization.
\nSometimes, the authors use quotes from other sources to provide more contextual information and not forget to note something important.\n
\nThe easiest to implement and at least powerful method for optimization training and inference of language models is pre-tokenization / pre-encoding. The idea behind the pre-tokenization is to tokenize the sequences beforehand neither do it on the fly, i.e during batching in training or inference respectively. Mathematically, such approach can potentially reduce data pre-processing time for training from $ O(n*e) $ to $ O(n) $, where $ n $ - number of sequences, $ e $ - number of epochs. In case of Cross-Validation from $ O(n*e*k) $ to $ O(n*k) $, where $ k $ - number of folds, and for inference from $ O(n*k) $ to $ O(n) $ respectively.
\n Note, that different tokenizers have different tokenization rules (Word Piece, Sentence Piece, Byte Pair Encoding, etc.), so authors recommend checking how the tokenizer works before doing pre-tokenization, for example, Sentence Piece tokenization realizes \"Subword regularization and BPE-dropout\" mechanisms, which add small augmentations during training or inference (similar to Test Time Augmentations), so doing tokenization on the fly with such tokenizers might be helpful to improve accuracy as well as the robustness of the language models.\n
\n>>> import sentencepiece as spm\n>>> s = spm.SentencePieceProcessor(model_file='spm.model')\n>>> for n in range(5):\n... s.encode('New York', out_type=str, enable_sampling=True, alpha=0.1, nbest_size=-1)\n...\n['▁', 'N', 'e', 'w', '▁York']\n['▁', 'New', '▁York']\n['▁', 'New', '▁Y', 'o', 'r', 'k']\n['▁', 'New', '▁York']\n['▁', 'New', '▁York']\n\n\n\nSource: Subword regularization and BPE-dropout\n
","metadata":{}},{"cell_type":"markdown","source":"\nDropout is a regularization technique for neural networks that drops a unit (along with connections) at training time with a specified probability $ p $ (a common value is 0.5). At test time, all units are present, but with weights scaled by $ p $ (i.e. becomes $ w*p $).
\nThe idea is to prevent co-adaptation, where the neural network becomes too reliant on particular connections, as this could be symptomatic of overfitting. Intuitively, dropout can be thought of as creating an implicit ensemble of neural networks.
\nSource: Dropout\n
\nIn spite of the easy implementation of Dropout regularization, it can be a serious bottleneck in large-scale models, because the Dropout implementation is based on masking (choose ~ $ p*100 $ % of neurons) and dot product (multiply selected neurons by zero) operations, hence turn Dropout off (set $ p $ to $ 0.0 $) may lead to quite significant increasing of model's performance during training, however, authors recommend turn Dropout off when there is no overfitting during training and the dataset is relatively large (one of the techniques to prevent overfitting).\n
","metadata":{}},{"cell_type":"markdown","source":"\nThere are a lot of kinds of Dropout: Monte-Carlo Dropout, DropConnect, Gaussian Dropout, and even StableDropout from the HuggingFace Transformers library (used in the DeBERTa model), etc. However, our provided implementation for turning Dropout off only works for Dropout from the PyTorch framework, despite this, the implementation can be easily re-implemented for another specific kind.\n
","metadata":{}},{"cell_type":"code","source":"import torch\nfrom torch import nn\nfrom transformers import AutoModel\nimport gc\n\n\n# extensions\ngc.enable()\n\n# utilities\n@torch.no_grad()\ndef turn_off_dropout(module: nn.Module) -> None:\n if isinstance(module, nn.Dropout):\n module.p = 0.0\n \n# initializing model\nmodel_path = \"distilbert-base-uncased\"\nmodel = AutoModel.from_pretrained(model_path)\n\n# turning Dropout off\nmodel.apply(turn_off_dropout)\nprint(model)\n\n# memory clearing\ndel model\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-09-10T16:14:35.328504Z","iopub.execute_input":"2022-09-10T16:14:35.328882Z","iopub.status.idle":"2022-09-10T16:14:50.370958Z","shell.execute_reply.started":"2022-09-10T16:14:35.328845Z","shell.execute_reply":"2022-09-10T16:14:50.370003Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n Source: TorchScript\n
","metadata":{}},{"cell_type":"markdown","source":"\nDeepSpeed is an open source deep learning optimization library for PyTorch. The library is designed to reduce computing power and memory use and to train large distributed models with better parallelism on existing computer hardware. DeepSpeed is optimized for low latency, high throughput training. It includes the Zero Redundancy Optimizer (ZeRO) for training models with 1 trillion or more parameters. Features include mixed precision training, single-GPU, multi-GPU, and multi-node training as well as custom model parallelism. The DeepSpeed source code is licensed under MIT License and available on GitHub.
\nThe team claimed to achieve up to a 6.2x throughput improvement, 2.8x faster convergence, and 4.6x less communication.\n
\nSource: Wikipedia\n
\nAlso, DeepSpeed functionality is integrated in the popular libraries: PyTorch Lightning, HuggingFace Transformers, HuggingFace Accelerator, etc.\n
\n\"I have tried DeepSpeed and I am able to fit a Longformer Large with BS 6 on my RTX 6000 24GB card rented from Jarvislabs . The best part is integrating DeepSpeed with huggingface trainer is super easy and although there is not much to it , I thought it would really help if I share a code example for it.\" - Tanul Singh's comment on the forum under Feedback Prize - Evaluating Student Writing competition.
","metadata":{}},{"cell_type":"markdown","source":"\n DeepSeed Compression is a library based on DeepSpeed. DeepSpeed Compression provides easy-to-use, fast, and powerful functionality for compression large-scale language models such as BLOOM (~176 billion parameters), YaLM (~100 billion parameters), GPT, etc. DeepSpeed Compression uses modern compression techniques, e.g Quantization, Pruning, Distillation, etc, due to them the accuracy of the compressed model doesn't significantly differ from the original model. A detailed description of the DeepSpeed Compression library can be found in the DeepSpeed Compression developers' blog post - DeepSpeed Compression: A composable library for extreme compression and zero-cost quantization.\n
![](\"https://www.microsoft.com/en-us/research/uploads/prod/2022/07/1400x788_Deepspeed_blog_hero_no_logo_V2-1920x1080.jpg\"){kind=logo}
\n \nThe below implementation doesn't show all functionality of the DeepSpeed library, if you are interested in deeply learning DeepSpeed's API, please refer to the official documentation - DeepSpeed's Documentation. DeepSpeed automatically does Gradient Clipping, Gradient Accumulation, Automatic Mixed Precision, and other approaches (if they are configured in DeepSpeed Configuration), so DeepSpeed library's users do not need to write and integrate them by themselves.\n
","metadata":{}},{"cell_type":"code","source":"!conda install -qq -y mpi4py \n!pip install -qq deepspeed","metadata":{"execution":{"iopub.status.busy":"2022-09-10T16:14:57.167053Z","iopub.execute_input":"2022-09-10T16:14:57.167492Z","iopub.status.idle":"2022-09-10T16:16:20.626308Z","shell.execute_reply.started":"2022-09-10T16:14:57.167456Z","shell.execute_reply":"2022-09-10T16:16:20.625094Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch import nn\nfrom torch.optim import AdamW\nfrom transformers import AutoModel\nimport deepspeed\nfrom typing import Iterator\nimport gc\nimport os\n\n\n# extensions\ngc.enable()\n\n# utilities\nno_decay_parameters = (\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\")\n\ndef get_decay_module_parameters(\n module: nn.Module, \n no_decay_parameters: Iterator[str] = no_decay_parameters,\n recurse: bool = True,\n) -> Iterator[nn.Parameter]:\n \n for name, parameter in list(module.named_parameters(recurse=recurse)):\n if name not in no_decay_parameters:\n yield parameter\n\n \ndef get_no_decay_module_parameters(\n module: nn.Module, \n no_decay_parameters: Iterator[str] = no_decay_parameters,\n recurse: bool = True,\n) -> Iterator[nn.Parameter]:\n\n for name, parameter in list(module.named_parameters(recurse=recurse)):\n if name in no_decay_parameters:\n yield parameter\n\n# config\nmodel_path = \"distilbert-base-uncased\"\nweight_decay = 0.01\nlr = 1e-5\ntrain_batch_size = 1\ngradient_accumulation_steps = 1\ndevice = \"cpu\"\n\n# initializing model\nmodel = AutoModel.from_pretrained(model_path)\n\n# initializing optimizer\nmodel_parameters = [\n {\"params\": get_decay_module_parameters(model), \"weight_decay\": weight_decay},\n {\"params\": get_no_decay_module_parameters(model), \"weight_decay\": 0.0},\n]\n\noptimizer = AdamW(params=model_parameters, lr=lr, weight_decay=weight_decay)\n\n# initializing learning rate scheduler\nscheduler = None\n\n# DeepSpeed config\ndeepspeed_arguments = {}\ndeepspeed_config = {\n \"zero_optimization\": {\n \"offload_param\": {\n \"device\": device,\n }\n },\n \"train_batch_size\": train_batch_size,\n \"gradient_accumulation_steps\": gradient_accumulation_steps,\n}\n\n# wrapping model via DeepSpeed\ndeepspeed.init_distributed(dist_backend=\"gloo\")\nmodel_engine, optimizer, training_dataloader, scheduler = deepspeed.initialize(\n args=deepspeed_arguments,\n model=model.requires_grad_(False), # WARNING: here is the issue I need to figure out.\n optimizer=optimizer,\n model_parameters=None,\n lr_scheduler=scheduler,\n config=deepspeed_config,\n dist_init_required=False,\n)","metadata":{"execution":{"iopub.status.busy":"2022-09-10T16:16:20.628518Z","iopub.execute_input":"2022-09-10T16:16:20.628917Z","iopub.status.idle":"2022-09-10T16:16:52.909709Z","shell.execute_reply.started":"2022-09-10T16:16:20.628877Z","shell.execute_reply":"2022-09-10T16:16:52.907500Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%script false --no-raise-error\n\n# training\nfor step, batch in enumerate(dataloader, 1):\n \n # prepare inputs and targets for the model and loss function respectively.\n \n # forward pass\n outputs = model_engine(inputs)\n \n # computing loss\n loss = loss_fn(outputs, targets)\n \n # backward pass\n model_engine.backward()\n \n # optimization step\n model_engine.step()","metadata":{"execution":{"iopub.status.busy":"2022-09-10T16:16:52.929285Z","iopub.execute_input":"2022-09-10T16:16:52.935221Z","iopub.status.idle":"2022-09-10T16:16:53.094655Z","shell.execute_reply.started":"2022-09-10T16:16:52.935175Z","shell.execute_reply":"2022-09-10T16:16:53.093245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# memory clearing\ndel model, model_engine, optimizer, scheduler, model_parameters\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-09-10T16:16:53.096203Z","iopub.execute_input":"2022-09-10T16:16:53.097001Z","iopub.status.idle":"2022-09-10T16:16:53.584676Z","shell.execute_reply.started":"2022-09-10T16:16:53.096942Z","shell.execute_reply":"2022-09-10T16:16:53.583811Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\nIn progress...\n
","metadata":{}},{"cell_type":"markdown","source":"\n In this paper, the authors described and implemented additional approaches for further optimizing large-scale models, especially Transformers-based models in the NLP tasks. Authors showed that even primitive and simple methods such as Pre-tokenization and Turning Dropout off can significantly reduce training and inference time, then authors investigated more powerful approaches such as TorchScript and DeepSpeed for faster inference and training of large-scale models respectively. Authors will extend this work by the Layers Fusing approach as soon as possible, and also authors want to discover more approaches for optimization in the future.\n
","metadata":{}},{"cell_type":"markdown","source":"\nIf you run into problems with running implementations, want to share personal results with some of the described approaches, or offer your methods for optimization, you can put comments under this article or write a personal message to one of the authors' social networks.\n
\nDuring writing this article, authors referred to this literature:\n
\n