CUDA OOM error when loading sharded checkpoint

barbet: <p>Hi all,</p>
<p>We fine-tuned Stability’s StableLM-7b using Huggingface’s <code>Trainer</code> API (with FSDP) and then saved the resulting checkpoints in the sharded format that is typical for large language models. Quite surprisingly, however, attempting to load the model for inference leads to a strange error when loading one of the checkpoints (<code>Unable to load weights from pytorch checkpoint file</code>)</p>
<p>We took some further investigative steps by making a simple <code>torch.load</code> call on the problem shard, and got a typical CUDA OOM error. The exceedingly strange thing about this OOM error is that we are working with a node with 8xA100s (80GB), and the given state dict is only 171kB (comprising only 7 layers of the model). So, you can imagine seeing the following error was quite a shock:</p>
<p><code>torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 29.31 GiB (GPU 0; 79.19 GiB total capacity; 55.76 GiB already allocated; 22.48 GiB free; 55.76 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></p>
<p>After looking into this further, I discovered a few threads discussing this issue, like <a href="https://discuss.pytorch.org/t/cuda-error-out-of-memory-when-load-models/38011">this one</a>, and attempted some of the fixes, namely loading the state dict on CPU first. After doing so, I received the following error:<br>
<code>RuntimeError: Trying to resize storage that is not resizable</code></p>
<p>So it seems that approach is out of the question. As I previously said, the strange thing here is that the first two shards load without issue, while the third and fourth cannot be loaded. Additionally, nothing seems particularly out of place in the shard-layer mapping JSON. I am stumped here. Perhaps it’s an error in Huggingface <code>Trainer</code>’s save function, but I figured I would start at the PyTorch source first.</p>

barbet: <p>It is also worth adding that the given device in the error (cuda:0) is completely empty when <code>torch.load</code> is called</p>

kwen2501: <p><a class="mention" href="/u/irisz">@irisz</a> I was wondering if you could you take a look. Thanks!</p>

irisz: <p>Could you share the code snippet for saving and loading?</p>
<p>If you are using HF trainer, the issue is they are running torch.save, which is not really suitable for larger model/multi-node scenarios for PyTorch FSDP.<br>
For really large models, should be using sharded_state_dict and the distributed FileSystemWriter which is optimized for handling multi-node, large model scenarios.<br>
Example distributed checkpoint save code here: <a href="https://github.com/pytorch/pytorch/blob/e71ab214226af1f9dbded944e939c6447e0e8f09/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py#L59" class="inline-onebox" rel="noopener nofollow ugc">pytorch/fsdp_checkpoint_example.py at e71ab214226af1f9dbded944e939c6447e0e8f09 · pytorch/pytorch · GitHub</a></p>

Kerem_Can_Turgutlu: <p>I have a related question, similarly I am training a 7B model using accelerate and FSDP with <code>StateDictType.SHARDED_STATE_DICT</code>.</p>
<p>What is the recommended way to load sharded <code>__{i}_{i}.distcp</code> optimizer and parameter state dict files on a cpu or a single gpu without needing to initialize torch distributed?</p>
<p>For training and resuming training distributed checkpoint files make sense but for inference there might not be a need to use FSDP and can we convert these files to something that can be read simply with  <code>torch.load</code>?</p>
<p>Following approach worked for me when saving full state dicts for both optimizer and the model.</p>
<pre><code class="lang-auto"># Save FSDP full optimizer state dict. 
full_optim_save_path = os.path.join(output_dir, "optimizer.bin")
optim_state = model.full_optim_state_dict(model, optimizer)
if accelerator.is_local_main_process:
    torch.save(optim_state, full_optim_save_path)
# Save FSDP full model state dict. 
cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with model.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
    model_state = model.state_dict()
    full_model_save_path = os.path.join(output_dir, "pytorch_model.bin")
    if accelerator.is_local_main_process:
        torch.save(model_state, full_model_save_path)
</code></pre>