Getting rid of "module." heritage

[BloodAxe]

What this PR does

We do not store state_dict of the model with "module." prefix anymore

Why this is important

• When training with DDP/DP mode, DistributedDataParallel/DataParallel wrappers adds this "module." prefix unintentionally to state_dict. So there goes discrepancy when the model is trained with or without DDP. Previously SG was solving this with a hack of adding artificial "module." prefix to all saved models in single GPU mode. However, this is a broken design

• Some of our checkpoints are saved with and some are saved without "module." which I believe the main cause we have this ugly NO_KEY_MATCHING checkpoint loading hack. The actual reasoning is lost in time, but I'm, pretty sure this is because of this discrepancy of saving models with and without DDP wrapper.

Solution

1. This PR introduce get_real_model method to return the "real" model if it is wrapped with DP/DDP. Whenever there is a need to save state_dict one should use this method to unwrap the model and save state dict of the real one.
2. Similarly, when loading a checkpoint same logic should be performed: get_real_model(net).load_state_dict(...)
3. Compatibility with existing checkpoints that may or may not has "module." prefix is secured buy checking whether all state_dict keys are starting with "module.". If so - we drop them and load checkpoint as usual. This handles the case when checkpoint was saved with DDP wrapper.

Testing

Not properly tested yet
Components requiring special attention

• Regular trainer
• KDTrainer
• EMA
• Pretrained models
• Unit Tests
• Integration Tests
• Shortened Recipe Tests

Risks of introducing breaking changes

No risk assessment has been done yet

Related issues

#1163
#1153

[Louis-Dupont]

Looks great, I am just not sure about the name get_real_model, I think it is pretty confusing.
Maybe unwrap_model? It's also not that good, but maybe slightly more explicit? Any better idea?

[BloodAxe]

unwrap_model is good, I like this name)

[Louis-Dupont]

LGTM

[ofrimasad]

LGTM
not sure if the Makefile was changed by mistake

[ofrimasad]

LGTM

[Louis-Dupont]

LGTM