fix gradient accumulation fusion for FSDP (#2371)

Signed-off-by: Lifu Zhang <lifuz@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by: Lifu Zhang <lifuz@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

fix gradient accumulation fusion for FSDP (#2371)
Signed-off-by: Lifu Zhang <lifuz@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by: Lifu Zhang <lifuz@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
d8f1e68f · Lifu Zhang · GitHub · c544ced2 · d8f1e68f · d8f1e68f
Unverified Commit d8f1e68f authored Nov 12, 2025 by Lifu Zhang Committed by GitHub Nov 12, 2025
3 changed files
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py