fix optimizer reset for relora sft

[winglian]

the relora optimizer reset wasn't pruning anything, this fixes that.

what seemed to be wrong with the original implementation is that iterating over the optimizer states isn't sufficient to determine the optimizer keys.

[winglian]

hmm, still doesn't seem to be correct,

[2024-11-21 21:17:08,843] [INFO] [axolotl.relora.reset_optimizer:72] [PID:2840] [RANK:0] Percent of optimizer states zeroed: 0.00                                                                                                                         
[2024-11-21 21:17:08,843] [INFO] [axolotl.relora.reset_optimizer:73] [PID:2840] [RANK:0] absolute n of optimizer states zeroed: 0                                                                                                                         
[2024-11-21 21:17:09,011] [INFO] [axolotl.relora.reset_optimizer:72] [PID:2841] [RANK:1] Percent of optimizer states zeroed: 0.00                                                                                                                         
[2024-11-21 21:17:09,011] [INFO] [axolotl.relora.reset_optimizer:73] [PID:2841] [RANK:1] absolute n of optimizer states zeroed: 0   

also, ooms after the first relora restart

[winglian]

reference implementation for Relora is here https://github.com/Guitaricet/relora/blob/main/peft_pretraining/training_utils.py#L267-L364