Autoparallel as an experiment into main #2054

xmfan · 2025-11-18T17:12:15Z

Pending description

🚧 WIP 🚧

cleanup files outside the experiment folder

Experiments like SimpleFSDP/PT2-Frontier/Autoparallel are all being developed at the same time, and SimpleFSDP/PT2-Frontier both run into issues with PP that requires the PP utilities from Autoparallel. We want to land the Autoparallel experiment into main to facilitate that sharing.

TODO - try converting model params into fake tensors - figure out init fn - integrate torchtitan configs for DP/TP to control autop Hack an init_fn for llama3 and observe loss decreasing with autoparallel """ [rank0]:[titan] 2025-06-16 16:24:16,593 - root - INFO - Training starts at step 1. [rank0]:[titan] 2025-06-16 16:24:23,544 - root - INFO - step: 1 loss: 8.1880 memory: 4.88GiB(6.16%) tps: 28 [rank0]:[titan] 2025-06-16 16:24:23,545 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40 [rank0]:[titan] 2025-06-16 16:24:23,842 - root - INFO - step: 2 loss: 8.1610 memory: 4.90GiB(6.20%) tps: 13,785 [rank0]:[titan] 2025-06-16 16:24:24,135 - root - INFO - step: 3 loss: 8.0871 memory: 4.90GiB(6.20%) tps: 14,006 [rank0]:[titan] 2025-06-16 16:24:24,433 - root - INFO - step: 4 loss: 7.9516 memory: 4.90GiB(6.20%) tps: 13,770 [rank0]:[titan] 2025-06-16 16:24:24,727 - root - INFO - step: 5 loss: 7.8552 memory: 4.90GiB(6.20%) tps: 13,959 [rank0]:[titan] 2025-06-16 16:24:25,023 - root - INFO - step: 6 loss: 7.7732 memory: 4.90GiB(6.20%) tps: 13,859 [rank0]:[titan] 2025-06-16 16:24:25,324 - root - INFO - step: 7 loss: 7.6987 memory: 4.90GiB(6.20%) tps: 13,664 [rank0]:[titan] 2025-06-16 16:24:25,617 - root - INFO - step: 8 loss: 7.6779 memory: 4.90GiB(6.20%) tps: 13,985 [rank0]:[titan] 2025-06-16 16:24:25,911 - root - INFO - step: 9 loss: 7.6043 memory: 4.90GiB(6.20%) tps: 13,962 [rank0]:[titan] 2025-06-16 16:24:26,207 - root - INFO - step: 10 loss: 7.5778 memory: 4.90GiB(6.20%) tps: 13,891 """ Adopt new autoparallel API with meta-init model Allows reverting a lot of the hacks in the original integration that were caused by not creating a model obj in the train.py due to passing a model_fn builder to autop. Fixes to align with latest autoparallel Add inductor config knobs for comms optimizations to torchtitan Make inductor always run compile passes basically, this is an annoying workaround for debugging iteratively. 1- you run the model, it compiles, but something weird happens 2- you enable some logging or tlparse, rerun. but inductor decides not to run your pass anymore, its results are cached. since (2) has confused me horribly on more than one occasion, i just disable caching for now Drop hacky llama3_init_fn and use autop init_weights feature Relying on meta-pytorch/autoparallel#20, this lets us automatically apply a user's init_weights fn to the autoparallel model. Verified this works with `CONFIG_FILE="./torchtitan/models/llama3/train_configs/debug_model.toml" ./run_train.sh --model.name llama3_auto_parallel --parallelism.tensor_parallel_degree 4 --training.dataset c4` ``` [rank0]:[titan] 2025-07-02 16:18:02,007 - root - INFO - Training starts at step 1. [rank0]:[titan] 2025-07-02 16:18:08,224 - root - INFO - step: 1 loss: 8.1848 memory: 1.09GiB(1.14%) tps: 77 tflops: 0.01 mfu: 0.00% [rank0]:[titan] 2025-07-02 16:18:08,224 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40 [rank0]:[titan] 2025-07-02 16:18:08,310 - root - INFO - step: 2 loss: 8.1619 memory: 1.15GiB(1.21%) tps: 48,138 tflops: 3.46 mfu: 0.35 % [rank0]:[titan] 2025-07-02 16:18:08,356 - root - INFO - step: 3 loss: 8.1140 memory: 1.15GiB(1.21%) tps: 88,440 tflops: 6.36 mfu: 0.64 % [rank0]:[titan] 2025-07-02 16:18:08,406 - root - INFO - step: 4 loss: 8.0099 memory: 1.15GiB(1.21%) tps: 82,626 tflops: 5.94 mfu: 0.60 % [rank0]:[titan] 2025-07-02 16:18:08,457 - root - INFO - step: 5 loss: 7.8928 memory: 1.15GiB(1.21%) tps: 81,594 tflops: 5.87 mfu: 0.59 % [rank0]:[titan] 2025-07-02 16:18:08,508 - root - INFO - step: 6 loss: 7.7758 memory: 1.15GiB(1.21%) tps: 79,607 tflops: 5.72 mfu: 0.58 % [rank0]:[titan] 2025-07-02 16:18:08,559 - root - INFO - step: 7 loss: 7.6221 memory: 1.15GiB(1.21%) tps: 81,448 tflops: 5.86 mfu: 0.59 % [rank0]:[titan] 2025-07-02 16:18:08,611 - root - INFO - step: 8 loss: 7.5578 memory: 1.15GiB(1.21%) tps: 79,732 tflops: 5.73 mfu: 0.58 % [rank0]:[titan] 2025-07-02 16:18:08,659 - root - INFO - step: 9 loss: 7.3851 memory: 1.15GiB(1.21%) tps: 85,655 tflops: 6.16 mfu: 0.62 % [rank0]:[titan] 2025-07-02 16:18:08,709 - root - INFO - step: 10 loss: 7.3361 memory: 1.15GiB(1.21%) tps: 81,855 tflops: 5.89 mfu: 0.60 % [rank0]:[titan] 2025-07-02 16:18:08,709 - root - INFO - Sleeping 2 seconds for other ranks to complete ``` fix lint

