https://community.databricks.com/t5/technical-blog/introducing-the-deepspeed-distributor-on-databricks/ba-p/59641 <P><span class="lia-inline-image-display-wrapper lia-image-align-right" image-alt="torchdistributed.png" style="width: 284px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6205i8AE54D16FF195CD6/image-dimensions/284x284?v=v2" width="284" height="284" role="button" title="torchdistributed.png" alt="torchdistributed.png" /></span><SPAN>The GPU shortage is real, and being able to scale up and optimize the training of large language models will help accelerate the delivery of an AI project.&nbsp;</SPAN><A style="font-family: inherit; background-color: #ffffff;" href="https://github.com/microsoft/DeepSpeed" target="_blank" rel="noopener">DeepSpeed</A> <SPAN>is a framework that can reduce GPU memory requirements by up to 4x.&nbsp;Setting up and triggering DeepSpeed jobs is complex, however, and Databricks is here to help.&nbsp;That is why we are pleased to announce the release of the new DeepSpeed Distributor,&nbsp;a fully open source extension to Spark that builds upon our work with both the <A href="https://docs.databricks.com/en/machine-learning/train-model/distributed-training/spark-pytorch-distributor.html" target="_self">TorchDistributor</A> and the <A href="https://docs.databricks.com/en/machine-learning/train-model/distributed-training/spark-tf-distributor.html" target="_self">spark-tensorflow-distributor</A>,&nbsp;Now you can trigger DeepSpeed powered training jobs with a simple PySpark function - the hassles of cluster configuration, communications and monitoring are all handled for you! In this article, we review the key features of DeepSpeed that you will need to advance your AI journey in the development of large language models (LLM) and computer vision (CV) applications.</SPAN></P> <DIV class="mce-toc"> <H2>Table of Contents</H2> <UL> <LI><A href="#mcetoc_1hmhf6r923f" target="_blank">Overcoming GPU bottlenecks</A></LI> <LI><A href="#mcetoc_1hmhf6r923g" target="_blank">Understanding DeepSpeed and the DeepSpeed Distributor</A> <UL> <LI><A href="#mcetoc_1hmhf6r923h" target="_blank">Model Precision Settings</A></LI> <LI><A href="#mcetoc_1hmhf6r923i" target="_blank">Optimizer and Gradient Settings</A></LI> </UL> </LI> <LI><A href="#mcetoc_1hmhf6r923j" target="_blank">How to trigger a DeepSpeed training job</A> <UL> <LI><A href="#mcetoc_1hmhf6r923k" target="_blank">Triggering a function from notebooks</A></LI> <LI><A href="#mcetoc_1hmhf6r923l" target="_blank">Triggering a Python script</A></LI> </UL> </LI> <LI><A href="#mcetoc_1hmhf6r923m" target="_blank">Tricks and Tips to be productive</A> <UL> <LI><A href="#mcetoc_1hmhf6r923n" target="_blank">Estimating Memory Requirements</A></LI> <LI><A href="#mcetoc_1hmhf6r923o" target="_blank">Setting Hugging Face Caches</A></LI> <LI><A href="#mcetoc_1hmhf6r923p" target="_blank">Configuring MLflow with DeepSpeed</A></LI> <LI><A href="#mcetoc_1hmhf6r923q" target="_blank">Tensorboard Integration</A></LI> <LI><A href="#mcetoc_1hmhf6r923r" target="_blank">Gradient Accumulation</A></LI> <LI><A href="#mcetoc_1hmhf6r923s" target="_blank">Gradient Checkpointing</A></LI> <LI><A href="#mcetoc_1hmhf6r923t" target="_blank">Finding the optimal configuration</A></LI> <LI><A href="#mcetoc_1hmhf6r923u" target="_blank">Monitoring and debugging</A></LI> </UL> </LI> <LI><A href="#mcetoc_1hmhf6r923v" target="_blank">Conclusion</A></LI> </UL> </DIV> <H2 id="mcetoc_1hmhf6r923f">Overcoming GPU bottlenecks</H2> <P>The major bottleneck in deep learning training today is the amount of RAM per GPU - they simply don’t have a lot of it.</P> <TABLE style="height: 400px; border-style: solid; width: 100%;" border="1" width="100%"> <TBODY> <TR> <TD width="25%" height="50px"> <P><STRONG>GPU Model</STRONG></P> </TD> <TD width="25%" height="50px"> <P><STRONG>VRAM (GB)</STRONG></P> </TD> <TD width="25%" height="50px"> <P><STRONG>AWS</STRONG></P> </TD> <TD width="25%" height="50px"> <P><STRONG>Azure</STRONG></P> </TD> </TR> <TR> <TD width="25%" height="50px"> <P><SPAN>T4</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>16</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>g4dn</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>NCasT4_v3</SPAN></P> </TD> </TR> <TR> <TD width="25%" height="50px"> <P><SPAN>V100</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>16</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>p3</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>NC_v3</SPAN></P> </TD> </TR> <TR> <TD width="25%" height="50px"> <P><SPAN>V100</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>32</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>p3dn</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>NDv2</SPAN></P> </TD> </TR> <TR> <TD width="25%" height="50px"> <P><SPAN>A10</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>24</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>g5</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>NVadsA10_v5</SPAN></P> </TD> </TR> <TR> <TD width="25%" height="50px"> <P><SPAN>A100</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>40</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>p4d</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>NDasrA100_v4</SPAN></P> </TD> </TR> <TR> <TD width="25%" height="50px"> <P><SPAN>A100</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>80</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>p4de</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>NDm_A100_v4</SPAN></P> </TD> </TR> <TR> <TD width="25%" height="50px"> <P><SPAN>H100</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>80</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>p5</SPAN></P> </TD> <TD width="25%" height="50px"> <P><SPAN>NDm_H100_v5</SPAN></P> </TD> </TR> </TBODY> </TABLE> <P class="lia-align-center">Common GPUs - their instance types and available RAM</P> <P>To address this issue, two main paradigms have emerged: <EM>data parallel</EM> and <EM>model parallel</EM> training.</P> <P>Data parallel training involves splitting the data across multiple GPUs. In this scenario, each GPU needs to contain a full copy of the model and training artifacts (the optimizer state and gradients), greatly reducing available RAM for training data.</P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Data Parallel Training - We still need all the Training Artifacts and the Full Model which takes up most of our VRAM" style="width: 712px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6183i3B43FEF45E17CF3C/image-size/large?v=v2&amp;px=999" role="button" title="brian_law_0-1707658364567.png" alt="Data Parallel Training - We still need all the Training Artifacts and the Full Model which takes up most of our VRAM" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">Data Parallel Training - We still need all the Training Artifacts and the Full Model which takes up most of our VRAM</span></span></P> <P><SPAN>On the other hand, model parallel training splits the model itself across multiple GPUs, freeing up RAM on each machine.&nbsp; In this scenario a new bottleneck appears - the available network bandwidth between the distributed GPUs. </SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Model Parallel Training - We have now split the model up but that means we need to get the data through all three GPUs to get a full forward and backward pass" style="width: 686px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6190i26706B8ACD670384/image-size/large?v=v2&amp;px=999" role="button" title="brian_law_0-1707659684027.png" alt="Model Parallel Training - We have now split the model up but that means we need to get the data through all three GPUs to get a full forward and backward pass" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">Model Parallel Training - We have now split the model up but that means we need to get the data through all three GPUs to get a full forward and backward pass</span></span></P> <P><SPAN>In the model parallel paradigm, the memory efficiency of DeepSpeed can be leveraged to allow us to train larger models without incurring exorbitant costs.</SPAN></P> <H2 id="mcetoc_1hmhf6r923g"><SPAN>Understanding DeepSpeed and the DeepSpeed Distributor</SPAN></H2> <P><SPAN>The DeepSpeed library consists of three key modules: the training of, inference with and compression of deep learning models. In this article we focus on model training.&nbsp;</SPAN></P> <P><SPAN>To understand how DeepSpeed fits and how to configure it, we must first unpack the training process. There are two constraints when training large language models and their computer vision cousins: the input / output speeds for loading training data and the amount of memory available per GPU card.&nbsp;For data loading optimizations, refer to our <A href="https://www.databricks.com/blog/contributing-spark-loader-for-hugging-face-datasets" target="_blank" rel="noopener">collaboration with Hugging Face</A> and the <A href="https://delta.io/blog/deltatorch-intro/" target="_self">Torchdelta</A>&nbsp;extension.</SPAN></P> <P><STRONG>For model training, there are three main memory hogs: the model and its weights, the optimizer state and gradients, and the size of the training data. </STRONG><SPAN>To reduce training time, we want to maximize the amount of data we load on the GPU. This requires configuring DeepSpeed to limit the RAM footprint of the model, optimizer and gradients. Let’s walk through those settings now.</SPAN></P> <H3 id="mcetoc_1hmhf6r923h"><SPAN>Model Precision Settings</SPAN></H3> <P><SPAN>To ensure a high quality chat experience for end users of our LLM applications we want to use the largest models we can. By default, model weights are stored as 32-bit floating point, so even 7 billion parameter models take up all the RAM on smaller GPUs and a considerable chunk on larger ones. </SPAN><STRONG>DeepSpeed includes optimizations that allow us to train at 16-bit floating point instead, greatly reducing model size. </STRONG><SPAN>When using 16-bit floating point precision, a good rule of thumb is that an average model requires approximately 2x its number of parameters (billions) in GPU RAM (GB) rather than 4x with the default </SPAN><I><SPAN>fp32</SPAN></I><SPAN>. Whilst it is possible to use lower precision, these techniques are still more experimental.&nbsp;</SPAN></P> <P><SPAN>There are two types of 16-bit precision to choose from: the normal <I>fp16</I> and <I>bfloat16</I>. When working with newer Nvidia GPUs (A10/A100/H100) it is recommended to use <I>bfloat16</I>, which can be set as follows in the JSON config file.</SPAN></P> <UL class="lia-list-style-type-circle"> <LI><LI-CODE lang="python">"bf16": { "enabled": true }</LI-CODE></LI> </UL> <UL class="lia-list-style-type-circle"> <LI> <P><SPAN>Older GPUs do not support </SPAN><I><SPAN>bfloat16</SPAN></I><SPAN> and hence we will fallback to the older less efficient </SPAN><I><SPAN>fp16</SPAN></I><SPAN> format, shown below.</SPAN></P> <LI-CODE lang="python">"fp16": { "enabled": "auto", "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, "hysteresis": 2, "min_loss_scale": 1 }</LI-CODE></LI> </UL> <P><STRONG>We recommend using </STRONG><STRONG><I>bfloat</I></STRONG><STRONG> compatible GPUs where possible.</STRONG><SPAN> It has been noted that </SPAN><I><SPAN>fp16</SPAN></I><SPAN> training can be unstable and may require manual intervention during the training to change learning rates and avoid </SPAN><I><SPAN>NaN</SPAN></I><SPAN> losses which is partially why </SPAN><I><SPAN>fp16</SPAN></I><SPAN> also has so many additional settings.; For more details check the </SPAN><A href="https://www.deepspeed.ai/docs/config-json/#bfloat16-training-options" target="_blank" rel="noopener"><SPAN>documentation</SPAN></A><SPAN> and Nvidia's guide on </SPAN><A href="https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#training" target="_blank" rel="noopener"><SPAN>mixed precision training</SPAN></A></P> <H3 id="mcetoc_1hmhf6r923i"><SPAN>Optimizer and Gradient Settings</SPAN></H3> <P><SPAN>The next area that we can look for memory savings is with the optimizer and gradient calculations. This can be done through leveraging the Zero Redundancy Optimizer (ZeRO) settings. ZeRO is the core of DeepSpeed and was first introduced in the </SPAN><A href="https://arxiv.org/pdf/1910.02054.pdf" target="_blank" rel="noopener"><SPAN>research paper</SPAN></A><SPAN>. The authors noted that a lot of the weights, states and gradients stored during training are redundant, allowing for them to be partitioned across GPUs and even onto CPU and disk. This does result in the need to move data around and increases network traffic, but the savings in GPU RAM are worth it.</SPAN></P> <P><SPAN>ZeRO has three settings:</SPAN><SPAN><BR /></SPAN></P> <UL> <LI style="font-weight: 400;" aria-level="1"><SPAN>Stage 1: Optimizer State Partitioning&nbsp;</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><SPAN>Stage 2: Optimizer and Gradient Partitioning&nbsp;</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><SPAN>Stage 3: Optimizer, Gradient and Parameter Partitioning&nbsp;</SPAN></LI> </UL> <P><SPAN>Beyond stage 3 we can also offload calculations onto the CPU and attached storage, though this can come with hefty performance constraints. For an understanding of the possible space savings, see the below diagram from the original ZeRO paper where Stage 1 = Pos, Stage 2 = Pos+g and Stage 3 = Pos+g+p.</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="brian_law_0-1707660628666.png" style="width: 521px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6191i72AAFBCF40510A28/image-dimensions/521x306?v=v2" width="521" height="306" role="button" title="brian_law_0-1707660628666.png" alt="brian_law_0-1707660628666.png" /></span></P> <P><SPAN>As we go from Stages 1 to 3, we can see that we get more and more VRAM savings. Whilst the formulas for calculating VRAM consumption are a little complex, here is a visualization of estimated VRAM requirements per GPU.&nbsp;</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Calculations assume 4 GPUs we can reduce further with" style="width: 550px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6192i82311D5F30F7DD79/image-dimensions/550x349?v=v2" width="550" height="349" role="button" title="brian_law_0-1707660684950.png" alt="Calculations assume 4 GPUs we can reduce further with" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">Calculations assume 4 GPUs we can reduce further with</span></span></P> <TABLE border="1" width="100%"> <TBODY> <TR> <TD width="32.5609756097561%"> <P><STRONG>Model Parameter Count (B)</STRONG></P> </TD> <TD width="20.975609756097562%"> <P><STRONG>Base VRAM (GB)</STRONG></P> </TD> <TD width="15.487804878048781%"> <P><STRONG>ZeRo 1 (GB)</STRONG></P> </TD> <TD width="15.487804878048781%"> <P><STRONG>ZeRo 2 (GB)</STRONG></P> </TD> <TD width="15.487804878048781%"> <P><STRONG>ZeRo 3 (GB)</STRONG></P> </TD> </TR> <TR> <TD width="32.5609756097561%"> <P><SPAN>7</SPAN></P> </TD> <TD width="20.975609756097562%"> <P><SPAN>112</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>49</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>38.5</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>28</SPAN></P> </TD> </TR> <TR> <TD width="32.5609756097561%"> <P><SPAN>13</SPAN></P> </TD> <TD width="20.975609756097562%"> <P><SPAN>208</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>91</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>71.5</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>52</SPAN></P> </TD> </TR> <TR> <TD width="32.5609756097561%"> <P><SPAN>34</SPAN></P> </TD> <TD width="20.975609756097562%"> <P><SPAN>544</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>238</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>187</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>136</SPAN></P> </TD> </TR> <TR> <TD width="32.5609756097561%"> <P><SPAN>70</SPAN></P> </TD> <TD width="20.975609756097562%"> <P><SPAN>480</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>210</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>165</SPAN></P> </TD> <TD width="15.487804878048781%"> <P><SPAN>120</SPAN></P> </TD> </TR> </TBODY> </TABLE> <P><SPAN>Beyond ZeRo 3, it is also possible to offload the optimizer states, gradients and parameters into either CPU RAM or attached NVMe drives. Whilst offload will cause a performance hit, it can allow small setups to train much larger models. </SPAN></P> <UL class="lia-list-style-type-circle"> <LI><SPAN><SPAN><SPAN>Offload can be set in the JSON configuration:<BR /></SPAN></SPAN></SPAN><LI-CODE lang="python"># cpu offload "offload_optimizer": {"device": "cpu"}, "offload_param": {"device": "cpu"},</LI-CODE><LI-CODE lang="python"># gpu offload "offload_optimizer": {"device": "nvme", "nvme_path": "/local_disk0/optimizer"}, "offload_param": {"device": "nvme", "nvme_path": "/local_disk0/param"}</LI-CODE></LI> </UL> <P><SPAN>On Databricks clusters, </SPAN><I><SPAN>/local_disk0</SPAN></I><SPAN> is the mount path for direct attached storage that you can use as a temporary cache. Offload also requires the installation of </SPAN><I><SPAN>libaio-dev</SPAN></I><SPAN>, an OS-level library that will require installing via init script (</SPAN><A href="https://docs.databricks.com/en/init-scripts/index.html" target="_blank" rel="noopener"><SPAN>AWS</SPAN></A><SPAN>/</SPAN><A href="https://learn.microsoft.com/en-us/azure/databricks/init-scripts/" target="_blank" rel="noopener"><SPAN>Azure</SPAN></A><SPAN>). When using offload, it is important to make sure that there is sufficient CPU RAM and NVMe storage. It is recommended to set the </SPAN><I><SPAN>autoscale local storage</SPAN></I><SPAN> setting when creating your cluster if you will use NVMe storage.</SPAN><SPAN><BR /></SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="brian_law_0-1707660898074.png" style="width: 400px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6193i384B81E42F7F82B6/image-size/medium?v=v2&amp;px=400" role="button" title="brian_law_0-1707660898074.png" alt="brian_law_0-1707660898074.png" /></span></P> <H2 id="mcetoc_1hmhf6r923j"><SPAN>How to trigger a DeepSpeed training job</SPAN></H2> <P><SPAN>Now that we understand the workings of key DeepSpeed configs, we can look at how to code and launch our pipelines. The beauty of the design around the DeepSpeed Distributor is that it doesn't require extensive code changes in order to leverage it within your workflows. Just like the TorchDistributor, the DeepSpeed flavor can be utilized in two ways:</SPAN><SPAN><BR /></SPAN></P> <UL> <LI style="font-weight: 400;" aria-level="1"><SPAN>Trigger a function written within the existing notebook context&nbsp;</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><SPAN>Trigger an external Python script</SPAN></LI> </UL> <H3 id="mcetoc_1hmhf6r923k"><SPAN>Triggering a function from notebooks</SPAN></H3> <P><SPAN>In this example, we assume that the training function is called </SPAN><I><SPAN>train</SPAN></I><SPAN> and it accepts the arguments </SPAN><I><SPAN>training_arguments</SPAN></I><SPAN> and </SPAN><I><SPAN>dataset.</SPAN></I></P> <UL class="lia-list-style-type-circle"> <LI><SPAN><SPAN>Running on single node:<BR /></SPAN></SPAN><LI-CODE lang="python">from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor distributor = DeepspeedTorchDistributor(numGpus=1, nnodes=1, localMode=True, useGpu=True, deepspeedConfig = deepspeed_dict) completed_trainer = distributor.run(train, training_arguments, dataset)</LI-CODE></LI> <LI> <P><SPAN>Running in distributed fashion:<BR /></SPAN></P> <LI-CODE lang="python">from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor distributor = DeepspeedTorchDistributor(numGpus=4, nnodes=2, localMode=True, useGpu=True, deepspeedConfig = deepspeed_dict) completed_trainer = distributor.run(train, training_arguments, dataset)</LI-CODE></LI> </UL> <P><SPAN>To elaborate on the parameters available to </SPAN><I><SPAN>DeepSpeedTorchDistributor</SPAN></I><SPAN>:</SPAN></P> <UL> <LI style="font-weight: 400;" aria-level="1"><I><SPAN>numGpus</SPAN></I><SPAN>: number of GPUs per node</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><I><SPAN>nnodes</SPAN></I><SPAN>: number of nodes to train on</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><I><SPAN>localMode</SPAN></I><SPAN>: if True, trains on the driver node only (for single node cluster) if False will train on the workers</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><I><SPAN>useGpu</SPAN></I><SPAN>: whether to train using GPUs</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><I><SPAN>deepspeedConfig</SPAN></I><SPAN>: JSON file with all configuration parameters</SPAN></LI> </UL> <P><SPAN>While </SPAN><I><SPAN>useGPU</SPAN></I><SPAN> is available, it is worth mentioning that DeepSpeed should ideally be used with GPUs.</SPAN></P> <H3 id="mcetoc_1hmhf6r923l"><SPAN>Triggering a Python script</SPAN></H3> <P><SPAN>In addition to executing a </SPAN><I><SPAN>train</SPAN></I><SPAN> function defined in a notebook, it is possible to use </SPAN><I><SPAN>DeepspeedTorchDistributor</SPAN></I><SPAN> to load and execute a training script by specifying the path to the Python file and passing command line arguments to it.</SPAN></P> <UL class="lia-list-style-type-circle"> <LI><SPAN><SPAN><SPAN>Running on single node:<BR /></SPAN></SPAN></SPAN><LI-CODE lang="python">from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor distributor = DeepspeedTorchDistributor(numGpus=4, nnodes=1, localMode=False, useGpu=True, deepspeedConfig = deepspeed_dict) completed_trainer = distributor.run('&lt;path_to_file&gt;/train.py', "--var1 value1", "--var2 value2")</LI-CODE></LI> <LI><SPAN><SPAN><SPAN>Running in distributed fashion:<BR /></SPAN></SPAN></SPAN><LI-CODE lang="python">from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor distributor = DeepspeedTorchDistributor(numGpus=4, nnodes=2, localMode=False, useGpu=True, deepspeedConfig = deepspeed_dict) completed_trainer = distributor.run('&lt;path_to_file&gt;/train.py', "--var1 value1", "--var2 value2")</LI-CODE></LI> </UL> <H2 id="mcetoc_1hmhf6r923m"><SPAN>Tricks and Tips to be productive</SPAN></H2> <P><SPAN>Now that you understand the key <I>DeepSpeed </I>configs and how to launch your training job, we will cover some other considerations that you should be aware of to ensure the smoothest possible experience.</SPAN></P> <H3 id="mcetoc_1hmhf6r923n"><SPAN>Estimating Memory Requirements</SPAN></H3> <P><SPAN>Managing and understanding the requirements to fine tune a particular model can be hard. There are a few helpful rules of thumb first to assess the viability of our training job before consuming any GPU resources.&nbsp;</SPAN></P> <P><SPAN>As discussed, the number of parameters in a model will dictate the amount of VRAM it and its training states will take. If we are training with </SPAN><I><SPAN>fp16</SPAN></I><SPAN> or </SPAN><I><SPAN>bfloat16</SPAN></I><SPAN> we can use a simple formula:&nbsp;</SPAN></P> <P><STRONG>2 X number of params (B) = VRAM required (GB)</STRONG></P> <P><SPAN>Thus, a 7 billion parameter model like </SPAN><I><SPAN>llama_v2_7b</SPAN></I><SPAN> will take approximately 14GB VRAM to load in </SPAN><I><SPAN>bfloat16</SPAN></I><SPAN> format.&nbsp;</SPAN></P> <P><SPAN>Next, we need to factor in the type of fine-tuning that we are doing. The memory footprint for model weights, weight gradients, adapter weights and optimizer state will depend on whether we are using LoRa, QLoRa, or full fine-tuning (excluding activations / input gradients).&nbsp; The total memory usage per parameter is summarized in the table below.</SPAN></P> <TABLE border="1" width="99.8780487804878%"> <TBODY> <TR> <TD width="19.878048780487806%"> <P><STRONG>Tuning Method</STRONG></P> </TD> <TD width="9.268292682926829%"> <P><STRONG>Weights</STRONG></P> </TD> <TD width="18.170731707317074%"> <P><STRONG>Weight Gradients</STRONG></P> </TD> <TD width="16.21951219512195%"> <P><STRONG>Optimizer State</STRONG></P> </TD> <TD width="17.4390243902439%"> <P><STRONG>Adapter Weights</STRONG></P> </TD> <TD width="18.902439024390244%"> <P><STRONG>Total</STRONG></P> </TD> </TR> <TR> <TD width="19.878048780487806%"> <P><SPAN>Full Fine-Tuning</SPAN></P> </TD> <TD width="9.268292682926829%"> <P><SPAN>2 bytes</SPAN></P> </TD> <TD width="18.170731707317074%"> <P><SPAN>2 bytes</SPAN></P> </TD> <TD width="16.21951219512195%"> <P><SPAN>8 bytes</SPAN></P> </TD> <TD width="17.4390243902439%"> <P><SPAN>N/A</SPAN></P> </TD> <TD width="18.902439024390244%"> <P><SPAN>12 bytes/parameter</SPAN></P> </TD> </TR> <TR> <TD width="19.878048780487806%"> <P><SPAN>LoRa Fine-Tuning</SPAN></P> </TD> <TD width="9.268292682926829%"> <P><SPAN>16 bits</SPAN></P> </TD> <TD width="18.170731707317074%"> <P><SPAN>~0.4 bits</SPAN></P> </TD> <TD width="16.21951219512195%"> <P><SPAN>~0.8 bits</SPAN></P> </TD> <TD width="17.4390243902439%"> <P><SPAN>~0.4 bits</SPAN></P> </TD> <TD width="18.902439024390244%"> <P><SPAN>17.6 bits/parameter</SPAN></P> </TD> </TR> <TR> <TD width="19.878048780487806%"> <P><SPAN>QLoRa Fine-Tuning*</SPAN></P> </TD> <TD width="9.268292682926829%"> <P><SPAN>4 bits</SPAN></P> </TD> <TD width="18.170731707317074%"> <P><SPAN>~0.4 bits</SPAN></P> </TD> <TD width="16.21951219512195%"> <P><SPAN>~0.8 bits</SPAN></P> </TD> <TD width="17.4390243902439%"> <P><SPAN>~0.4 bits</SPAN></P> </TD> <TD width="18.902439024390244%"> <P><SPAN>5.6 bits/parameter</SPAN></P> </TD> </TR> </TBODY> </TABLE> <P class="lia-align-center"><I><SPAN>*Note that currently QLoRa and DeepSpeed ZeRo 3 are incompatible.</SPAN></I></P> <P class="lia-align-left">Combining our formula for model size memory requirements with the additional fine-tuning memory requirements, we can estimate the total RAM required for a training run:</P> <TABLE border="1" width="99.8780487804878%"> <TBODY> <TR> <TD width="13.536585365853659%"> <P><STRONG>Model</STRONG></P> </TD> <TD width="11.829268292682926%"> <P><STRONG>Model Size (</STRONG><STRONG><I>fp32</I></STRONG><STRONG>)</STRONG></P> </TD> <TD width="15.121951219512194%"> <P><STRONG>Model Size (</STRONG><STRONG><I>bfloat16</I></STRONG><STRONG>)</STRONG></P> </TD> <TD width="24.75609756097561%"> <P><STRONG>Full Finetune Requirements (</STRONG><STRONG><I>bfloat16</I></STRONG><STRONG>)</STRONG></P> </TD> <TD width="16.951219512195124%"> <P><STRONG>LoRa Requirements</STRONG></P> </TD> <TD width="17.682926829268293%"> <P><STRONG>QLoRa Requirements</STRONG></P> </TD> </TR> <TR> <TD width="13.536585365853659%"> <P><SPAN>llama_v2_7b</SPAN></P> </TD> <TD width="11.829268292682926%"> <P><SPAN>28GB</SPAN></P> </TD> <TD width="15.121951219512194%"> <P><SPAN>~14GB</SPAN></P> </TD> <TD width="24.75609756097561%"> <P><SPAN>~84GB</SPAN></P> </TD> <TD width="16.951219512195124%"> <P><SPAN>~15.4GB</SPAN></P> </TD> <TD width="17.682926829268293%"> <P><SPAN>~4.55GB</SPAN></P> </TD> </TR> <TR> <TD width="13.536585365853659%"> <P><SPAN>llama_v2_13b</SPAN></P> </TD> <TD width="11.829268292682926%"> <P><SPAN>52GB</SPAN></P> </TD> <TD width="15.121951219512194%"> <P><SPAN>~25GB</SPAN></P> </TD> <TD width="24.75609756097561%"> <P><SPAN>~150GB</SPAN></P> </TD> <TD width="16.951219512195124%"> <P><SPAN>~28.6GB</SPAN></P> </TD> <TD width="17.682926829268293%"> <P><SPAN>~8.45GB</SPAN></P> </TD> </TR> <TR> <TD width="13.536585365853659%"> <P><SPAN>MNPT_7b</SPAN></P> </TD> <TD width="11.829268292682926%"> <P><SPAN>28GB</SPAN></P> </TD> <TD width="15.121951219512194%"> <P><SPAN>~14GB</SPAN></P> </TD> <TD width="24.75609756097561%"> <P><SPAN>~84GB</SPAN></P> </TD> <TD width="16.951219512195124%"> <P><SPAN>~15.4GB</SPAN></P> </TD> <TD width="17.682926829268293%"> <P><SPAN>~4.55GB</SPAN></P> </TD> </TR> </TBODY> </TABLE> <P><SPAN>Now that we size up the viability of our job with our available resources, we can look into the rest of the training setup.</SPAN></P> <H3 id="mcetoc_1hmhf6r923o"><SPAN>Setting Hugging Face Caches</SPAN></H3> <P><SPAN>When using Hugging Face libraries it is important to understand how it caches. Caching can generally be set at an individual item level (a specific call to dataset or model), at a library level (for all dataset or transformer objects) or at a global level (for all Hugging Face libraries).</SPAN></P> <P><SPAN>There is also a hierarchy in that setting a cache when instantiating an item will override the global default. All the Hugging Face libraries default to <I>~/.cache</I> which is the root OS folder on databricks nodes. This can quickly result in the root OS folder filling up and causing the cluster to crash. <STRONG>To alleviate pressure on root, it is recommended to set the cache directories to either a DBFS directory or /</STRONG><STRONG><I>local_disk0</I></STRONG><STRONG> path.</STRONG> This can be configured through environment variables to ensure they are set prior to instantiating your model and dataset classes.</SPAN></P> <UL class="lia-list-style-type-circle"> <LI><EM><EM><EM>Configuring environment variables:<BR /></EM></EM></EM><LI-CODE lang="python">import os # Setting Hugging Face cache to /local_disk0 os.environ['HF_HOME'] = '/local_disk0/hf_home' os.environ['HF_DATASETS_CACHE'] = '/local_disk0/hf_home' os.environ['TRANSFORMERS_CACHE'] = '/local_disk0/hf_home'</LI-CODE></LI> </UL> <P><SPAN>Whilst caching on DBFS guarantees persistence even after clusters get shut down, caching on </SPAN><I><SPAN>local_disk0</SPAN></I><SPAN> can offer better performance for frequently accessed data like a Hugging Face dataset. </SPAN><STRONG>A rule of thumb is to cache transformer models to DBFS since they are loaded only once, while datasets can be cached to </STRONG><STRONG><I>`local_disk0`</I></STRONG><STRONG>.</STRONG></P> <H3 id="mcetoc_1hmhf6r923p"><STRONG><SPAN>Configuring MLflow with DeepSpeed</SPAN></STRONG></H3> <P><SPAN>MLflow has strong support for transformers, and we recommend installing the latest version (</SPAN><A href="https://mlflow.org/docs/latest/index.html" target="_blank" rel="noopener"><SPAN>2.10.1</SPAN></A><SPAN>) for the best experience. It is beyond this blog to explain the basic concepts of MLflow, but there are a few tricks to using it with DeepSpeed that we will cover.&nbsp;</SPAN></P> <P>When operating in a standard notebook environment, the Python session is initiated with a login token for MLflow. When running DeepSpeed, however, individual GPUs will each have a separate Python process that does not inherit these credentials. To proceed, we can save these parameters to Python variables using <A href="https://docs.databricks.com/en/dev-tools/databricks-utils.html" target="_blank" rel="noopener"><I>dbutils</I></A>, then assign them to environment variables within the function that <I>DeepspeedTorchDistributor</I> will distribute.</P> <UL class="lia-list-style-type-circle"> <LI><EM>Getting and settings Databricks host, token:<BR /></EM><LI-CODE lang="python"># logging host and token import os browser_host = dbutils.notebook.entry_point.getDbutils().notebook().getContext().browserHostName().get() db_host = f"https://{browser_host}" db_token = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get() os.environ['DATABRICKS_HOST'] = db_host os.environ['DATABRICKS_TOKEN'] = db_token</LI-CODE></LI> </UL> <P><SPAN>We also need to create and set the MLflow experiment, as only the primary notebook process can automatically create new experiments. Finally, when doing our MLflow logging, we should ensure that it is run only from the head node of the training process. The head node is the one that has <EM>global_rank</EM> 0 in Nvidia parlance.&nbsp;</SPAN></P> <UL class="lia-list-style-type-circle"> <LI><SPAN><SPAN><SPAN>Getting GPU rank with PyTorch<BR /></SPAN></SPAN></SPAN><LI-CODE lang="python">global_rank = torch.distributed.get_rank()</LI-CODE></LI> </UL> <H3 id="mcetoc_1hmhf6r923q"><SPAN>Tensorboard Integration</SPAN></H3> <P><SPAN>Whilst MLflow provides logging and monitoring capabilities for visualizing training runs, Tensorboard can still be useful for monitoring your run, and can be launched in Databricks by running the following commands.</SPAN></P> <UL class="lia-list-style-type-circle"> <LI>Launching Tensorboard in Databricks Notebooks:<BR /><LI-CODE lang="python">%load_ext tensorboard # This sets up our tensorboard settings experiment_log_dir = &lt;log-directory&gt;</LI-CODE><LI-CODE lang="python">%tensorboard --logdir $experiment_log_dir # This starts Tensorboard</LI-CODE></LI> </UL> <P><SPAN>In order to use Tensorboard during a training run, execute this code in its own notebook. Keep in mind that when writing logs to DBFS the results may not be viewable in Tensorboard until the run is finished.</SPAN></P> <H3 id="mcetoc_1hmhf6r923r"><SPAN>Gradient Accumulation</SPAN></H3> <P><SPAN>If we didn’t have GPU VRAM constraints then we would set larger batches to get through our training data faster. Due to VRAM constraints we commonly are restricted to batch sizes in the single digits. Gradient accumulation tells the algorithm to avoid doing a weight update till it has gone through sufficient data records. We can hence reduce our VRAM requirements at a small cost to performance.</SPAN></P> <UL class="lia-list-style-type-circle"> <LI>Setting gradient accumulation in the DeepSpeed JSON config:<BR /><LI-CODE lang="python">"gradient_accumulation_steps": 4</LI-CODE></LI> </UL> <H3 id="mcetoc_1hmhf6r923s"><SPAN>Gradient Checkpointing</SPAN></H3> <P><SPAN>As we discussed earlier, the gradient calculations during a training run are one major source of VRAM consumption. By default, we store all the activation values for our network on the forward pass since they will be used for the backward pass. Gradient checkpointing strategically stores only a portion of the activations. It does mean that inevitably some will have to be recalculated to slow down our training loop (for more details see <A href="https://arxiv.org/abs/1604.06174" target="_blank" rel="noopener">here</A>).&nbsp;<BR /></SPAN></P> <UL class="lia-list-style-type-circle"> <LI><SPAN><SPAN><SPAN>Turning on gradient checkpointing:<BR /></SPAN></SPAN></SPAN><LI-CODE lang="python">model.gradient_checkpointing_enable()</LI-CODE></LI> </UL> <P><STRONG>In our experiments, using gradient checkpointing with </STRONG><STRONG><I>ZeRo 3</I></STRONG><STRONG> offload on 4x A10Gs saved approximately 6GB VRAM per GPU. </STRONG></P> <H3 id="mcetoc_1hmhf6r923t"><STRONG><SPAN>Finding the optimal configuration</SPAN></STRONG></H3> <P>The goal of configuring your DeepSpeed training loop is to have it run as efficiently as possible. One of the most frustrating things that can happen is an out-of-memory error deep into your training run.</P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Typical out-of-memory (OOM) exception on GPUs" style="width: 917px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6198iB98DBEA87DCCDD31/image-size/large?v=v2&amp;px=999" role="button" title="brian_law_0-1707741306041.png" alt="Typical out-of-memory (OOM) exception on GPUs" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">Typical out-of-memory (OOM) exception on GPUs</span></span></P> <P><SPAN>To limit the chance of this happening, we recommend following these steps to establish an optimal configuration:</SPAN><SPAN><BR /></SPAN></P> <OL> <LI style="font-weight: 400;" aria-level="1"><SPAN>Estimate your memory requirements using the table above</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><SPAN>For model precision, use </SPAN><I><SPAN>bfloat16</SPAN></I></LI> <LI style="font-weight: 400;" aria-level="1"><SPAN>For optimizer and gradient settings, use </SPAN><I><SPAN>ZeRo 3</SPAN></I></LI> </OL> <P><SPAN>Try a training run and see how it goes. If you find yourself running into out-of-memory issues, progressively add offload and set gradient accumulation steps and turn on gradient checkpointing. The training process may run more slowly, but it will be more likely to complete.</SPAN></P> <H3 id="mcetoc_1hmhf6r923u"><SPAN>Monitoring and debugging</SPAN></H3> <P><SPAN>The best way to monitor and debug DeepSpeed training runs is to look at the <I>Metrics</I> and <I>Driver logs</I> tabs in the Databricks web UI:</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Find hardware utilization metrics and system logs from the Cluster UI" style="width: 905px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6199i5B6D76B6D6311B52/image-size/large?v=v2&amp;px=999" role="button" title="brian_law_1-1707741353341.png" alt="Find hardware utilization metrics and system logs from the Cluster UI" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">Find hardware utilization metrics and system logs from the Cluster UI</span></span></P> <P><SPAN>The metrics tab has separate dropdowns for hardware settings, including GPU and CPU utilization. </SPAN><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="GPU utilization metrics in Databricks" style="width: 999px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6200iA435C3EFCB96A286/image-size/large?v=v2&amp;px=999" role="button" title="brian_law_2-1707741377461.png" alt="GPU utilization metrics in Databricks" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">GPU utilization metrics in Databricks</span></span></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Load and CPU utilization metrics in Databricks" style="width: 999px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6201iE5C716E3B647FC17/image-size/large?v=v2&amp;px=999" role="button" title="brian_law_3-1707741399820.png" alt="Load and CPU utilization metrics in Databricks" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">Load and CPU utilization metrics in Databricks</span></span></P> <P><SPAN>The driver logs will give us the same debug details that would appear in a terminal session should the excerpts displayed in the notebook prove insufficient.</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Additional details in driver logs" style="width: 999px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/6202i931DA9E4EB9F7B31/image-size/large?v=v2&amp;px=999" role="button" title="brian_law_4-1707741425620.png" alt="Additional details in driver logs" /><span class="lia-inline-image-caption" onclick="event.preventDefault();">Additional details in driver logs</span></span></P> <H2 id="mcetoc_1hmhf6r923v"><SPAN>Conclusion</SPAN></H2> <P><SPAN>The DeepSpeed Distributor on Databricks marks a pivotal advancement in efficiently training large language models by significantly reducing GPU memory requirements. It simplifies resource configurations and management, enabling scalable and optimized AI project development.&nbsp; Try it out on Databricks today using the code found <A href="https://github.com/Data-drone/Deepspeed_distributor" target="_blank" rel="noopener">here</A>.</SPAN></P> Tue, 13 Feb 2024 14:51:52 GMT brian_law 2024-02-13T14:51:52Z https://community.databricks.com/t5/technical-blog/introducing-the-deepspeed-distributor-on-databricks/ba-p/59641 <P>Are you training large deep learning models? Do you want to make the most of the compute resources available to you? Have you gone through the pain of configuring and setting up compute clusters? Well this article is just for you! With the new DeepSpeed distributor you all the above and more quicker and simpler! Best of all, it is in open source too for the benefit of the whole community. Take your GPU-enabled cluster and accelerate your model training and finetuning today!&nbsp;&nbsp;</P> Tue, 13 Feb 2024 14:51:52 GMT https://community.databricks.com/t5/technical-blog/introducing-the-deepspeed-distributor-on-databricks/ba-p/59641 brian_law 2024-02-13T14:51:52Z https://community.databricks.com/t5/technical-blog/introducing-the-deepspeed-distributor-on-databricks/bc-p/101651#M408 <P><a href="https://community.databricks.com/t5/user/viewprofilepage/user-id/1714">@brian_law</a>&nbsp;In the section above on HuggingFace cache you recommend using DBFS for the model cache. Can a shared path (such as a Workspace of Volume) be used as well?</P> Tue, 10 Dec 2024 18:29:54 GMT https://community.databricks.com/t5/technical-blog/introducing-the-deepspeed-distributor-on-databricks/bc-p/101651#M408 pemidexx 2024-12-10T18:29:54Z