State_dict() and CUDA OOM

cerisara: <p>Hi,</p>
<p>I’m training a model on 4xGPU A100 80GB and it trains fine, but is really occupying nearly all the VRAM. Now, when I want to save the model, I cannot directly use torch.save(model) because of the ProcessGroup, so I use torch.save(model.state_dict()); but this creates an OOM crash, because it allocates new VRAM for the state_dict: this seems rather useless!<br>
Do you know of any way to save the state_dict without allocating any new VRAM? Thanks a lot!</p>
<p>Christophe</p>

ptrblck: <p>What is your host RAM used for? Maybe you could delete unneeded data before moving the <code>state_dict</code> to the host.</p>

cerisara: <p>Thank you! It’s the classical weights + gradients + activations (w/ gradient checkpointing) + optimizer state (but I use SGD), and I want to checkpoint the model during training. But there’s not enough VRAM to hold 2 times the weights in bf16 (I’m using sharding for training) on a single GPU, so I’m not sure it’s possible to do that… That’s why I’d like to directly copy the weights from the GPU to the disk, or through the CPU, without duplicating the weights on the GPU… Thanks a lot for your ideas, which are always super-useful <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"></p>
<p>OK but I see your point: move the weights onto the CPU RAM, and then call state_dict() and save it… But I guess calling state_dict() will also duplicate the weights, and I’m afraid the CPU RAM cannot host 2 times the weights, but I’ll have to do some maths to know… My point was that it seems pretty useless to duplicate the weights just to save them, I don’t really understand why it works like that?</p>

ptrblck: <p>I don’t fully understand where the duplication would come from. Calling <code>model.state_dict()</code> will return references to the parameters and buffers, so no copy will be created unless you explicitly call <code>copy.deepcopy</code> on it. Moving the model from the GPU to the CPU will also just move the data without creating duplicates.</p>

