Pstjohn/stop and go test non validation

[pstjohn]

Explicitly test that train inputs and outputs are consistent when we use PreemptionCallback to handle a training interrupt.
Currently it seems that doing this changes the validation step schedule, but otherwise training should be identical.

STOP OUTPUT:

Sanity checking Validation: iteration 1/2
Sanity checking Validation: iteration 2/2
Training epoch 0, iteration 0/9 | lr: 0 | global_batch_size: 2 | global_step: 0 | reduced_train_loss: 4.87
Training epoch 0, iteration 1/9 | lr: 2e-06 | global_batch_size: 2 | global_step: 1 | reduced_train_loss: 4.771 | consumed_samples: 4
[NeMo I 2024-12-02 19:12:44 preemption:87] Received signal 12, initiating graceful stop
[NeMo I 2024-12-02 19:12:44 preemption:67] Preemption detected, saving checkpoint and exiting
2024-12-02 19:12:44,487 _dedup_tensors.py:46 INFO p:MainProcess t:MainThread: Duplicate keys to remove: {}

RESUME OUTPUT:

Restored all states from the checkpoint at /tmp/tmpn8luhenb/TestESM2StopAndGoCheckpointNotAtValidation/checkpoints/epoch=0-step=2-val_loss=0.00-last/weights
Training epoch 0, iteration 3/9 | lr: 6e-06 | consumed_samples: 8 | global_batch_size: 2 | global_step: 3 | reduced_train_loss: 4.783
Training epoch 0, iteration 4/9 | lr: 8e-06 | consumed_samples: 10 | global_batch_size: 2 | global_step: 4 | reduced_train_loss: 4.785
Validation: iteration 1/2
Validation: iteration 2/2
2024-12-02 19:12:50,690 _dedup_tensors.py:46 INFO p:MainProcess t:MainThread: Duplicate keys to remove: {}
Training epoch 0, iteration 5/9 | lr: 1e-05 | consumed_samples: 12 | global_batch_size: 2 | global_step: 5 | reduced_train_loss: 4.672 | val_loss: 4.711
Training epoch 0, iteration 6/9 | lr: 1.2e-05 | consumed_samples: 14 | global_batch_size: 2 | global_step: 6 | reduced_train_loss: 4.7 | val_loss: 4.711
Training epoch 0, iteration 7/9 | lr: 1.4e-05 | consumed_samples: 16 | global_batch_size: 2 | global_step: 7 | reduced_train_loss: 4.621 | val_loss: 4.711
Training epoch 0, iteration 8/9 | lr: 1.6e-05 | consumed_samples: 18 | global_batch_size: 2 | global_step: 8 | reduced_train_loss: 4.532 | val_loss: 4.711
Validation: iteration 1/2
Validation: iteration 2/2
Training epoch 0, iteration 9/9 | lr: 1.8e-05 | consumed_samples: 20 | global_batch_size: 2 | global_step: 9 | reduced_train_loss: 4.497 | val_loss: 4.544
`Trainer.fit` stopped: `max_steps=10` reached.

CONTINUOUS OUTPUT:

Sanity checking Validation: iteration 1/2
Sanity checking Validation: iteration 2/2
Training epoch 0, iteration 0/9 | lr: 0 | global_batch_size: 2 | global_step: 0 | reduced_train_loss: 4.87
Training epoch 0, iteration 1/9 | lr: 2e-06 | global_batch_size: 2 | global_step: 1 | reduced_train_loss: 4.771 | consumed_samples: 4
Training epoch 0, iteration 2/9 | lr: 4e-06 | global_batch_size: 2 | global_step: 2 | reduced_train_loss: 4.87 | consumed_samples: 6
Training epoch 0, iteration 3/9 | lr: 6e-06 | global_batch_size: 2 | global_step: 3 | reduced_train_loss: 4.783 | consumed_samples: 8
Validation: iteration 1/2
Validation: iteration 2/2
Training epoch 0, iteration 4/9 | lr: 8e-06 | global_batch_size: 2 | global_step: 4 | reduced_train_loss: 4.785 | consumed_samples: 10 | val_loss: 4.758
Training epoch 0, iteration 5/9 | lr: 1e-05 | global_batch_size: 2 | global_step: 5 | reduced_train_loss: 4.672 | consumed_samples: 12 | val_loss: 4.758
Training epoch 0, iteration 6/9 | lr: 1.2e-05 | global_batch_size: 2 | global_step: 6 | reduced_train_loss: 4.7 | consumed_samples: 14 | val_loss: 4.758
Training epoch 0, iteration 7/9 | lr: 1.4e-05 | global_batch_size: 2 | global_step: 7 | reduced_train_loss: 4.621 | consumed_samples: 16 | val_loss: 4.758
Validation: iteration 1/2
Validation: iteration 2/2
Training epoch 0, iteration 8/9 | lr: 1.6e-05 | global_batch_size: 2 | global_step: 8 | reduced_train_loss: 4.532 | consumed_samples: 18 | val_loss: 4.573
Training epoch 0, iteration 9/9 | lr: 1.8e-05 | global_batch_size: 2 | global_step: 9 | reduced_train_loss: 4.497 | consumed_samples: 20 | val_loss: 4.573
Validation: iteration 1/2
Validation: iteration 2/2
`Trainer.fit` stopped: `max_steps=10` reached.

[pstjohn]

/build-ci

[sichu2023]

Good to go. Only some comments to help me understand it better.

[pstjohn]

/build-ci

[pstjohn]

/build-ci