meta-pytorch / meta-pytorch/autoparallel

Move AutoParallel activation checkpointing out of the compile backend

Open
#447 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
97
Forks
27
Avg merge
5d 18h
Merged PRs (30d)
3

Description

AutoParallel currently applies FSDP AC tags before compile_with_joint_descriptors, but the auto_parallel backend also runs the AC joint pass. The backend pass preserves the earlier FSDP tags and applies new tags, which can cause FSDP recompute to be represented twice in the backward graph. It may be removed by DCE, but this is a fragmented and fragile path. AC should be applied once during AutoParallel placement, before compilation, and the backend should only configure Inductor bucketing/overlap context.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Trace AutoParallel placement through compile_with_joint_descriptors and the auto_parallel backend's AC joint pass. Confirm where FSDP AC tags are applied, then verify that placement performs AC once while the backend only configures Inductor bucketing and overlap context; the duplicate backward-graph representation should no longer occur.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, compilers
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.