cerisara: <p>Thanks for the explanation; but if state_dict() does not allocate any memory, I don’t understand why I have an OOM error when I call it, and I don’t otherwise. Here is the error trace to show where it allocates memory, may it be useful? Thanks a lot!</p>
<p>P.S. On 2nd thoughts, it might be because when asking to save the model, it’s unsharding it, so trying to copy back all the other GPU shards into the main one, and it OOM because it still has other gradients, etc. in VRAM… So this would mean that I should only save the model after releasing all gradients, etc… !?!</p>
<pre><code class="lang-auto">│   209 │   │   if rank==0:                                                    │
│   210 │   │   │   dt = time.time() - t0                                      │
│   211 │   │   │   if dt&gt;1800: # 30 minutes                                   │
│ ❱ 212 │   │   │   │   states = model.state_dict()                            │
│   213 │   │   │   │   torch.save(states, wd2+"fsdpmod."+str(it0)+".pt")      │
│   214 │   │   │   │   t0 = time.time()                                       │
│   215 │   │   │   │   it0 += 1                                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1818 in state_dict        │
│                                                                              │
│   1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)      │
│   1816 │   │   for name, module in self._modules.items():                    │
│   1817 │   │   │   if module is not None:                                    │
│ ❱ 1818 │   │   │   │   module.state_dict(destination=destination, prefix=pre │
│   1819 │   │   for hook in self._state_dict_hooks.values():                  │
│   1820 │   │   │   hook_result = hook(self, destination, prefix, local_metad │
│   1821 │   │   │   if hook_result is not None:                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1818 in state_dict        │
│                                                                              │
│   1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)      │
│   1816 │   │   for name, module in self._modules.items():                    │
│   1817 │   │   │   if module is not None:                                    │
│ ❱ 1818 │   │   │   │   module.state_dict(destination=destination, prefix=pre │
│   1819 │   │   for hook in self._state_dict_hooks.values():                  │
│   1820 │   │   │   hook_result = hook(self, destination, prefix, local_metad │
│   1821 │   │   │   if hook_result is not None:                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1818 in state_dict        │
│                                                                              │
│   1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)      │
│   1816 │   │   for name, module in self._modules.items():                    │
│   1817 │   │   │   if module is not None:                                    │
│ ❱ 1818 │   │   │   │   module.state_dict(destination=destination, prefix=pre │
│   1819 │   │   for hook in self._state_dict_hooks.values():                  │
│   1820 │   │   │   hook_result = hook(self, destination, prefix, local_metad │
│   1821 │   │   │   if hook_result is not None:                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1818 in state_dict        │
│                                                                              │
│   1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)      │
│   1816 │   │   for name, module in self._modules.items():                    │
│   1817 │   │   │   if module is not None:                                    │
│ ❱ 1818 │   │   │   │   module.state_dict(destination=destination, prefix=pre │
│   1819 │   │   for hook in self._state_dict_hooks.values():                  │
│   1820 │   │   │   hook_result = hook(self, destination, prefix, local_metad │
│   1821 │   │   │   if hook_result is not None:                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1818 in state_dict        │
│                                                                              │
│   1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)      │
│   1816 │   │   for name, module in self._modules.items():                    │
│   1817 │   │   │   if module is not None:                                    │
│ ❱ 1818 │   │   │   │   module.state_dict(destination=destination, prefix=pre │
│   1819 │   │   for hook in self._state_dict_hooks.values():                  │
│   1820 │   │   │   hook_result = hook(self, destination, prefix, local_metad │
│   1821 │   │   │   if hook_result is not None:                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1818 in state_dict        │
│                                                                              │
│   1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)      │
│   1816 │   │   for name, module in self._modules.items():                    │
│   1817 │   │   │   if module is not None:                                    │
│ ❱ 1818 │   │   │   │   module.state_dict(destination=destination, prefix=pre │
│   1819 │   │   for hook in self._state_dict_hooks.values():                  │
│   1820 │   │   │   hook_result = hook(self, destination, prefix, local_metad │
│   1821 │   │   │   if hook_result is not None:                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1815 in state_dict        │
│                                                                              │
│   1812 │   │   if hasattr(destination, "_metadata"):                         │
│   1813 │   │   │   destination._metadata[prefix[:-1]] = local_metadata       │
│   1814 │   │                                                                 │
│ ❱ 1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)      │
│   1816 │   │   for name, module in self._modules.items():                    │
│   1817 │   │   │   if module is not None:                                    │
│   1818 │   │   │   │   module.state_dict(destination=destination, prefix=pre │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/nn/modules/module.py:1722 in                   │
│ _save_to_state_dict                                                          │
│                                                                              │
│   1719 │   │   │   │   module                                                │
│   1720 │   │   """                                                           │
│   1721 │   │   for hook in self._state_dict_pre_hooks.values():              │
│ ❱ 1722 │   │   │   hook(self, prefix, keep_vars)                             │
│   1723 │   │                                                                 │
│   1724 │   │   for name, param in self._parameters.items():                  │
│   1725 │   │   │   if param is not None:                                     │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/utils/_contextlib.py:115 in decorate_context   │
│                                                                              │
│   112 │   @functools.wraps(func)                                             │
│   113 │   def decorate_context(*args, **kwargs):                             │
│   114 │   │   with ctx_factory():                                            │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                               │
│   116 │                                                                      │
│   117 │   return decorate_context                                            │
│   118                                                                        │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/_state_dict_utils.py:669 in   │
│ _pre_state_dict_hook                                                         │
│                                                                              │
│   666 │   │   StateDictType.LOCAL_STATE_DICT: _local_pre_state_dict_hook,    │
│   667 │   │   StateDictType.SHARDED_STATE_DICT: _sharded_pre_state_dict_hook │
│   668 │   }                                                                  │
│ ❱ 669 │   _pre_state_dict_hook_fn[fsdp_state._state_dict_type](              │
│   670 │   │   fsdp_state,                                                    │
│   671 │   │   module,                                                        │
│   672 │   │   *args,                                                         │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/_state_dict_utils.py:271 in   │
│ _full_pre_state_dict_hook                                                    │
│                                                                              │
│   268 │   in ``nn.Module``.                                                  │
│   269 │   """                                                                │
│   270 │   _common_pre_state_dict_hook(module, fsdp_state)                    │
│ ❱ 271 │   _common_unshard_pre_state_dict_hook(                               │
│   272 │   │   module,                                                        │
│   273 │   │   fsdp_state,                                                    │
│   274 │   │   offload_to_cpu=fsdp_state._state_dict_config.offload_to_cpu,   │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/_state_dict_utils.py:143 in   │
│ _common_unshard_pre_state_dict_hook                                          │
│                                                                              │
│   140 │   Performs the pre-state_dict tasks shared by all state_dict types t │
│   141 │   ``_unshard_fsdp_state_params()``. FULL_STATE_DICT and SHARDED_STAT │
│   142 │   """                                                                │
│ ❱ 143 │   _enter_unshard_params_ctx(                                         │
│   144 │   │   module,                                                        │
│   145 │   │   fsdp_state,                                                    │
│   146 │   │   writeback=False,                                               │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/_state_dict_utils.py:109 in   │
│ _enter_unshard_params_ctx                                                    │
│                                                                              │
│   106 │   │   offload_to_cpu=offload_to_cpu,                                 │
│   107 │   │   with_grads=with_grads,                                         │
│   108 │   )                                                                  │
│ ❱ 109 │   fsdp_state._unshard_params_ctx[module].__enter__()                 │
│   110                                                                        │
│   111                                                                        │
│   112 @no_type_check                                                         │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/contextlib.py:135 in __enter__                                     │
│                                                                              │
│   132 │   │   # they are only needed for recreation, which is not possible a │
│   133 │   │   del self.args, self.kwds, self.func                            │
│   134 │   │   try:                                                           │
│ ❱ 135 │   │   │   return next(self.gen)                                      │
│   136 │   │   except StopIteration:                                          │
│   137 │   │   │   raise RuntimeError("generator didn't yield") from None     │
│   138                                                                        │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/_unshard_param_utils.py:198   │
│ in _unshard_fsdp_state_params                                                │
│                                                                              │
│   195 │   # No need to call `wait_stream()` since we unshard in the computat │
│   196 │   # stream directly                                                  │
│   197 │   computation_stream = torch.cuda.current_stream()                   │
│ ❱ 198 │   _unshard(state, handles, computation_stream, computation_stream)   │
│   199 │   if with_grads:                                                     │
│   200 │   │   _unshard_grads(handles)                                        │
│   201                                                                        │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py:329 in      │
│ _unshard                                                                     │
│                                                                              │
│    326 │   │   │   event.synchronize()                                       │
│    327 │   with torch.cuda.stream(unshard_stream):                           │
│    328 │   │   for handle in handles:                                        │
│ ❱  329 │   │   │   handle.unshard()                                          │
│    330 │   │   │   handle.post_unshard()                                     │
│    331                                                                       │
│    332                                                                       │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/flat_param.py:918 in unshard  │
│                                                                              │
│    915 │   │   │   )                                                         │
│    916 │   │   │   self._use_unsharded_flat_param(unsharded_flat_param)      │
│    917 │   │   │   return                                                    │
│ ❱  918 │   │   unsharded_flat_param = self._alloc_padded_unsharded_flat_para │
│    919 │   │   padded_unsharded_flat_param = self._all_gather_flat_param(uns │
│    920 │   │   self._use_unsharded_flat_param(padded_unsharded_flat_param)   │
│    921                                                                       │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/flat_param.py:944 in          │
│ _alloc_padded_unsharded_flat_param                                           │
│                                                                              │
│    941 │   │   flat_param = self.flat_param                                  │
│    942 │   │   unsharded_flat_param = self._get_padded_unsharded_flat_param( │
│    943 │   │   self._check_storage_freed(unsharded_flat_param)               │
│ ❱  944 │   │   _alloc_storage(unsharded_flat_param, flat_param._padded_unsha │
│    945 │   │   return unsharded_flat_param                                   │
│    946 │                                                                     │
│    947 │   def _get_padded_unsharded_flat_param(self) -&gt; torch.Tensor:       │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/utils/_contextlib.py:115 in decorate_context   │
│                                                                              │
│   112 │   @functools.wraps(func)                                             │
│   113 │   def decorate_context(*args, **kwargs):                             │
│   114 │   │   with ctx_factory():                                            │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                               │
│   116 │                                                                      │
│   117 │   return decorate_context                                            │
│   118                                                                        │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/distributed/fsdp/_utils.py:79 in               │
│ _alloc_storage                                                               │
│                                                                              │
│    76 │   │   │   tensor_storage_size == 0,                                  │
│    77 │   │   │   f"Tensor storage should have been resized to be 0 but got  │
│    78 │   │   )                                                              │
│ ❱  79 │   │   tensor._typed_storage()._resize_(size.numel())                 │
│    80 │   return not already_allocated                                       │
│    81                                                                        │
│    82                                                                        │
│                                                                              │
│ /gpfslocalsup/pub/anaconda-py3/2023.03/envs/pytorch-gpu-2.0.0+py3.10.9/lib/p │
│ ython3.10/site-packages/torch/storage.py:764 in _resize_                     │
│                                                                              │
│    761 │                                                                     │
│    762 │   # For internal use only, to avoid deprecation warning             │
│    763 │   def _resize_(self, size):                                         │
│ ❱  764 │   │   self._untyped_storage.resize_(size * self._element_size())    │
│    765 │                                                                     │
│    766 │   @classmethod                                                      │
│    767 │   def _free_weak_ref(cls, *args, **kwargs):                         │
╰──────────────────────────────────────────────────────────────────────────────╯
OutOfMemoryError: CUDA out of memory. Tried to allocate 236.00 MiB (GPU 0; 79.15
GiB total capacity; 76.94 GiB already allocated; 121.19 MiB free; 77.82 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>

ptrblck: <aside class="quote no-group" data-username="cerisara" data-post="5" data-topic="193214">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cerisara/48/428_2.png" class="avatar"> cerisara:</div>
<blockquote>
<p>but if state_dict() does not allocate any memory, I don’t understand why I have an OOM error when I call it, and I don’t otherwise.</p>
</blockquote>
</aside>
<p>Yes, it doesn’t as also seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">model = models.resnet152().cuda()
print("{:.2f}MB".format(torch.cuda.memory_allocated()/1024**2))
# 230.27MB

sd = model.state_dict()
print("{:.2f}MB".format(torch.cuda.memory_allocated()/1024**2))
# 230.28MB

clone = copy.deepcopy(sd)
print("{:.2f}MB".format(torch.cuda.memory_allocated()/1024**2))
# 462.24MB
</code></pre>
<aside class="quote no-group" data-username="cerisara" data-post="5" data-topic="193214">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cerisara/48/428_2.png" class="avatar"> cerisara:</div>
<blockquote>
<p>it might be because when asking to save the model, it’s unsharding it, so trying to copy back all the other GPU shards into the main one</p>
</blockquote>
</aside>
<p>That might be the case as I didn’t know you are using <code>FSDP</code>.<br>
I believe the recommendation is to turn on <code>cpu_offload</code> in <a href="https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type"><code>torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type</code></a>, which should avoid the CUDA OOM.</p>

cerisara: <p>Thanks a lot for the link, that’s exactly what I need! Thanks! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>