Loading weights without `map_location` with `strict=True` doesnt throw error

distributed
1

Suppose I’m running code with torch.nn.distributed on two GPUs. If the pre-trained weights were saved using cuda:0, then loading the same weights on different GPUs fails and doesn’t even throw error despite having strict = True

checkpoint = torch.load(args.pretrain)
checkpoint['model_state_dict']['ln_final.bias'].device  # device(type='cuda', index=0)

msg = model.load_state_dict(checkpoint['model_state_dict'], strict=True) 
# always return <All keys matched successfully>

If the above process is run via torch.nn.distributed on two GPUs, cuda:0 and cuda:1, the models weights will still show <All keys matched successfully> for both GPUs but no weights will be loaded for cuda:1 unless checkpoint = torch.load(args.pretrain, map_location='cpu')

Took some time time to figure this out. Is this an expected behavior?

2

strict=True only checks that the key names line up, it doesn’t look at devices at all, so <All keys matched successfully> is expected here. load_state_dict copies the values into your existing parameters with copy_, which works across devices, so I’d double-check how you’re seeing that cuda:1 got no weights. What does bite you without map_location is that every rank deserializes the whole checkpoint onto cuda:0 first, which is why the DDP tutorial loads with map_location mapped to the local rank. I work on aicost.tools and mostly stare at the GPU bill side of things, but map_location='cpu' (or f'cuda:{rank}') is the safe default either way.

3

I checked the values, values got updated on cuda:0 but didnt on cuda:1

Thank you, your explanation does make sense. I guess, I just assumed, load_state_dict to throw error on device incompatibility rather displaying <All keys matched successfully> without loading anything.

4

I also just encountered this issue. Trained models on GPU zero and when calculating additional metrics using the trained models only the runs on cuda:0 would perform as expected while the runs on cuda:1 would make random predictions. I only get expected results when explicitly mapping the weights to cuda:1