Bugfix in Sequence Parallelism (#178)

Fix bug in Sequence Parallelism Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Bugfix in Sequence Parallelism (#178)
Fix bug in Sequence Parallelism Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
30212170 · Kirthi Shankar Sivamani · GitHub · 0792ded4 · 30212170
Unverified Commit 30212170 authored Apr 27, 2023 by Kirthi Shankar Sivamani Committed by GitHub Apr 27, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

transformer_engine/pytorch/module.py transformer_engine/pytorch/module.py +1 -1

No files found.
--- a/transformer_engine/pytorch/module.py
+++ b/transformer_engine/pytorch/module.py
@@ -1761,7 +1761,7 @@ class _Linear(torch.autograd.Function):
                fp8_meta["scaling_fwd"].scale_inv,
                tex.FP8FwdTensors.GEMM1_WEIGHT,
                fp8_dtype_forward,
-                inputmat,
+                inputmat_total,
                fp8_meta["scaling_fwd"].scale_inv,
                tex.FP8FwdTensors.GEMM1_INPUT,
                fp8_dtype_forward,