lets existing torchtitan knobs which govern DP/TP mesh creation and mesh size influence the sharding constraints of autoparallel, allowing it to support these different sharding configurations.

Signed-off-by: Edward Z. Yang <[email protected]>

- fix passing of "none" (not None) to control bucketing passes

since meta-pytorch/autoparallel#52 landed

prints an (internal, vpn) only link for each profile trace file that's saved to manifold. Just search for 'trace' in your job logs on mast, and click one of the rank links. e.g. [trainer37|5]:[titan] 2025-08-07 14:21:01,227 - root - INFO - Finished dumping profiler traces in 5.22 seconds: https://interncache-all.fbcdn.net/manifold/perfetto-artifacts/tree/ui/index.html#!/?url=https://interncache-all.fbcdn.net/manifold/torchtrain_datasets/tree/outputs/torchtitan-64-whc-jv2j4mp/profile_trace/iteration_20/rank37_trace.json

This PR makes bucket sizes for all-gather and reduce-scatter to be of the same size for 1d FSDP.

IMO we should just add the loss in the model and let autoparallel parallelize it for us. But for now, let's follow how the other models are implemented

just add `--experimental.enable_simplefsdp_passes` and do not try to combine it with other `bucket_*` or `reorder_*` options.

This command should now run `CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" ./run_train.sh --model.name deepseekv3_auto_parallel` However it doesn't actually do anything with autoparallel yet. Next step is to attach local_map to the model so that autoparallel can run.

Validated debugmodel llama3 works, but ds3 crashes becuase of `build_optimizers_with_moe_load_balancing` doing stuff that traverses the original model structure, only now its an AutoParallelModule which isn't compatible, we'll have to disable this optimization for now and think about what to do. Note: paths have changed, update your run commands: `CONFIG_FILE=./torchtitan/models/llama3/train_configs/debug_model.toml ./run_train.sh --model.name llama3_auto_parallel --parallelism.tensor_parallel_degree 4` Failing (ds3): `CONFIG_FILE=././torchtitan/models/deepseek_v3/train_configs/debug_model.toml ./run_train.sh --model.name deepseekv3_auto_parallel`

…1684)

as titled, this pr adds entry to simplefsdp's autobucketing pass in autoparallel. original code is in: pytorch/pytorch#160282 The main code for autobucketing pass will be added to autoparallel repo.

needs to merge in lock step with meta-pytorch/autoparallel#233

wconstab · 2025-11-18T17:21:57Z

torchtitan/train.py

            gc_freq=job_config.training.gc_freq, debug=job_config.training.gc_debug
        )

+        # TODO(whc)


