Initial FSDP2 support

[S1ro1]

FSDP2

What does this add?

This PR adds support for enabling FSDP2 in Accelerate as another way to parallelize your training 🚀

Why is it needed?

FSDP2 offers a lot of advantages over current FSDP1 implementation, mostly not user facing, but this enables further development of features in Accelerate such as combining TP + FSDP.
For users, this enables small improvements in memory usage

What parts of the API does this impact?

User-facing:

Allows users to specify --fsdp-version=2 to use FSDP2 instead of current FSDP1, this also changes some of the configuration options for FullyShardedDataParallelPlugin. Old configurations are still supported as before.

Basic Usage Example(s):

from accelerate import FullyShardedDataParallelPlugin, Accelerator

fsdp_plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2
    # other options...
)
accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

When would I use it, and when wouldn't I?

Using FSDP2 as a default is reccomended, with swapping to FSDP1 if something isn't supported.

[HuggingFaceDocBuilderDev]

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.

[muellerzr]

Great start! Left some initial comments

[muellerzr]

Suggested change

	"reshard_after_forward set to bool. This is not supported in FSDP1, please set to a `str` or an instance of `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`"
	f"reshard_after_forward set to {self.reshard_after_forward}. This is not supported with FSDP1, please set to a `str` or an instance of `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`"

[muellerzr]

I feel like we can simplify this

Suggested change

	if self.fsdp_version != 2 and isinstance(self.cpu_offload, CPUOffloadPolicy):
	raise ValueError(
	"cpu_offload set to `torch.distributed.fsdp.CPUOffloadPolicy`. This is not supported in FSDP1, please set to a `bool` or an instance of `torch.distributed.fsdp.CPUOffload`"
	)
	if self.fsdp_version == 2 and not isinstance(self.cpu_offload, CPUOffloadPolicy):
	raise ValueError(
	"cpu_offload set to `bool` or `torch.distributed.fsdp.CPUOffload`. This is not supported in FSDP2, please set to an instance of `torch.distributed.fsdp.CPUOffloadPolicy`"
	)
	if isinstance(self.cpu_offload, CPUOffloadPolicy):
	err = "`cpu_offload` set to `torch.distributed.fsdp.CPUOffloadPolicy`."
	if self.fsdp_version != 2:
	raise ValueError(f"{err} This is not supported in FSDP1, please set to a `bool` or an instance of `torch.distributed.fsdp.CPUOffload`")
	else:
	raise ValueError(f"{err} This is not supported in FSDP2, please set to an instance of `torch.distributed.fsdp.CPUOffloadPolicy`")

[muellerzr]

Note that in general, users will ignore warnings (and will get annoyed it bloats the logs). So we should instead use logger.warn if we don't want to explicitly raise an error about this.

[muellerzr]

Instead of doing many warning.warn,s let's accumulate them all and do one big logger.warn at the end

[muellerzr]

Suggested change

	if self.fsdp_version == 2 and self.forward_prefetch is not None:
	raise ValueError(
	"forward_prefetch is not yet implemented in FSDP2, set to None or use `fsdp_version` set to 1"
	)
	if self.fsdp_version == 2 and self.forward_prefetch is not None:
	raise ValueError(
	"forward_prefetch is not yet implemented in FSDP2, set to None or use `fsdp_version=1`"
	)

[muellerzr]

Will this import lead to issues on old pytorch versions?

[kmehant]

@muellerzr @S1ro1

I see potential collaboration on this thread with my PRs

1. Support TP + FSDPv2 / HSDP or just FSDPv2 / HSDP #3395
2. [RFC] Support FSDP2 #3231

With changes on this PR I understand that the design is to convert FSDP1 args to FSDP2 args and put to use. With my initial discussions started at #3231, In the PR #3395 I propose to have FSDP2 as a separate distributed type training protocol alongsides FSDP1. Also, to support more complex combinations of parallelisms I as well propose a design as described in the PR's description.

Looking forward.

[stevhliu]

Good breakdown of their differences, and the table at the end really helps with the comparison as well! 👏 My only feedback is:

The main feature of FSDP2 is DTensor, but I didn't really realize it's significance until the following section that illustrated the limitation of FSDP1. Now that first sentence "Simpler internal implementation, where each Parameter is a separate DTensor" is more impactful to me.

I wonder if it'd be better to start with that section so users understand what DTensor is and how it enables FSDP2 to overcome the limitations of FSDP1. Then you can discuss the other benefits of FSDP2 in the following section.

[stevhliu]

The doc-builder only supports TIP or WARNING

Suggested change

	> [!NOTE]
	> [!TIP]

[SunMarc]

cc @a-r-r-o-w

[SunMarc]

Thanks for this huge work ! Left a couple of comments

[SunMarc]

Good idea but maintaining that might be a bit of a pain. You already have added the deprecation msg for some args + not sure what will happen with this tool you will fully deprecate fsdp_sharding_strategy. wdyt @muellerzr ?

[muellerzr]

For now we should maintain until PyTorch is fully removed with the old or accelerate 2.0 IMO. It's fully breaking on the pytorch side and that hasn't sat well with me

[kmehant]

Thank you for this work! Couple of comments/questions.

Wouldn't this approach conflict with the way transformers does it with delay_optimizer_creation (https://github.com/huggingface/transformers/blob/348f3285c5114159d2ff4933b4b8ae36866d01a7/src/transformers/trainer.py#L2314) knob?

Use of transformers would need dropping use of this knob for fsdpv2.

[S1ro1]

Yes, exactly, there would be no need for this knob for FSDP2 anymore. I don't think that it is a breaking change, rather an improvement.

[byi8220]

+1 on this. I'm guessing transformers will require changes to get this working with Trainer.

[S1ro1]

Absolutely, if you do create a PR feel free to mention me, if not I'll take a look next week ish.

[byi8220]

Thanks! I actually do need to get Trainer to work with FSDP2, so if I have time I might take a look over the weekend if I have time. I'm not sure how much effort getting everything together would be.

Hopefully it can literally be a 1 line change in transformers, but I ran into some OOM issue when I tried hacking it together, so maybe it won't be that easy.

[S1ro1]

@byi8220 I haven't done any research/work on using LoRA yet. This PR was just an initial work that supports simple FSDP2, I'll try to check how it interacts with LORA in the near future. Thanks for the checks on that. I think that preliminary merging your PR for transformers integration is fine, as it should support all that is currently supported by Accelerate, so I approved for now.

[byi8220]

SG. I'll try to find a second GPU and run the tests before I officially raise the PR.

At worst this shouldn't break anything, since this is a no-op unless the user explicitly opts in.

I'll try to check how it interacts with LORA in the near future. Thanks for the checks on that.

I've mostly been hacking around, and I'm pretty unfamiliar with FSDP, so there is a high chance I am misinterpreting, or that I am not setting up my trainer properly. However everything I am suspecting as memory wastage seems logical to me (need to wrap separately, need to be clever with mmap). I might file an issue if I can make a clean reproduction

[S1ro1]

I'd suggest trying to do LORA with accelerate only first, do trainer after. As of the mmap stuff: getting that correctly has been a hassle in my limited experience with this bug/feature. Again, would leave it out of this and try to make simplest reproduction. If you won't come up with anything I'll try to look at this ~next week.

My guess is maybe we have to wrap the lora modules in isolation, and replicate what lora_fsdp_wrap_policy is doing?

I'd assume this is not the case with FSDP2, FSDP2 also known as per-parameter sharding allows each parameter to have its own metadata and allocates for those (I'm 99% sure) - you can check more about this here. Therefore this is not needed as sharding "only" influences the distributed semantics. I'm pretty certain this is needed only for FSDP1. For FSDP2 we can just shard as we would usually.

[byi8220]

I'll try to find a second GPU and run the tests before I officially raise the PR.

PR was raised, passed 1/2 new tests. The other test had to be skipped (its FSDP 1 counterpart is skipped)

I'd suggest trying to do LORA with accelerate only first, do trainer after.

SG, I'll try to get something small and raise a bug for it soon. Getting fsdp2+lora in trainer working is a high priority blocker for me, so I will look into this tomorrow

For FSDP2 we can just shard as we would usually.

Makes sense, I experimented with this a bit and found no savings.

As of the mmap stuff: getting that correctly has been a hassle in my limited experience with this bug/feature.

I can relate...

[byi8220]

Filed a bug about this, #3474

[kmehant]

Should we make this 2.6.0?

Since imports such as

accelerate/src/accelerate/utils/dataclasses.py

Line 1987 in d7c741a

from torch.distributed.fsdp import CPUOffloadPolicy, OffloadPolicy

would break for any version below 2.6.0

[kmehant]

Thank you for this check, you might also want to protect imports since the imports (such as

accelerate/src/accelerate/utils/dataclasses.py

Line 1987 in d7c741a

from torch.distributed.fsdp import CPUOffloadPolicy, OffloadPolicy

) seem to break first even before we hit this condition.

[S1ro1]

Thank you for the sanity checks. Will do more thorough version checks in the follow-up PRs.

[kmehant]

Thanks @S1ro1, appreciate your review to PR #3499 since I needed these on my setup. Thanks.