FSDP failed to save model checkpoints

username: weiqis
cooked: <p>Hi, I’m training a large LM on 8 A100-80GB GPUs using FSDP in HuggingFace’s Trainer. I specified the FSDP parameters as following:</p>
<pre><code class="lang-auto">fsdp: full_shard auto_wrap
fsdp_config:
  fsdp_transformer_layer_cls_to_wrap:
  - LlamaDecoderLayer
</code></pre>
<p>But when saving the model checkpoint, FSDP gives the following warning due to CUDA out of memory:</p>
<pre><code class="lang-auto">/home/ubuntu/miniconda3/envs/finetune-clm/lib/python3.10/site-
packages/torch/distributed/fsdp/_state_dict_utils.py:312: UserWarning: Failed to clone() tensor with
name _fsdp_wrapped_module.model.layers.59.mlp.gate_proj.weight on rank 6. This may mean that 
this state_dict entry could point to invalid memory regions after returning from state_dict() call if this 
parameter is managed by FSDP. Please check clone implementation of 
_fsdp_wrapped_module.model.layers.59.mlp.gate_proj.weight. Error: CUDA out of memory. Tried to 
allocate 228.00 MiB (GPU 6; 79.20 GiB total capacity; 75.32 GiB already allocated; 75.25 MiB free; 
77.23 GiB reserved in total by PyTorch) If reserved memory is &gt;&gt; allocated memory try setting 
max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and 
PYTORCH_CUDA_ALLOC_CONF
</code></pre>
<p>However, there are plenty of memory spaces left during training:</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| NVIDIA-SMI 520.61.05    Driver Version: 520.61.05    CUDA Version: 11.8     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA A100-SXM...  On   | 00000000:05:00.0 Off |                    0 |
| N/A   50C    P0   250W / 400W |  45141MiB / 81920MiB |     99%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA A100-SXM...  On   | 00000000:06:00.0 Off |                    0 |
| N/A   51C    P0   220W / 400W |  47039MiB / 81920MiB |     99%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA A100-SXM...  On   | 00000000:07:00.0 Off |                    0 |
| N/A   48C    P0   328W / 400W |  52993MiB / 81920MiB |     99%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA A100-SXM...  On   | 00000000:08:00.0 Off |                    0 |
| N/A   52C    P0   413W / 400W |  54209MiB / 81920MiB |    100%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   4  NVIDIA A100-SXM...  On   | 00000000:09:00.0 Off |                    0 |
| N/A   48C    P0    94W / 400W |  47175MiB / 81920MiB |     99%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   5  NVIDIA A100-SXM...  On   | 00000000:0A:00.0 Off |                    0 |
| N/A   43C    P0   379W / 400W |  53687MiB / 81920MiB |    100%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   6  NVIDIA A100-SXM...  On   | 00000000:0B:00.0 Off |                    0 |
| N/A   43C    P0   158W / 400W |  49709MiB / 81920MiB |    100%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   7  NVIDIA A100-SXM...  On   | 00000000:0C:00.0 Off |                    0 |
| N/A   59C    P0   421W / 400W |  54341MiB / 81920MiB |    100%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
</code></pre>
<p>It is just a warning and the training process can continue. But the saved model checkpoints are in bad shape and cannot be loaded. That means I will not be able to resume from an intermediate checkpoints.</p>
<p>When I tested the training job on a smaller GPU using a smaller model, FSDP can save model checkpoints without any problem, even when the GPU memory was tighter (less than 1GB free memory during training). But on the 80GB A100 GPUs, it complains CUDA OOM although there are almost 30GB free memory left during training. Are there anything I can do to get rid of this problem? Or did I miss anything?</p>
<p>My training env:</p>
<ul>
<li>torch: 2.0.0+cu118</li>
<li>cuda: 11.8</li>
<li>transformers: 4.48.1</li>
</ul>
<p>Thanks in advance!</p>

username: eqy
cooked: <p>Could you check if the CPU offloading workaround as described <a href="https://github.com/pytorch/pytorch/issues/98823" rel="noopener nofollow ugc">here</a> works for you?</p>

username: weiqis
cooked: <p>Awesome! Thanks <a class="mention" href="/u/eqy">@eqy</a> ! The workaround trick mentioned in <a href="https://github.com/pytorch/pytorch/issues/98823" rel="noopener nofollow ugc">the issue</a> indeed solved my problem.</p>
<p>To anyone who is also facing this problem: the cause of this problem is that when HuggingFace Trainer calls <a href="https://github.com/huggingface/transformers/blob/04ab5605fbb4ef207b10bf2772d88c53fc242e83/src/transformers/trainer.py#L2798" rel="noopener nofollow ugc"><code>model.state_dict()</code></a> during the saving step with a FSDP model, it triggers a  <code>_full_post_state_dict_hook</code> that will be executed after the call. In this post hook it did a clone of the state dict, which caused the CUDA OOM error.</p>
<p>To solve this problem, the most elegant way is probability utilizing the <code>FSDP.set_state_dict_type</code> with <code>offload_to_cpu=True</code> (sorry I cannot plug in more links as a new user), or using the <code>FSDP.state_dict_type</code> context manager to set cpu offload when calling <code>model.state_dict()</code>. However, since I’m utilizing FSDP through HuggingFace Trainer and the FSDP API is pretty limited, a quick and dirty workaround (as suggested in the github issue) is to just change the line</p>
<pre><code class="lang-auto">state_dict[fqn] = state_dict[fqn].clone().detach()
</code></pre>
<p>to</p>
<pre><code class="lang-auto">state_dict[fqn] = state_dict[fqn].cpu().clone().detach()
</code></pre>
<p>in <code>_full_post_state_dict_hook</code>.</p>
<p>It is located in line 309 of <code>torch/distributed/fsdp/_state_dict_utils.py</code><br>
in <code>torch=2.0.0</code>, or line 2221 of <code>torch/distributed/fsdp/fully_sharded_data_parallel.py</code><br>
in <code>torch=1.13.0</code>.</p>