enable LoRA + FSDP2

[weifengpy]

how to run it

• 7B: tune download meta-llama/Llama-2-7b-hf --output-dir /tmp/Llama-2-7b-hf --hf-token <HF_TOKEN> && tune run --nproc_per_node 8 lora_finetune_fsdp2 --config llama2/7B_lora
• 70B: tune download meta-llama/Llama-2-70b-hf --output-dir /tmp/Llama-2-70b-hf --hf-token <HF_TOKEN> && tune run --nproc_per_node 8 lora_finetune_fsdp2 --config llama2/70B_lora

recipe tests: pytest tests/recipes/test_lora_finetune_fsdp2.py -m integration_test
unit test: pytest tests/torchtune/utils/test_distributed.py -k test_state_dict

Highlights

• +12% tokens per second (see wandb snapshot) with +0.8% memory reserved
• 3.2 x speedup in model init: from 250s to 77s

FSDP2 changes

• init models on device=meta instead of cpu
• activation checkpointing happens before register FSDP hooks

checkpoint changes

• loading: load HF checkpoint with mmap=True. convert plain tensor into DTensor
• saving: gather full state dict on rank0 cpu. convert DTensor into plain tensor
• resume finetuning from previous checkpoints

optional

• foreach=True for adam + DTensor through optimizer.foreach=True since pytorch PR have not land yet [DTensor] Turn on foreach implementation of optimizer for DTensor by default pytorch#123394

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/855

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8fbbc4b with merge base 135cf2e ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[awgu]

I think one main question for discussion is what the torchtune folks feel about defining an explicit initializer for the RoPE theta buffer, which would be required to do meta-device init (which should speed up initialization!).

[kartikayk]

My noob understanding here is that you want to account for the Linear module before the TransformerDecoderLayer - is that right? If so, is there a better way to do this? It wasn't immediately obvious to me that reversing this list gets me that outcome - but maybe that's because I haven't played around with the modules function? Disregard if this is a pretty standard way to achieve this

[awgu]

I wanted to quickly mention my opinion:

What we want is a post-order traversal of the modules so that we visit lower modules before higher ones. nn.Module.modules() always gives reverse post-order, so reversing it gives post-order.

FSDP1 hides all of this under the auto_wrap_policy argument, which does the post-order traversal for the user. Personally, I did not want to have an auto_wrap_policy argument for FSDP2 because the name is a misnomer. Auto wrapping is not automatically choosing which modules to wrap for the user -- rather, it just performing a post-order traversal and applying a user-defined policy to determine if a module is wrapped or not. What I am open to though is having some other higher-level utility (say, apply_fsdp()) that does the same thing as auto_wrap_policy but exists outside of the function fully_shard.

[ebsmothers]

Thanks for this thorough explanation, it helps to understand this code block a lot. I agree it'd be nice to have an apply_fsdp utility but not a major concern here. Can we add code comments here to get the point you described across? (Basically that we are iterating over lower-order modules to higher-level modules and wrapping individual transformer layers. And I assume the separate wrapping of trainable LoRA weights is more related to the point you mentioned today about lower memory, rather than the flat param rationale of grads being allocated per shard?) Users will take this recipe code as a starting point so the more explicit we are here the easier they'll find it to extend.

[kartikayk]

Got you, thanks so much for the explanation! I think something that would be super helpful would be document here in the form of comments the relationship between:

• the modules on which we call fully_shard
• init on meta device
• calling initialize_parameters and reset_parameters

Also I think there was a technical reason with FSDP1 to call the function reset_parameters. Is that still true? Or can we standardize this with initialize_parameters in the modules code? Happy to chat about this offline!

[kartikayk]

Does this need to be LoRALinear? Or am I misunderstanding?

[weifengpy]

nn.Linear with required_grad=True are referring to lora_a and lora_b

Trunk FSDP1 is wrapping lora_a and lora_b separately thus we are do the same wrapping for parity. But as you mentioned, we can wrap LoRALinear instead so lora_a and lora_b are communicated together

[kartikayk]

Just to clarify, we special handle RoPE because the buffer is not being loaded from a state dict, right?

[weifengpy]

that's correct

[ebsmothers]

Similar comment here, let's document what's happening so that users can easily understand why we initialize these modules separately.

[kartikayk]

This is a bit unfortunate but we have two different meanings of "full model" IIUC:

• One is related to FSDP i.e. full vs sharded - is that right?
• Other is LoRA i.e. full model vs LoRA adapters

