bug fix for using return_layernorm_output=True (#1382)

the current implementation would release the output of ln, leading to an error if setting `return_layernorm_output=True`. Signed-off-by: Liyuan Liu <[email protected]> Co-authored-by: Tim Moon <[email protected]>
NVIDIA · Jan 8, 2025 · 61cf102 · 61cf102
1 parent b898cbe
commit 61cf102
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -373,7 +373,7 @@ def forward(
                 ub=ub_obj_lnout if ub_overlap_ag else None,
                 extra_output_tensor=ln_out if ub_overlap_ag else None,
             )
-            if not is_grad_enabled:
+            if not is_grad_enabled and not return_layernorm_output:
                 clear_tensor_data(ln_out_total)
 
             if bias_gelu_nvfusion: