✌️ Add support for FSDP2 #3317
Merged
✌️ Add support for FSDP2 #3317
+54
−26
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
I think it's ok to bump high since trl is also beta
So it seems that |
|
I tried myself: looks good! and for the record, FSDP is way faster in my case: |
|
Hi we had a meeting with @qgallouedec and he mentioned that i should update two thoughts on FSDP2.
|
@qgallouedec Would you mind sharing the script you ran? Cause I got errors trying to use fsdp2 with this PR. # accelerate launch --config_file fsdp.yaml error.py
import torch
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTConfig, SFTTrainer
from accelerate import PartialState
from peft import get_peft_model, LoraConfig
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map={"": PartialState().process_index},
attn_implementation="flash_attention_2",
torch_dtype=torch.bfloat16,
use_cache=False,
)
# apply_liger_kernel_to_qwen2(model)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token_id = 151643
dummy_dataset = Dataset.from_dict({"text": ["Dummy dataset"] * 32, })
training_args = SFTConfig(
output_dir="trainer_output",
num_train_epochs=1,
per_device_train_batch_size=2,
gradient_accumulation_steps=2,
report_to="none",
bf16=True,
gradient_checkpointing=True,
gradient_checkpointing_kwargs={"use_reentrant": False}
)
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dummy_dataset,
processing_class=tokenizer,
)
trainer.train()
trainer.save_model()error: |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
FSDP2 fails with LoRA and the following error: Update: fixed by huggingface/accelerate#3545 |
@lewtun should we merge in the meantime? |
|
@lewtun @qgallouedec yea maybe we need to write some custom summoning code like this |
qgallouedec
approved these changes
May 5, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do?
This PR adds support for FSDP2 by (a) adding an
accelerateconfig and (b) updatingprepare_fsdp().Useful migration guide: https://huggingface.co/docs/accelerate/main/en/concept_guides/fsdp1_vs_fsdp2
The FSDP2 API is still in beta, so we can hold off on merging this if we feel the
accelerateversion is too high to bump right now. The alternative would be to check the accelerate version being used and only allow FSDP2 to be run for versions > 1.6.0In particular,FULL_STATE_DICTis not supported in FSDP2 yet, so it is not possible to savetransformersfriendly checkpoints. This is currently being worked on inaccelerate, so we could hold off until thenUpdate: now supported on
accelerate@mainand tested with:A quick test on SFT shows the two FSDP versions are almost identical, modulo grad norms:
Commands to test with on 2 GPUs
accelerate launch --config_file examples/accelerate_configs/fsdp2.yaml --num_processes 2 trl/scripts/sft.py \ --model_name_or_path Qwen/Qwen2-0.5B \ --dataset_name trl-lib/Capybara \ --learning_rate 2.0e-5 \ --num_train_epochs 1 \ --packing \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 8 \ --gradient_checkpointing \ --eos_token '<|im_end|>' \ --logging_steps 25 \ --eval_strategy steps \ --eval_steps 100 \ --output_dir Qwen2-0.5B-SFT # LoRA accelerate launch --config_file examples/accelerate_configs/fsdp2.yaml --num_processes 2 trl/scripts/sft.py \ --model_name_or_path Qwen/Qwen2-0.5B \ --dataset_name trl-lib/Capybara \ --learning_rate 2.0e-4 \ --num_train_epochs 1 \ --packing \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 8 \ --gradient_checkpointing \ --eos_token '<|im_end|>' \ --logging_steps 25 \ --eval_strategy steps \ --eval_steps 100 \ --use_peft \ --lora_r 32 \ --lora_alpha 16 \ --output_dir Qwen2-0.5B-SFTWorking SFT variants
Before submitting
Pull Request section?
to it if that's the case.
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.