add act checkpoint at sub layer level

[cli99]

This PR allows users to specify what module to do activation checkpoint through fsdp when activation_checkpointing is true in fsdp_config. It checks the field activation_checkpointing_target in model_config, if nothing is specified but activation_checkpointing is set to true, then MPTBLOCK is act checkpointed. To checkpoint the submodules in ATTN_CLASS_REGISTRY, FFN_CLASS_REGISTRY,NORM_CLASS_REGISTRY, add the corresponding names to activation_checkpointing_target.
For example,

activation_checkpointing_target: 
    - grouped_query_attention

checkpoints the activation of GroupedQueryAttention.
If multiple submodules are up for act checkpointing, add their names to activation_checkpointing_target

[dakinggg]

Two high level comments:
(1) Please add some tests. Not sure how hard it is to test that act ckpt works the way you think it does, but at least test that everything is constructed the way you expect.
(2) Should we add this to all HF models too? I don't see why not. That would go here:

llm-foundry/llmfoundry/models/hf/hf_fsdp.py

Line 200 in 84c86e3

model.activation_checkpointing_fn = lambda module: isinstance(

[cli99]

Two high level comments: (1) Please add some tests. Not sure how hard it is to test that act ckpt works the way you think it does, but at least test that everything is constructed the way you expect. (2) Should we add this to all HF models too? I don't see why not. That would go here:

llm-foundry/llmfoundry/models/hf/hf_fsdp.py

Line 200 in 84c86e3

model.activation_checkpointing_fn = lambda module: isinstance(

(1) are there some existing act ckpt tests to leverage? Or shall we just check to make sure the module names passing works as expected?
(2) sub-block act checkpointing requires we know the module definitions inside the block , it's model dependent.

[dakinggg]

(1) Nope. Probably ok to just test that the module passing works as expected. Would be great if we actually tested activation checkpointing, but that does not exist right now, no.
(2) Right, so maybe you have to provide the full import path of the module you want to checkpoint? Is that reasonable?

[vchiley]

@dakinggg re (2) this member method is a member method of the MPT model.

enabling this for HF models would requires a custom str to class conversion for a list for every hf model class before it is used in this call...
probably outside the scope of this PR 🤷‍♂️
(not sure how maybe you have to provide the full import path of the module you want to checkpoint would work)

[vchiley]

lgtm

I'd wait for @dakinggg approval before merging

[dakinggg]

@vchiley you would do something like import('string name of class') to get the class. but yeah, doesn't need to be part of this PR. @cli99 can you make a jira for adding custom act ckpt to generic HF models and put it in the LLM Foundry backlog jira epic please?

[dakinggg]

Also would like basic unit tests testing the use of this param before merging

[cli99]

@dakinggg , added a unit test