Skip to content

Categorize autograd engine, DDP reducer, and gradient helpers as BACKWARD#18

Closed
Chenguang-Zhu wants to merge 1 commit into
facebookresearch:mainfrom
Chenguang-Zhu:export-D103992607
Closed

Categorize autograd engine, DDP reducer, and gradient helpers as BACKWARD#18
Chenguang-Zhu wants to merge 1 commit into
facebookresearch:mainfrom
Chenguang-Zhu:export-D103992607

Conversation

@Chenguang-Zhu

Copy link
Copy Markdown
Member

Summary:
Three new BACKWARD rules in the categorization rule table, all using prefix matching against the demangled C++ symbol form (real frame names from PyTorch memory snapshots carry parameter signatures, so exact-name matching silently misses every real-world frame):

  1. C++ autograd engine / graph-tasktorch::autograd::Engine::, python::PythonEngine::, Node::operator(), PyNode::, CppNode<, GraphTask::, AccumulateGrad::, deleteNode, utils::LambdaPostHook. Account for a large share of allocations the old categorizer logged as UNKNOWN.

  2. DDP reducer bookkeepingc10d::Reducer::mark_variable_ready (also covers _dense overload), mark_bucket_ready, autograd_hook, initialize_buckets, initialize_local_used_map, rebuild_buckets, set_static_graph. Allreduce comm is categorized as NET separately.

  3. Grad-scale helpers_instantiate_filtered_grads (Distributed Shampoo's post-backward grad materialization).

Differential Revision: D103992607

…WARD

Summary:
Three new BACKWARD rules in the categorization rule table, all using prefix matching against the demangled C++ symbol form (real frame names from PyTorch memory snapshots carry parameter signatures, so exact-name matching silently misses every real-world frame):

1. **C++ autograd engine / graph-task** — `torch::autograd::Engine::`, `python::PythonEngine::`, `Node::operator()`, `PyNode::`, `CppNode<`, `GraphTask::`, `AccumulateGrad::`, `deleteNode`, `utils::LambdaPostHook`. Account for a large share of allocations the old categorizer logged as UNKNOWN.

2. **DDP reducer bookkeeping** — `c10d::Reducer::mark_variable_ready` (also covers `_dense` overload), `mark_bucket_ready`, `autograd_hook`, `initialize_buckets`, `initialize_local_used_map`, `rebuild_buckets`, `set_static_graph`. Allreduce comm is categorized as NET separately.

3. **Grad-scale helpers** — `_instantiate_filtered_grads` (Distributed Shampoo's post-backward grad materialization).

Differential Revision: D103992607
@meta-cla meta-cla Bot added the cla signed label May 8, 2026
@meta-codesync

meta-codesync Bot commented May 8, 2026

Copy link
Copy Markdown

@Chenguang-Zhu has exported this pull request. If you are a Meta employee, you can view the originating Diff in D103992607.

@meta-codesync

meta-codesync Bot commented May 8, 2026

Copy link
Copy Markdown

This pull request has been merged in 3155b30.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant