[SimpleFSDP] add manual bucketing pass #1881

ruisizhang123 · 2025-10-15T17:41:25Z

This PR adds support for aten-level manual bucketing in SimpleFSDP+aot_eager backend. Dependent on PyTorch PR

TODO List:

We should have better way of handling region info other than a list of str FQNs in current manual_bucketed_modules. It would be very easy to miss some of model modules. (cc. @xmfan @SherlockNoMad )
Currently, the reordering happens under the hood and overlap with last/next compute. We should allow users to specify which module they want to reorder.
Loss difference on multi-node training
DSV3 manual bucketing

I'll address the TODO items in follow up PRs. Let's start with this simple FSDP+TP+llama3 PR.

Performance (FSDP2 under eager mode, SimpleFSDP uses aot_eager backend)

Llama 3-8B

Performance (All Batch_size = 1). (The slower TPS on Single Node is sort of as expected, since FSDP2 handles copy-in/out in two different streams, whereas SimpleFSDP handles copy-in/out in the same stream)

Node	Method	Parallelism	Memory	TPS	Trace
1-Node (8H100)	SimpleFSDP	FSDP=8	40.96GiB(43.12%)	7,227	LINK
1-Node (8H100)	FSDP2-eager	FSDP=8	47.82GiB(50.35%)	7,380	LINK
8-Node (64H100)	SimpleFSDP	FSDP=64	29.37GiB	4,984
8-Node (64H100)	FSDP2	FSDP=64	31.41GiB	5,097
1-Node (8H100)	SimpleFSDP	FSDP=4 TP=2	28.28GiB(29.77%)	5,881	LINK
1-Node (8H100)	FSDP2	FSDP=4 TP=2	35.33GiB(37.20%)	5,898	LINK
8-Node (64H100)	SimpleFSDP	FSDP=8 TP=8
8-Node (64H100)	FSDP2	FSDP=8 TP=8

Example SimpleFSDP 1D overlapping trace:

Example SimpleFSDP 2D overlapping trace:

Bitwise Loss:

FSDP-only:

FSDP+TP:

tianyu-l

Looks nice. Had some comments.

torchtitan/experiments/simple_fsdp/backend.py

tianyu-l · 2025-10-28T23:08:59Z

torchtitan/experiments/simple_fsdp/tests/integration_tests.py

-            "1D+aot_eager_autobucketing",
-            "1d_aot_eager_autobucketing",
-        ),
+        # TODO(ruisizhang123): add back after autobucketing pass is mature


shall we add a manual bucketing test?

we should also add one in the loss unit test.

I have a few to do items for reordering. I think it'd be better to add the tests after the API is stable?

tianyu-l · 2025-10-28T23:11:43Z

torchtitan/experiments/simple_fsdp/job_config.py

-    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing"""
+    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing """
+
+    manual_bucketed_modules: list[str] = field(default_factory=list)


we need to have instructions about this field. E.g. it's not super obvious what this means "tok_embeddings,layers.[0-5],norm+output", as it involves regex I have a guess, but users might not.

btw, are the list separated by ,?

The list is separated by ,; but I didn't do explicit spilting here. essentially, it's similar to filter_fqns here

Should we add fsdp_ prefix? Or do we imagine this field will be use for other use cases, if so what are the use cases?

torchtitan/experiments/simple_fsdp/backend.py

xmfan · 2025-10-30T04:57:41Z

torchtitan/experiments/simple_fsdp/backend.py

+            manual_overlap_bucketing,
+        )
+
+        torch._inductor.config.allow_buffer_reuse = False


what happens by default?

In bucketing, we shouldn't allow buffer reuse; otherwise newly created comm copy-in/copy-out buffers will reuse prev buffer, which messed up the copied out data value and made the loss nan.

xmfan · 2025-10-30T04:59:25Z

torchtitan/experiments/simple_fsdp/job_config.py

 class Compile:
    model_backend_override: str | None = None
-    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing"""
+    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing """


should make this subclass torchtitan.config.job_config.Compile

