NVIDIA / NVIDIA/Megatron-LM

Leverage TE wgrad accumulation fusion in MFSDP v2

Open
#6,942 4 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.