JointGraph-based Training Prototype

[SherlockNoMad]

This is an e2e prototype to run llama3-simplefsdp using export-y aot_autograd workflow.

Setup: shard_dp = 2, tp = 4.

MVP

• [Done] Start with a simpleFSDP model, enable TP + FSDP
• [Done] Apply aot_export_joing_with_descriptor on parallelized module with DTensor input to get the joint graph
• [Done] Apply min_cut_partitioner to get forward and backward graph module
• [Need verification] Apply prefect/bucketing graph passes on fw_gm and bw_gm to reorder/group the communication collectives
• [In Progress] Run the fw and bw graph with JointGraphRunner

Issues

• fwd_rng_state show up in the aot_export_joint grpah input pytorch#164559
• [DTensor] Improve Sharding propagation error message pytorch#164543
• What's input order for aot_export_joint graph? using model.parameter() 's order as input seems wrong.

Repro steps:
Needs to be run on top of pytorch/pytorch#164557
NGPU=8 CONFIG_FILE="./torchtitan/models/llama3/train_configs/debug_model.toml" with-proxy ./run_train.sh --model.name joint_graph_runner.llama3 --compile.enable --parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=4

Sample output:
P1975157784: rank0_autograd_function_0fea2786.py
P1975158481: rank1_autograd_function_28587623.py