It's additional config extended from job_config.Comfile. not sure wdym here.

something like class Compile(torchtitan.config.job_config.Compile)

fegin · 2025-10-30T17:04:10Z

torchtitan/experiments/simple_fsdp/backend.py

+    elif backend_name == "aot_eager_manualbucketing":
+        # Perform manual optimization in aten fx-level and execute code in aot_eager backend
+        # The manualbucketing logic is here:
+        bucketing_modules = compile_config.manual_bucketed_modules


This local variable is not used.

fegin · 2025-10-30T17:10:15Z

torchtitan/experiments/simple_fsdp/job_config.py

-    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing"""
+    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing """
+
+    manual_bucketed_modules: list[str] = field(default_factory=list)


Should we add fsdp_ prefix? Or do we imagine this field will be use for other use cases, if so what are the use cases?

tianyu-l · 2025-10-30T22:27:45Z

torchtitan/experiments/simple_fsdp/job_config.py

+    Manual bucket modules based on user specified FQNs
+    Abbreviations are supported to make specifying modules easier.
+    Currently, the following abbreviations are available:
+    (1) layers.[0-2] -> [layers.0], [layers.1], [layers.2]


Right now user has to know how many layer a particular flavor of model has, when applying manual bucketing. Do you think we can improve the UX by automatically resolving the number of layers?

I even think we shouldn't expose this option in toml. In toml user should just need to specify bucketing_mode = "none", "transformer_block", "auto"
And if it's transformer_block, we explicitly iterate over all the transformerblocks and pass the expanded fqns in manual_overlap_bucketing. That means manual_overlap_bucketing don't need to be smart about abbreviations.

Happy to hear people's thoughts.

I mean we could have another "manual" mode supporting Manual bucket modules if people really want to override, but a good default of transformer block level bucketing should be enabled more easily.

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 15, 2025

ruisizhang123 marked this pull request as draft October 15, 2025 17:41

ruisizhang123 force-pushed the ruisi/manual_bucket_pass branch 3 times, most recently from c20775e to a5c4027 Compare October 23, 2025 21:27

ruisizhang123 force-pushed the ruisi/manual_bucket_pass branch 4 times, most recently from 8fa2426 to 71cb39b Compare October 27, 2025 04:51

ruisizhang123 mentioned this pull request Oct 28, 2025

[simplefsdp] add manual bucketing pass pytorch/pytorch#165487

Open

ruisizhang123 force-pushed the ruisi/manual_bucket_pass branch from 71cb39b to 27bcc7d Compare October 28, 2025 07:06

ruisizhang123 marked this pull request as ready for review October 28, 2025 07:06

ruisizhang123 changed the title ~~[WIP][SimpleFSDP] add manual bucketing pass~~ [SimpleFSDP] add manual bucketing pass Oct 28, 2025

ruisizhang123 requested a review from tianyu-l October 28, 2025 07:08

ruisizhang123 force-pushed the ruisi/manual_bucket_pass branch from 27bcc7d to 3c46d64 Compare October 28, 2025 17:57

tianyu-l reviewed Oct 28, 2025

View reviewed changes

ruisizhang123 force-pushed the ruisi/manual_bucket_pass branch from 3c46d64 to d62eb25 Compare October 30, 2025 04:44

xmfan reviewed Oct 30, 2025

View reviewed changes

ruisizhang123 requested a review from tianyu-l October 30, 2025 05:02

add manual bucketing pass

1453136

ruisizhang123 force-pushed the ruisi/manual_bucket_pass branch from d62eb25 to 1453136 Compare October 30, 2025 05:21

fegin reviewed Oct 30, 2025

View reviewed changes

tianyu-l reviewed Oct 30, 2025

View reviewed changes

[SimpleFSDP] add manual bucketing pass #1881

Are you sure you want to change the base?

[SimpleFSDP] add manual bucketing pass #1881

Conversation

ruisizhang123 commented Oct 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ruisizhang123 commented Oct 15, 2025 •

edited

Loading