Fix: resume_training: True and EMACallback active

[dtronmans]

Purpose

• Previous problem: when EMACallback was enabled, training for a few epochs, then training again with resume_training True in the config and passing the previous checkpoint using the weights argument was leading to an error.

Specification

• The error was the following:

│ ❱  70 │   │   │   for ema_v, model_v in zip(                                                     │
│    71 │   │   │   │   self.state_dict_ema.values(),                                              │
│    72 │   │   │   │   model.state_dict().values(),                                               │
│    73 │   │   │   │   strict=True,                                                               │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: zip() argument 2 is longer than argument 1

• By allowing state dict loading that was not strict (load matching keys and log keys that only exist in either the EMA state dict or in the .ckpt file) I found that there were duplicate _node keys in the live EMA state that were filtered out of the saved checkpoint
• Fix: no more zip(..., strict=True) for EMA update, just match entries by key and log the asymmetric entries
• Ignore _node alias-only asymmetries during EMA loading since these keys are filtered out of saved checkpoints

Dependencies & Potential Impact

Deployment Plan

Testing & Validation

• test_resume_training_with_ema_does_not_crash: train 1 epoch with EMA enabled and overfit_batches=1, start a new run with resume_training=True, pass the previous checkpoint through train(weights=...), assert it does not crash
• Manual validation: