Incompatibility in model checkpoints between multi-GPU/single-GPU

[zphang]

Summary: If you train a model with multi-GPU, then try to load it with single-GPU, the weights will not be successfully loaded. The run will print warnings but will not fail.

Background:

• When a PyTorch module is wrapped in DataParallel, the original module will be stored as multi_gpu_model.module.
• This means that calling state_dict on multi_gpu_model will add a "module." prefix to every key in the state_dict.
• In other words, state_dicts are not compatible between single-gpu and multi-gpu models.

What this affects:

• Two-phase training runs that are run from two separate commands (via load_target_train_checkpoint) will incorrectly load the model weights, reverting to the pretrained weights if one phase uses multi-GPU and the other uses single-GPU
• E.g. If intermediate task training is on multi-GPU and target task training is on single-GPU, the target task run will fail to load the intermediate-trained weights, and revert to pretrained weights. The logs will be filled with "parameter missed" warnings, but will not fail

Proposed solution:

• Don't call state_dict on a DataParallel model. Write a function that determines whether a model is DaraParallel, and call model.state_dict or model.module.state_dict correspondingly

def get_state_dict_for_saving(model: nn.Module) -> nn.Module:
    if isinstance(model, nn.DataParallel):
        return model.module.state_dict()
    else:
        return model.state_dict()

[sleepinyourhat]

Eek—does anyone have bandwidth to put together a fix?

[pruksmhc]

I can take a look at this this weekend

[pruksmhc]

Update: It's actually not quite as simple as just putting the above function to modify saving, since modifying saving will not work if we want to restart a job from multi-GPU (because the saved model will now always have model.state_dict() while the model in multi-GPU expects module.model.state_dict()

[zphang]

Could you instead call model.module.load_state_dict?

[zphang]

This is an automatically generated comment.

As we update jiant to v2.x, jiant v1.x has been migrated to https://github.com/nyu-mll/jiant-v1-legacy. As such, we are closing all issues relating to jiant v1.x in this repository.

If this issue is still affecting you in jiant v1.x, please follow up at nyu-mll/jiant-v1-legacy#1087.

If this issue is still affecting you in jiant v2.x, reopen this issue or create a new one.