The current function is a bit confusing since we pass in the adapter state_dict and this is referred to as full_sd in the function itself.

[weifengpy]

good catch. you are right "full" is opposing to "sharded" here. What do you think if rename to "local"? load_from_local_model_state_dict ?

[ebsmothers]

Can we add docstrings for the functions added in this file?

[kartikayk]

Where is device_mesh and placements information coming from?

[weifengpy]

it's from fully_shard(model, mesh). we are using default mesh where every rank serves FSDP, since there is no 2D/3D parallasim involved

after fully_shard(model, mesh), model.parameters() is converted from plain tensor to DTensor with mesh

[weifengpy]

synced with @kartikayk on action items

• add docstring for meta init, state dict, migration to DCP
• place FSDP2 in receipes/dev/
• attach wandb loss curve
• pending feedbacks on unit tests

[ebsmothers]

This is looking great! Most of my comments are minor and around documentation.

I saw there was also some discussion around the location for this recipe. I am inclined to agree with the point you and @kartikayk discussed around putting it in recipes/dev. I don't want that to be a long-term home for it, but also want to make sure we don't break folks who are using the current recipe and aren't on the latest version of PyTorch yet.

[ebsmothers]

Thanks for this thorough explanation, it helps to understand this code block a lot. I agree it'd be nice to have an apply_fsdp utility but not a major concern here. Can we add code comments here to get the point you described across? (Basically that we are iterating over lower-order modules to higher-level modules and wrapping individual transformer layers. And I assume the separate wrapping of trainable LoRA weights is more related to the point you mentioned today about lower memory, rather than the flat param rationale of grads being allocated per shard?) Users will take this recipe code as a starting point so the more explicit we are here the easier they'll find it to extend.

[ebsmothers]

Similar comment here, let's document what's happening so that users can easily understand why we initialize these modules separately.

[ebsmothers]

Just fyi might need a merge, I think these are actually now defined in the version of recipe we have in main

[ebsmothers]

To make sure I understand, our distributed LoRA recipe will now only work on torch >= 2.4.0, is that correct?

[weifengpy]

that's correct. FSDP2 will be released in 2.4.0, although it has been in nightly for a while

[ebsmothers]

Thanks for clarifying! In that case I think it makes sense to create dev/recipes/lora_finetune_fsdp2.py (or something like that) in the short-term, then migrate to replace recipes/lora_finetune_distributed.py once we've socialized and gotten enough users onto >= 2.4.0. Let me know how that sounds to you.

[weifengpy]

proposing a middle ground: type hinting with "FSDPModule" without try...except

def load_from_full_model_state_dict(
    model: "FSDPModule",

[ebsmothers]

We do also have gpu_test in torchtune, can we use that here for the sake of consistency?

[ebsmothers]

Can we add docstrings for the functions added in this file?

[ebsmothers]

If we catch missing and unexpected keys from load_state_dict with strict=False what format will the keys be in? Previously with FSDP1 the keys contained all the info about FSDP wrapping. E.g.model.layers.0._fsdp_flat_param.attn.q_proj.weight (probably not exactly right but something like that). Will that still be the case here?

[weifengpy]

for FSDP2, it's clean FQNs without FSDP prefix. For example, layers.0.attn.q_proj.lora_a.weight

FSDP2 is clean because 1) fully_shard register hooks instead wrap nn.Module, 2) fully_shard set module.__class__ = type(f"FSDP{cls.__name__}", (FSDPModule, cls), dct) https://fburl.com/i20yr3s2

[ebsmothers]

This is great! I think this means we can actually do validation of LoRA state dict load more cleanly (note that we actually have two separate utilities for this for the single-device vs distributed case because of this FSDP prefix issue). Not a concern for this PR but this will allow us to clean up our code a bit

[ebsmothers]

Might be useful to add some comments for this code block

[ebsmothers]

A couple small comments on the configs but otherwise no major concerns. Really excited to see this in our library!

[ebsmothers]

Probably need to do a find and replace of lora_finetune_distributed -> lora_finetune_fsdp2 in all three config files

[ebsmothers]

Can update this header to mention that this config is for the recipe using FSDP2 (I know the config file is the same, but nice visibility to just explicitly call it out at the top of the file)

[weifengpy]

good catch! updated to lora_finetune_fsdp2 and mentioned FSDP2

[awgu]

By the way, another option here is to just make two passes if that is clearer.

for module in model.modules():
    if <LoRA adapter>:
        fully_shard(module)
for module in model.modules():
    if <transformer block>:
        fully_shard(module)