Ignore parameters from a checkpoint

[moinnadeem]

Motivation

Some GLUE tasks require mid-training: for tasks with small, noisy datasets:

1. First, pre-train BERT on some generic language modeling corpus. (ie. Wikipedia)
2. Then, fine-tune BERT on a task-specific corpus that helps logical inference. For instance, fine-tune BERT on a 3-way classification task on a large, labeled dataset.
3. Finally, fine-tune BERT on your small, relevant task-specific corpus. For instance, a 2-way classification task.

For #3, we load the checkpoint obtained after completing #2.

This implies that BERT was trained with a task-specific classification head (in code, there is now a module.model.classifier.weight and module.model.classifier.bias) that exists that is irrelevant to us. This is because:
a) The classification weights are from a task that doesn't matter to us
b) The classification weights are a 3-way classification task, and we now have a 2-way classification task.

Because of B), the matrix sizes will not match up, and training will crash. In code, the checkpoint from B is something like batch_size x 3, and our model has a matrix of shape batch_size x 2.

This motivates ignoring certain parameters from a checkpoint.

Solutions

I see three proposed solutions here:

1. As a part of our checkpointing code, we can decide to ignore certain user-specfied parameters when saving the checkpoint.
2. As a part of our checkpoint loading code, we can ignore certain user-specfied parameters when loading the checkpoint.
3. As a part of checkpoint loading, for any keys with the same name, we can ignore any checkpoint parameters that have different shapes than the model weights, and print a warning.

Before continuing, please decide which one you like the best so I don't bias you.

To unblock myself previously, I have implemented #2: it will delete user-specified weights from the dictionary. This consists of a parameter in checkpointing called ignore_model_keys that can be specified with a list of the weights to delete from the state_dict. See the code:

• hparams API
• checkpoint loading implementation
• sample YAML file

We should decide:

1. Which solution do we prefer?
2. What does the API look like for our solution? If it is # 2, are we happy with the API and implementation that I have come up with?

Once we have decided, I'm happy to code up the pull request!

[ravi-mosaicml]

I think I'm leaning towards #2. I would rather the checkpoint saver save everything, and then selectively load, rather than discard information. I would rather that parameters not being loaded are explicitly listed rather than implicitly ignored.