Arde/fsdp activation checkpointing

[arde171]

What does this PR do?

Currently, HF Trainer didn't support FSDP activation checkpointing. This PR provides support to FSDP activation checkpointing.
Please see the details about the FSDP activation checkpointing here.
I saw the improvement in training performance for the large LLM models (e.g., LLAMA 70B) with FSDP activation checkpointing as compared to the existing gradient_checkpointing option. It's easy to enable FSDP activation_checkpointing.

we just need to add "activation_checkpointing": "True" to enable the FSDP activation_checkpointing as shown in below example fsdp_config.json file.

fsdp_config.json

{
  "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
  ...
  "activation_checkpointing": "True"
}

Please see the below PR for more details about FSDP activation checkpointing in accelerate repo:
PR: huggingface/accelerate#1891

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[ArthurZucker]

Cc @pacman100 if you think this is relevant 🤗

[pacman100]

Thank you @arde171 for adding this. Please raise an error if both activation_checkpointing in FSDP config and training arg gradient_checkpointing are set to True. The error should mention that both can't be set to True and to use FSDP's checkpointing logic when using FSDP.

[pacman100]

Thank you @arde171!