Leverage TE wgrad accumulation fusion in MFSDP v2
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
MFSDP v2 currently materializes local wgrads, then copies them into the reduce-scatter input buffer.
In `experimental/module.py::_reduce_gradient_groups()`, MFSDP v2 allocates the partial-gradient buffer and calls `copy_gradients_to_partial_buffer()`. `experimental/parameter_group.py::copy_gradients_to_partial_buffer()` then copies each `unsharded.grad` into the corresponding partial-buffer view before reduce-scatter.
Transformer Engine wgrad accumulation fusion can write wgrads directly into a preallocated destination buffer. Using the MFSDP v2 reduce-scatter input views as that destination should eliminate this extra full-gradient copy, as well as the separate transient autograd-gradient buffer.
TE enables this through its wgrad-accumulation option and a preallocated `main_grad`-style destination.
**Done when:** supported TE linear layers write directly to MFSDP v2’s reduce-scatter input buffer; gradients match the existing path; and a benchmark reports memory and step-time impact.
Contributor guide
Research direction
Read experimental/module.py::_reduce_gradient_groups() and experimental/parameter_group.py::copy_gradients_to_partial_buffer() first, then trace the supported Transformer Engine linear-layer path and its wgrad-accumulation option. Verify gradients against the existing path and run a benchmark reporting memory and step-time impact; done means supported TE layers write directly to the MFSDP v2 reduce-scatter input buffer.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100