# How to avoid fragmentation during load\_state\_dict(load(...)) when used with an old model obj?

**URL:** <https://discuss.pytorch.org/t/how-to-avoid-fragmentation-during-load-state-dict-load-when-used-with-an-old-model-obj/34848>\
**Category:** Uncategorized\
**Created:** [January 16, 2019, 11:18pm UTC](https://discuss.pytorch.org/t/how-to-avoid-fragmentation-during-load-state-dict-load-when-used-with-an-old-model-obj/34848 "2019-01-16T23:18:20Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![stas](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/stas/32/8409_2.png) [@stas](https://discuss.pytorch.org/u/stas)\
**Post date:** [January 16, 2019, 11:18pm UTC](https://discuss.pytorch.org/t/how-to-avoid-fragmentation-during-load-state-dict-load-when-used-with-an-old-model-obj/34848/1 "2019-01-16T23:18:20Z")

</div>

Currently, when loading of the saved model is used as checkpoints, this code causes GPU RAM fragmentation:

```auto
torch.save(model.state_dict(), PATH)
model.load_state_dict(torch.load(PATH))

```

1. `torch.load` loads the model on cuda allocating new RAM
2. `load_state_dict` replaces the old model with the new one, while freeing the previously allocated gpu RAM

The following memory allocation sketch shows how a hole is created in these 2 steps, which is very likely to lead to memory fragmentation.

```auto
1. [allocated mem][prev model][new model][free]
2. [allocated mem][---hole---][new model][free]

```

If the checkpoints are used multiple times and the model is ever-growing, there will be lots of free, but non-reusable sections in the memory. leading to a lot of wasted GPU RAM.

Is there a way to add an alternative re-loading function which first frees up the previous dict and then replaces it with the new one? Looking at the code it appears it’d be very complicated to implement. Unless perhaps it’d be enough to remove it from CUDA.

Another approach would be to have a torch.load() split into two parts, so that (1) the first part loading into general RAM, (2) then load\_state\_dict assigning it and (3) switching it to the right cuda device (I’m assuming cuda in this example).

---

<div class="post-metadata">

**Author:** ![stas](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/stas/32/8409_2.png) [@stas](https://discuss.pytorch.org/u/stas)\
**Post date:** [January 23, 2019, 5:10am UTC](https://discuss.pytorch.org/t/how-to-avoid-fragmentation-during-load-state-dict-load-when-used-with-an-old-model-obj/34848/2 "2019-01-23T05:10:26Z")

</div>

So the understanding that came from working on [GPU RAM fragmentation diagnostics](https://discuss.pytorch.org/t/gpu-ram-fragmentation-diagnostics/34073) is that this won’t cause an actual fragmentation, since an average model will be 100MB+ and a whole bunch of whole gpu memory pages will be freed and re-used later through remapping of free pages.

So the only issue here is if there is not enough memory left to load the new model, without first unloading the old one, in which case the card will not be able to do much work anyway, other than perhaps a very simple inference. So I suppose this is a very low priority for devs to spend their time on.

All is good then.

Thanks again to @colesbury for his very insightful [answer](https://discuss.pytorch.org/t/gpu-ram-fragmentation-diagnostics/34073/6).