GPU memory spikes with FSDP and torch.compile

ipoletaev: <p>Hi torch users,</p>
<p>Any ideas on whether it is expected to see an FSDP model (after being compiled) to re-compile things after certain number of training steps? Maybe someone faced similar problems? The code follows usual training procedure with FSDP:</p>
<pre><code class="lang-auto">model = ...
model = fsdp.FullyShardedDataParallel(...) # Wraps only decoder blocks in FSDP units
model = torch.compile(model)
optimizer = ...
...
# training loop
loss = model(X, Y)
loss.backward(loss)
optimizer.step()
optimizer.zero_grad(set_to_none=True)
</code></pre>
<p>Here is an example of a GPT2 test runs on 2 nodes , 8 GPUs each in FSDP mode:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/7/d75e1e5211713968d41236aab88df9eddb08dec6.png" data-download-href="https://discuss.pytorch.org/uploads/default/d75e1e5211713968d41236aab88df9eddb08dec6" title="Screenshot 2023-05-03 at 4.00.39 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d75e1e5211713968d41236aab88df9eddb08dec6_2_601x500.png" alt="Screenshot 2023-05-03 at 4.00.39 PM" data-base62-sha1="uJeiQtsa7DGr5EAHlD86TXoTSEm" width="601" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d75e1e5211713968d41236aab88df9eddb08dec6_2_601x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d75e1e5211713968d41236aab88df9eddb08dec6_2_901x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d75e1e5211713968d41236aab88df9eddb08dec6_2_1202x1000.png 2x" data-dominant-color="FEFDFD"><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">Screenshot 2023-05-03 at 4.00.39 PM</span><span class="informations">2106×1750 176 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>You can see above that all of a sudden around ~180 seconds (the code is in the training loop, no evaluation or any other changes to the model) VRAM usage spikes.</p>
<p>Compiled model is - no doubt - faster than its default counterpart. But the latter at least  has no such strange in-flight-vram-changing behaviors. Why is it bad? Because under high load where GPU memory usage is around 100%, such sudden spikes could result in OOMs.</p>
<p>Thank you for any insights!</p>

eqy: <p>Do you have a sense for whether ~180 seconds as indicated here would still be during a warmup step in your workload? In that case the behavior could be expected as a full forward pass is needed at a minimum to observe peak memory usage for tensors alone, and possibly one or more full forward+backward passes for <code>torch.compile</code> to trace the graph and JIT compile kernels. These JIT compiled kernels would then be expected to occupy additional device memory after they are generated.</p>

ipoletaev: <p>Thanks for the quick reply.</p>
<p>Not really. On different model sizes this behavior still manifests itself but at a different point in training, however - usually close to the beginning. E.g. at a training step 42 <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20">  If I do a dummy forward pass on random input before entering the loop - the issue does not go away.</p>
<p>I wonder if there is a way to enable some logging to see what happens? OR maybe someone know a magical place to put the <code>print()</code> to in order to see and say “yes, it is JIT compiler doing things here” or “no, it is something else”. Related question: if I replace activation function implementation (that is - a simple python function) with its autograd’s version - the spike is smaller. Does that suggest a better idea of what it could be?</p>

marksaroufim: <p>Are your input sequences always the same shape? If the shapes are dynamic then recompilations are expected</p>
<p>As far as logging goes <code>TORCH_COMPILE_DEBUG=1 python train.py</code> is the single most useful flag available</p>

ipoletaev: <p>Yes! Sequence length was the first thing I looked into. Thanks for the flag suggestion, I’ll try it.</p>

ipoletaev: <p><a class="mention" href="/u/marksaroufim">@marksaroufim</a> so around the spike time the only things this flag helps to log is:</p>
<pre><code class="lang-auto">[2023-05-04 17:19:26,426] torch._dynamo.eval_frame: [DEBUG] skipping &lt;lambda&gt; .../.venv/lib/python3.8/site-packages/torch/_dynamo/utils.py
[2023-05-04 17:19:26,426] torch._dynamo.eval_frame: [DEBUG] skipping _remove_id .../.venv/lib/python3.8/site-packages/torch/_dynamo/utils.py
</code></pre>
<p>Nothing else…</p>