i think we can just move this into the parallelize thing in autop experiment

addressed locally, will push when github is back up

wconstab · 2025-11-18T17:22:58Z

torchtitan/train.py

-        # Build the collection of model converters. No-op if `model.converters` empty
-        model_converters = build_model_converters(job_config, parallel_dims)
-        model_converters.convert(model)
+        with torch.device("meta"):


i'm not sure what to do with this. worst case we could 'move' the model to meta inside the autop parallelism thing, if that works.

i guess i'm surprised that titan builds the 'real' model though. i thought titan always built a fake model, parallelized it, then moved it to device. i'll take a look at this

Hmm, it is just right above this line? We are using meta device. It seems to me this change is not necessary.

yeah idk what happened here, without this change we run fine, im just gonna remove it

wconstab · 2025-11-18T17:23:36Z

torchtitan/tools/profiling.py

-            logger.info(
-                f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
-            )
+            log_str = f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"


this one is not actually autop-specific. it's just 'mast specific'. we might benefit from just landing it. but we can always drop this patch from what we land to main and keep an 'autop' branch with just this patch for doing mast runs

@tianyu-l @wwwjn what's the best place to put this? maybe as a flag to train.py?

Agree this is not autop-specific. Maybe we can put it in torchtitan mast launch scripts?

For internal specific code, we should follow PyTorch convention to put it under torchtitan/fb and replace it when launching internally. Unfortunately, we have not built any flow for this yet.

as i said above, to decouple this, its fine to just delete this when landing to main. we can keep an 'autoparallel' branch alive with just this minimal set of unlanded to main things, to support autoparallel mast launches. later, we can move it to fb/ folder if we have a way to do that

xmfan · 2025-11-18T17:31:52Z

torchtitan/models/moe/moe.py

        )


