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.
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 callscopy_gradients_to_partial_buffer().experimental/parameter_group.py::copy_gradients_to_partial_buffer()then copies eachunsharded.gradinto 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.