Fault-tolerant checklist

[awaelchli]

🚀 Fault-tolerant training progress tracker

Fault-tolerant training in PL can be activated by setting the environment variable.

PL_FAULT_TOLERANT_TRAINING=1

Description of the architecture around fault-tolerant training

• progress tracking in loops

• progress restart of loops

• capturing state of iterable datasets

  • non-random state, multiple workers
  • random state
  • sampler inside of iterable dataset
  •  batch sampler
  • testing of items above

• capturing state of map-style dataset

  • single worker global random state (add fault-tolerance for global random state in map-style datasets #8950)
  • multiple workers global random state (add fault-tolerance for global random state in map-style datasets #8950)
  • local generators
  • random sampler state, after updates in pytorch 1.10
  • testing of items above (partially done in add fault-tolerance for global random state in map-style datasets #8950)

• mid epoch restart

  • failure on mid training epoch loop

    • failure on mid training epoch loop with 1 optimizer
    • failure on mid training epoch loop with multiple optimizers multiple optimizer restart with fault-tolerant training #9537
    • failure between optimizer step and scheduler step

  • failure on mid validation epoch loop: [Feat] Add Fault Tolerant Training for ValidationLoop. #9563, Support skipping to validation #9681

    • reset on epoch end validation with 1 dataloaders
    • reset on epoch end validation with multiple dataloaders
    • reset on validation using val_check_interval: https://github.com/PyTorchLightning/pytorch-lightning/pull/9563/files#diff-c2e3dba8e80f865d593f0310cdf80eaeccdb80c9a250f7eda9edd466b7710874R989-R995

• epoch end restart

  • failure on training epoch loop end
    • failure between optimizer step and scheduler step
    • failure before logging metrics
  • no failure on training epoch end

• reloading multiple train dataloaders with individual random states (see add fault-tolerance for global random state in map-style datasets #8950 for todos in added tests)

• Fault-tolerant training with multiple dataloaders (Fault tolerant training with multiple dataloaders #11349)

• parametrize the important tests in test_auto_restart.py with num_workers > 0 and run as part of slow tests Slow CI #9086

• end-user guide - how to use and description of limitations [doc] Add Fault Tolerant Documentation Page #9256

• benchmark iteration/s for small and large models with fault-tolerant enabled vs. not enabled

• add test case for different batch structures (dict, list, etc.) fix state extraction from batch when fault-tolerant training #9281

• fix progress bar tracking on restart fix progress bar restart with fault-tolerant training enabled #9310

• Add logic to resume OptimizerLoop? Is this needed/wanted?

• Fix num_workers > 0 causing repeated random state when resuming

• Fix TODOs
  https://github.com/PyTorchLightning/pytorch-lightning/blob/ce00053002a1bb5385f7e44fceea97af50313d4c/tests/loops/test_loops.py#L899-L900
  https://github.com/PyTorchLightning/pytorch-lightning/blob/5841ca97825bd9786ab84d70f0abfa6e673528b4/tests/trainer/connectors/test_signal_connector.py#L27
  https://github.com/PyTorchLightning/pytorch-lightning/blob/5841ca97825bd9786ab84d70f0abfa6e673528b4/tests/utilities/test_auto_restart.py#L1181-L1183

Open questions:

• Do we need to capture torch.cuda.get_rng_state too?

Supplementary documents:
State of fault-tolerant training in PyTorch Lightning

cc @Borda @carmocca @justusschock @awaelchli @ninginthecloud

[tchaton]

Hey @ananthsub,

As fault tolerant training is quite complex, would you like someone on your side to be involved too ?

Best,
T.C

[aazzolini]

Do we have a design document for the proposed work?
Fault tolerance is very easy to get wrong. We have spent a lot of time at Facebook ensuring fault tolerance in our internal toolkit so we have some idea of multiple things that could go wrong.
It would be great if we went through a careful design review for this one.

[ananthsub]

@awaelchli @tchaton - we should cover metrics here as well. I am not clear around the assumptions for where/when metric states are synced, whether they're synced to rank 0 before saving, and whether we assume that we save the checkpoint from rank zero only.

https://github.com/PyTorchLightning/pytorch-lightning/blob/35876bb75f27eb8f44220afd4bfa757a0432d233/pytorch_lightning/trainer/connectors/checkpoint_connector.py#L137-L141

This does not hold for deepspeed or other use cases where we save a part of the lightning module state dict across ranks?

[tchaton]

Do we have a design document for the proposed work?
Fault tolerance is very easy to get wrong. We have spent a lot of time at Facebook ensuring fault tolerance in our internal toolkit so we have some idea of multiple things that could go wrong.
It would be great if we went through a careful design review for this one.

Hey @aazzolini,

I entirely agree, this is definitely very error prone ! We are working on this with @awaelchli and we will share it with you asap.
We would be happy to have recorded meeting to present the code too for ease onboarding.
We will really appreciate if Facebook could participate on this feature support, as I believe PyTorch needs improvement to make this process simpler.

Hey @ananthsub,

We didn't investigate the impact of sharding on TorchMetrics yet and added fault tolerant for non-sharded model for our fault tolerant V0.
IMO, we should have a mechanism to prevent sharding by default for Metric as each rank is accumulating their own state.

My intuition is that this assumption should still hold. Before saving, we accumulate the states on all ranks. So each ranks will contain the accumulated states. On reload, we just reset the metric on non-rank 0.

But it might become more intricate if metrics states get sharded, which might results in state collisions.

Best,
T.C