FSDP is using more GPU memory than DDP

username: XYG_123
cooked: <p>Hi team,</p>
<p>I have one model, trained 2 times. One under DDP and one under FSDP. And I have enabled activation checkpointing for them both. With DDP, I can use <code>batch_size==512</code>. However, with FSDP, I can only use <code>batch_size==256</code>.</p>
<p>After some profile, I found <code>SplitWithSizesBackward</code> is suspicious. There are large chunks of <code>cudaMalloc</code> there.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/a/6a1f7d35d0db9e493c3e0c5941929fe5ebe18d00.png" data-download-href="https://discuss.pytorch.org/uploads/default/6a1f7d35d0db9e493c3e0c5941929fe5ebe18d00" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/a/6a1f7d35d0db9e493c3e0c5941929fe5ebe18d00_2_690x171.png" alt="image" data-base62-sha1="f8O4LPkQj8gjANMrsGjcVTMCu9G" width="690" height="171" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/a/6a1f7d35d0db9e493c3e0c5941929fe5ebe18d00_2_690x171.png, https://discuss.pytorch.org/uploads/default/original/3X/6/a/6a1f7d35d0db9e493c3e0c5941929fe5ebe18d00.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/a/6a1f7d35d0db9e493c3e0c5941929fe5ebe18d00.png 2x" data-dominant-color="CDB7CA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">773×192 19.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Can anyone share some insights about this function or this pattern? I would appreciate that.</p>

username: XYG_123
cooked: <p>The OOM happens during the backward.</p>
<pre><code class="lang-auto">   total_loss.backward()
  File "/usr/local/lib/python3.8/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/usr/local/lib/python3.8/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 308.00 MiB. GPU 7 has a total capacty of 39.56 GiB of which 188.81 MiB is free. Process 457865 has 39.37 GiB memory in use. Of the allocated memory 36.82 GiB is allocated by PyTorch, and 1.24 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
[2024-12-02 23:10:31,719][thrift_logger][INFO] - Thrift loggers for 0 topics are closed
&gt;=1 process exited with error return code 1, which will cause training to hang. Killing remaining process groups: [464460, 464461, 464463, 464464]
</code></pre>

username: ening
cooked: <p>I also find this problem. When I use fsdp and three gpus to fine-tune llama3 vision model, I find every gpu(fsdp) has a higher memory than others(every gpu in ddp and single gpu). I get OOM error and it happens in _flat_param.py when flattening tensors. What’s more, even I can load model when using fsdp, it has higher memory than ddp.<br>
Can you have some solution to solve this problem or it may be a bug?</p>

username: ptrblck
cooked: <p>Double post from <a href="https://discuss.pytorch.org/t/fsdp-memory-higher-than-ddp-or-single-gpu/215337">here</a>.</p>