+def _moe_forward(


this was needed to expose something to local_map, i'm gonna try to limit this change to AP

wconstab · 2025-11-18T17:49:00Z

torchtitan/config/job_config.py

+    # "aten" (default), "inductor", "none"
+    comms_bucket_reorder_strategy: str = "aten"
+
+    autop_force_bf16: bool = False


Is there a way for an experiment to add a config knob without polluting the top-level file?

https://github.com/pytorch/torchtitan/blob/main/docs/extension.md#extending-jobconfig

wconstab · 2025-11-18T17:49:32Z

torchtitan/config/job_config.py

    """

+    # "aten" (default), "inductor", "none"
+    comms_bucket_reorder_strategy: str = "aten"


this one might be worth landing in main. it's purely inductor-specific and used by simple-fsdp as well.
cc @ruisizhang123 @fegin @tianyu-l

If it is shared by 2 different models/experiments, I think it is okay to add it to the core job_config. This will be used by full dtensor and compiler toolki iiuc.

It's inductor specific bucketing. Maybe we should upstream inductor bucketing code to pytorch?

What happens if users run core FSDP2 with this option? Is it a no op?

I think we encounter a tricky case where multiple experiments would share config that's not in core. I would say the "right" way for now might be just duplicating this config into their own custom job_config.py, but deeper reason is that we need to reinvent the config system -- the idea is to let each component have its own config, rather than sharing a central config.

This seems a bit urgent, as we are hitting such issues from different angles, recently.

cc @ailzhang

Nothing would happened. Because the graphs saw by compiler are only compute graphs. Then, the bucketing pass would not taking into effect (as no comms are bucketed/reordered).

Maybe we should have a config class specific for pt2-frontier lolll

wconstab · 2025-11-18T17:50:05Z

torchtitan/components/optimizer.py

+                blocks = model_part.layers.values()
+            else:
+                # TODO: fix autoparallel to preserve the module dict
+                blocks = model_part.layers.children()


anyone remember the history on this? do we still need it?

Yeah, I needed this for DSv3, autoparallel doesn't transparently wrap the original model. The original model in torchtitan uses self.layers as a nn.ModuleDict, but I think the AP wrapper has self.layers as a nn.Module or something

If this is solely used by autoparallel, this should be moved to experiments/auto_parallel.

fegin · 2025-11-18T21:46:13Z

torchtitan/train.py

-        # Build the collection of model converters. No-op if `model.converters` empty
-        model_converters = build_model_converters(job_config, parallel_dims)
-        model_converters.convert(model)
+        with torch.device("meta"):


Hmm, it is just right above this line? We are using meta device. It seems to me this change is not necessary.

fegin · 2025-11-18T21:46:28Z

torchtitan/train.py


        # build model (using meta init)
        model_args = self.train_spec.model_args[job_config.model.flavor]
+        model_cls = self.train_spec.model_cls


Remove this.

fegin · 2025-11-18T21:46:50Z

torchtitan/train.py

            ensure_pp_loss_visible(parallel_dims, job_config, color)
        else:
-            # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel
+            # apply Autoparallel


Merge conflict?

fegin · 2025-11-18T21:48:35Z

torchtitan/tools/profiling.py

-            logger.info(
-                f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
-            )
+            log_str = f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"


For internal specific code, we should follow PyTorch convention to put it under torchtitan/fb and replace it when launching internally. Unfortunately, we have not built any flow for this yet.

fegin · 2025-11-18T21:49:09Z

torchtitan/models/deepseek_v3/__init__.py

+        # use_flex_attn=True,
+        # attn_mask_type="block_causal",
    ),
    "236B": DeepSeekV3ModelArgs(
        vocab_size=102400,
        dim=5120,
        inter_dim=12288,
        moe_inter_dim=1536,
        n_layers=60,
        n_dense_layers=1,
        n_heads=128,
        moe_args=MoEArgs(
            num_experts=160,
            num_shared_experts=2,
            top_k=6,
            score_func="softmax",
            route_norm=False,
            route_scale=16.0,
            score_before_experts=False,
        ),
        n_expert_groups=8,
        n_limited_groups=3,
        q_lora_rank=1536,
        kv_lora_rank=512,
        qk_nope_head_dim=128,
        qk_rope_head_dim=64,
        v_head_dim=128,
-        use_flex_attn=True,
-        attn_mask_type="block_causal",
+        # use_flex_attn=True,
+        # attn_mask_type="block_causal",
    ),
    "671B": DeepSeekV3ModelArgs(
        vocab_size=129280,
        dim=7168,
        inter_dim=18432,
        moe_inter_dim=2048,
        n_layers=61,
        n_dense_layers=3,
        n_heads=128,
        moe_args=MoEArgs(
            num_experts=256,
            num_shared_experts=1,
            top_k=8,
            score_func="sigmoid",
            route_norm=True,
            route_scale=2.5,
            score_before_experts=False,
        ),
        n_expert_groups=8,
        n_limited_groups=4,
        q_lora_rank=1536,
        kv_lora_rank=512,
        qk_nope_head_dim=128,
        qk_rope_head_dim=64,
        v_head_dim=128,
-        use_flex_attn=True,
-        attn_mask_type="block_causal",
+        # use_flex_attn=True,
+        # attn_mask_type="block_causal",


You have to revert all of these.

fegin · 2025-11-18T21:51:00Z

torchtitan/config/job_config.py

    """

+    # "aten" (default), "inductor", "none"
+    comms_bucket_reorder_strategy: str = "aten"


If it is shared by 2 different models/experiments, I think it is okay to add it to the core job_config. This will be used by full dtensor and compiler toolki iiuc.

fegin · 2025-11-18T21:53:03Z

torchtitan/components/optimizer.py

+                blocks = model_part.layers.values()
+            else:
+                # TODO: fix autoparallel to preserve the module dict
+                blocks = model_part.layers.children()


If this is solely used by autoparallel, this should be moved to experiments/auto_parallel.

ruisizhang123 · 2025-11-18T22:12:36Z

torchtitan/experiments/auto_parallel/README.md

+
+`CONFIG_FILE="./torchtitan/models/llama3/train_configs/debug_model.toml" ./run_train.sh --model.name llama3_auto_parallel --parallelism.tensor_parallel_degree 4`
+
+Use simplefsdp's autobucketing pass:


I feel like it's quite confusing to put simplefsdp's inductor autobucketing pass here. Should we considering moving part of inductor bucketing utilities to pytorch and the algorithm to simplefsdp folder?

cc. @tianyu-l @wconstab

ruisizhang123 · 2025-11-18T22:13:45Z

torchtitan/models/deepseek_v3/model/model.py

        self.attention_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
        self.ffn_norm = nn.RMSNorm(model_args.dim, eps=model_args.norm_eps)
-
        self.moe_enabled = layer_id >= model_args.n_dense_layers


nit: add empty line?

xmfan · 2025-11-19T00:09:41Z

torchtitan/experiments/auto_parallel/deepseek_v3/parallelize_deepseekv3.py

+    return functional_router_forward
+
+
+def _moe_forward(


I intentionally keep these comments around, because it's harder to tell the delta otherwise

wconstab and others added 30 commits July 11, 2025 12:46

Autoparallel support for DP-only, DP+TP, or TP-only

e6d2caf

lets existing torchtitan knobs which govern DP/TP mesh creation and mesh size influence the sharding constraints of autoparallel, allowing it to support these different sharding configurations.

Update CLI inductor configs for bucketing/reordering

68476b3

add back llama3_autoparallel_init_fn

9ee9f75

Track API change from new AOTAutograd interface

f6e4099

Signed-off-by: Edward Z. Yang <[email protected]>

Support forcing the model into bf16 for perf debugging

4d7ee8a

Integrate MixedPrecision with AutoParallel and fix example_inputs

b801d0b

Use in-place compile API

b099cf9

Signed-off-by: Edward Z. Yang <[email protected]>

Fix bucketing pass configs

b3587d9

- fix passing of "none" (not None) to control bucketing passes

Support both eager and autoparallel init based on model.name

42c2c07

Remove llama3 init weights hack

d93845e

since meta-pytorch/autoparallel#52 landed

Support new compile API from autoparallel PR #77

6c782eb

Fix bucket sizes for AutoParallel 1D (#1545)

4712163

This PR makes bucket sizes for all-gather and reduce-scatter to be of the same size for 1d FSDP.

Add support for loss parallel (#1546)

3f04d22

IMO we should just add the loss in the model and let autoparallel parallelize it for us. But for now, let's follow how the other models are implemented

Add config for running simple-fsdp bucketing/reordering passes

8e50870

just add `--experimental.enable_simplefsdp_passes` and do not try to combine it with other `bucket_*` or `reorder_*` options.

[dsv3] patch graph break fix, works up until sharding rules

1233902

update simplefsdp pass config

4f8677b

[dsv3] disable MoE while we fix local_map, works up until optimizer

714cc5b

tweak ds3 model.py to reflect main branch for DS3 baseline can run (#…

bfa9f7f

…1684)

add simplefsdp's autobucketing pass entry (#1658)

75fb2eb

as titled, this pr adds entry to simplefsdp's autobucketing pass in autoparallel. original code is in: pytorch/pytorch#160282 The main code for autobucketing pass will be added to autoparallel repo.

[dsv3] 1D AP w/ local_map

8769396

[dsv3] Turn off Flex for AP

db22479

Merge branch 'main' into autoparallel

87ef4e0

Update to new model registration API

9dc0bd8

Whc/knobs (#1994)

c6e25bd

needs to merge in lock step with meta-pytorch/autoparallel#233

Merge remote-tracking branch 'origin/main' into autoparallel

26410e8

lint

e6ea814

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

wconstab reviewed Nov 18, 2025

View reviewed changes

xmfan commented Nov 18, 2025

View reviewed changes

wconstab reviewed Nov 18, 2025

View reviewed changes

xmfan added 2 commits November 18, 2025 11:50

undo moe patching

7abede8

move inductor config into experiment folders

d2e76b7

fegin reviewed Nov 18, 2025

View reviewed changes

ruisizhang123 reviewed Nov 18, 2025

View reviewed changes

fix local_map moe patch

472b4ad

xmfan commented Nov 19, 2025

View reviewed changes

xmfan added 5 commits November 18, 2025 16:23

move flex disables into experiment folder

ac0def9

fix newline

a24ef07

no longer necessary train.py changes

da611e4

restore comment

6cc8caa

temporarily extend hacky optimizer stuff to make dsv3 ap 1d run again

d54a6d4


		`CONFIG_FILE="./torchtitan/models/llama3/train_configs/debug_model.toml" ./run_train.sh --model.name llama3_auto_parallel --parallelism.tensor_parallel_degree 4`

		Use simplefsdp's autobucketing pass:

		)


		def _moe_forward(

Autoparallel as an experiment into main #2054

Are you sure you want to change the base?

Autoparallel as an experiment into main #2054

Uh oh!

Conversation

xmfan commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

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!

tianyu-l Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xmfan Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

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

xmfan commented Nov 18, 2025 •

edited

Loading

tianyu-l Nov 18, 2025 •

edited

Loading

ruisizhang123 Nov 19, 2025 •

edited

Loading

xmfan Nov 18, 2025 •

edited

Loading