Save LR scheduler and grad scaler state in checkpoints

[barrosyan]

Fixes #77.

Resuming from a checkpoint re-initialized the learning rate scheduler and the AMP grad scaler, so a resumed run diverged from an uninterrupted one.

Changes:

• Trainer.save_model now stores lr_scheduler_state_dict (when a scheduler is set) and grad_scaler_state_dict.
• Trainer.load_checkpoint restores both. If a checkpoint has no scheduler state (older checkpoints), it logs a warning and keeps the scheduler at its initial state.
• Rename the misspelled optimizer_state_dit key to optimizer_state_dict. Loading still falls back to the old key, so existing checkpoints keep working.
• Add tests/benchmark/test_trainer_checkpoint.py: one test resumes mid-schedule and checks the LR, scheduler step and grad scale; one loads a checkpoint written with the legacy key. The first test fails without this change.