[fix auto-microbatch] FSDP reshard and cleanup after OOM to fix the cuda memory leak

[bigning]

What does this PR do?

There are two issues with the current auto-microbatch.

1. nccl timeout or hang
2. it finds a lower device train micro-batch size or still OOM with device_train_microbatch_size = 1

This PR targets at 2nd issue. PR 3016 (WIP) targets the 1st issue.

For the 2nd memory issue, it's because if OOM happens during fsdp fwd/bwd after unshard, we need to manually reshard and cleanup. Otherwise, there is memory leaks like this:

There is an existing FSDP backward callback _post_backward_final_callback which tries to reshard and cleanup at the end of fsdp backward. We just call this API after OOM happens.

test

unit test

python -m composer.cli.launcher -n 2 -m pytest -m gpu tests/trainer/test_fsdp.py -k test_fsdp_reshard_after_oom
python -m composer.cli.launcher -n 2 -m pytest -m gpu tests/trainer/test_fsdp.py -k test_fsdp_same_state_after_oom_reshard

end-to-end test

llama2 13B finetuning on 16 A100-40G GPUs with global train batch size = 1024

• baseline: llama2-13b-baseline-new-Ln2FYB OOMed even with device-train-microbatch-size = 1
• test: llama2-13b-test-16-40g-r7z22-IBHkBe run well with device-train-microbatch-size = 2

llama2 70B finetuning on 32 H100-80G GPUs with global train batch size = 1024

• baseline: llama2-70b-8192-32-h100-80g-baseline-npn6Qj OOMed even with device-train-microbatch-size = 1
• test: llaama2-70b-8192-32-h100-80g-test-vbLmLL chose device-train-microbatch-size=2 and run well before evaluation. After evaluation, it has nccl timeout error, which is a different issue. We'll debug and fix the timeout in a separate PR.
• fixed device-train-microbatch-size=4: llama2-70b-32-h100-80g-fixed-4-KKDQ7I OOMed.

[mvpatel2000]

This is great! A few minor nits.

Would like @cli99 's thoughts too

[cli99]

Nice test. Looks good to me.

[awgu]

Thanks for calling this out and great catch! I think that calling the _post_backward_final_callback sounds good to me.

[bigning]

Thanks for calling this out and great catch! I think that calling the _post_backward_final_callback sounds good to me.

Thank you @awgu for the quick response!

[mvpatel2000]

This is a really nice PR. Super great catch

[mvpatel2000]

nit