ub_overlap_ag_fprop disabled during activation recomputation due to is_grad_enabled check

Describe the bug

When activation recomputation is enabled, the forward pass is re-executed under torch.no_grad(). This causes is_grad_enabled to be False, which unconditionally disables the Userbuffers all-gather overlap (ub_overlap_ag_fprop) during the recompute forward pass.Relevant code at layernorm_linear.py#L195:

ub_overlap_ag_fprop = (
    ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output
)

The intent of checking is_grad_enabled here appears to be: "only overlap communication in forward if we will later need the backward pass." However, during activation recomputation, the forward pass is part of backward (it reconstructs activations needed for gradient computation). The recompute forward still performs the same GEMM + all-gather operations, so disabling comm overlap here causes a significant performance regression — the all-gather and GEMM become serialized instead of overlapped.This means enabling activation recomputation (e.g., via Megatron-LM's --recompute-activations) silently degrades TP communication overlap performance, which is counterproductive since recomputation is typically used in memory-constrained scenarios where maximizing overlap is critical.Steps/Code to reproduce bug

import torch
import transformer_engine.pytorch as te

# Minimal setup with TP overlap enabled
# (requires multi-GPU with userbuffers configured)

model = te.LayerNormLinear(
    hidden_size, ffn_hidden_size,
    params_dtype=torch.bfloat16,
    # ... with ub_overlap_ag enabled via env/config
)

# Case 1: Normal forward — ub_overlap_ag_fprop is active
with torch.enable_grad():
    out = model(inp)  # comm overlap works ✓

# Case 2: Recompute forward — ub_overlap_ag_fprop is disabled
with torch.no_grad():
    out = model(inp)  # comm overlap disabled ✗ (is_grad_enabled=False)

In practice, this is triggered by Megatron-LM's activation recomputation via CheckpointFunction, which calls the forward pass inside torch.no_grad().Expected behaviorub_overlap_ag_fprop should remain active during the recompute forward pass. The comm+GEMM overlap is a pure performance optimization for the forward computation itself — it does not depend on whether gradients are being tracked.A possible fix:

# Option 1: Remove the is_grad_enabled check entirely for fprop overlap
ub_overlap_ag_fprop = (
    ub_overlap_ag_fprop and not return_layernorm_output
)

# Option 2: Also allow overlap during recompute phase
ub_overlap_ag_fprop = (
    ub_overlap_ag_fprop
    and (is_grad_enabled or in_fp8_activation_recompute_phase())
    and not return_layernorm_output
)

Start at transformer_engine/pytorch/module/layernorm_linear.py#L195 and trace how ub_overlap_ag_fprop is handled during the CheckpointFunction recompute path. Reproduce with a multi-GPU setup using userbuffers and activation recomputation, then verify that all-gather and GEMM overlap remains enabled during recompute without changing normal forward behavior.