Checkpointing Issue about openfold

[gahdritz]

If you're going to be training on more than one GPU, we recommend DeepSpeed, partly for this reason. The native PyTorch DDP implementation doesn't interact very well with activation checkpointing (if you're curious, look up the "find_unused_parameters" option). Basically you have to choose between having activation checkpoints and using every single parameter in the model for every forward pass, which isn't ideal. DeepSpeed doesn't have such requirements. If you really want to avoid DeepSpeed (which also has memory/runtime advantages over the native stuff), change strategy = "ddp" to strategy = DDPPlugin(find_unused_parameters=False) in train_openfold.py (importing DDPPlugin from pytorch_lightning.plugins.training_type in version 1.5.0)