Workaround AC HOP mutation issue when tracing token dispatch

[xmfan]

FIXES #1935

Stacked PRs:

• ->Workaround AC HOP mutation issue when tracing token dispatch #1984

tlparse: https://fburl.com/sqxd6c0w

---

Workaround AC HOP mutation issue when tracing token dispatch

TORCH_COMPILE_FORCE_DISABLE_CACHES=1 HF_TOKEN=<token> HF_HUB_DISABLE_XET=1 CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml" with-proxy ./run_train.sh --model.name simple_fsdp.deepseek_v3

This is a problem for SimpleFSDP where we want to fullgraph the entire model, these "mutation" cause graph break

It is less of a problem outside SimpleFSDP, because we don't currently compile token dispatch

[tianyu-l]

These shouldn't be exposed to single-device model code. Plus, I don't think it will work if EP is not used.

If it's getting too hard, maybe we should use local_map / to_local to re-implement MoE.

[ruisizhang123]

Thank you for the fix! Do you think it would require fewer user-side changes if we reimplemented apply_ac as a graph pass?

[ruisizhang123]

minor: toml file config changed.

[xmfan]

This change is only needed if you use compile(torch.utils.checkpoint(, so graph pass wouldn't need it. but if you use both eager and graph-based, you will need this again

[ruisizhang123]

This change is only needed if you use compile(torch.utils.checkpoint(, so graph pass wouldn't need it. but if you use both eager and graph-based, you will need this again

Yes, what I meant is that if we're going for a compiler-based approach to distributed parallelism in simplefsdp, it would make sense to have a specialized apply_ac function that’s also compiler-based. (and users are not allowed to use eager checkpoint to implement ac)