PyTorch: allocation on wrong device when loading checkpoint

[Marks101]

Problem description

When storing the state dict of a transformer_engine model on a device with index A and afterwards load it on another device with index B, some memory still gets allocated on device A. The reason for this seems the be the pickle.loads in https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/module/base.py#L411

This issue has a large practical impact in distributed training: if device 0 on a host writes the model checkpoint and when resuming the training all other devices read it, all 8 devices on the host suddenly allocate on device 0. In our case we are not able to resume trainings due to OOMs.

Minimal example

import argparse
import torch
import transformer_engine as te

parser = argparse.ArgumentParser()
parser.add_argument("--device", default=0, type=int)
parser.add_argument("--load", action="store_true")
args = parser.parse_args()

torch.cuda.set_device(args.device)

hidden_size = 1024
model = te.pytorch.Linear(hidden_size, hidden_size).cuda()

if args.load:
 state = torch.load("model.pt")
 model.load_state_dict(state)
 print("After load, max_memory_reserved on device 0:", torch.cuda.max_memory_reserved(0))

for _ in range(10):
 with te.pytorch.fp8_autocast(enabled=True):
 output = model(torch.randn((32, hidden_size)).cuda())

torch.save(model.state_dict(), "model.pt")

The issue is triggered this way:

>>> python test_fp8_mem.py --device 0
>>> python test_fp8_mem.py --device 1 --load
After load, max_memory_reserved on device 0: 23068672

Although, torch is set to use device 1 in the second call, memory on device 0 gets allocated as well. The overall CUDA context cannot be freed anymore, thus a significant part of the device is blocked.

Workaround

Restricting CUDA_VISIBLE_DEVICES to a single GPU per process works for us as a quick and dirty workaround

[mikolajblaz]

This issue comes up also when using the new distributed checkpoint format (extra state is saved only from 1 of the devices in every TP group)

[ksivaman]

@Marks101 There are 2 parts to this problem:

1. The model weights: These make up a large chunk of the allocated memory (~21 out of 23 MB in the toy example above) and can simply be allocated directly on the correct device by passing the map_location="cuda" argument to torch.load in your script, bypassing the need for any allocation on device 0.
2. The fp8 metadata: These are custom fp8 related tensors that need to be stored and make up only a small portion of the memory (~2 out of 23 MB in the toy example above). These are saved via PyTorch's custom get_extra_state method in TE which unfortunately does not receive the map_location argument. So these extra tensors would still be redundantly allocated on device 0 before being moved to the correct device during loading of the state dict. Often in practical scenarios this memory is negligible compared to the memory of the actual parameters, however, I have created Store FP8 checkpointing data in CPU #351 as a workaround.

[Marks101]

You solution looks good. That should do the job. Thank you very much!

[ksivaman]

Fixed in #351

[Marks101]

Thank you so much for processing this so quickly!