[PyTorch] Preserve FP8 recompute state for inner autocast

tests/pytorch/test_numerics.py

@@ -648,7 +648,15 @@ def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, recipe, fp8_m
 
 
 def _test_e2e_full_recompute(
-    bs, dtype, config, fp8, recipe, fp8_model_params=False, recompute=False, use_reentrant=True
+    bs,
+    dtype,
+    config,
+    fp8,
+    recipe,
+    fp8_model_params=False,
+    recompute=False,
+    use_reentrant=True,
+    inner_autocast=False,
 ):
     reset_rng_states()
     FP8GlobalStateManager.reset()
@@ -685,10 +693,17 @@ def _test_e2e_full_recompute(
         te_inp_hidden_states.retain_grad()
     te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q)
 
-    with autocast(enabled=fp8, recipe=recipe):
+    forward = block
+    if inner_autocast:
+
+        def forward(*args, **kwargs):
+            with autocast(enabled=fp8, recipe=recipe):
+                return block(*args, **kwargs)
+
+    with autocast(enabled=fp8 and not inner_autocast, recipe=recipe):
         if recompute:
             te_out = te_checkpoint(
-                block,
+                forward,
                 te_inp_hidden_states,
                 attention_mask=te_inp_attn_mask,
                 checkpoint_core_attention=False,
@@ -697,7 +712,7 @@ def _test_e2e_full_recompute(
                 use_reentrant=use_reentrant,
             )
         else:
-            te_out = block(
+            te_out = forward(
                 te_inp_hidden_states,
                 attention_mask=te_inp_attn_mask,
                 checkpoint_core_attention=False,
@@ -787,6 +802,48 @@ def test_gpt_full_activation_recompute(
         )
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkeypatch):
+    """Check recompute numerics when FP8 autocast starts inside the checkpointed callable."""
+    if not use_reentrant:
+        # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion.
+        monkeypatch.setenv("NVTE_BIAS_GELU_NVFUSION", "0")
+
+    dtype = torch.bfloat16
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+    config = model_configs["126m"]
+
+    outputs, names = _test_e2e_full_recompute(
+        1,
+        dtype,
+        config,
+        True,
+        fp8_recipe,
+        recompute=False,
+        use_reentrant=use_reentrant,
+    )
+    outputs_recompute, _ = _test_e2e_full_recompute(
+        1,
+        dtype,
+        config,
+        True,
+        fp8_recipe,
+        recompute=True,
+        use_reentrant=use_reentrant,
+        inner_autocast=True,
+    )
+
+    for name, ref, test in zip(names, outputs, outputs_recompute):
+        torch.testing.assert_close(
+            test,
+            ref,
+            msg=f"Mismatch in tensor {name}",
+            rtol=0.125,
+            atol=0.0675,
+        )
+
+
 def _test_e2e_checkpointing_get_model(config, dtype):
     sigma = 0.023
     init_method = init_method_normal(sigma)

transformer_engine/pytorch/distributed.py

@@ -256,9 +256,12 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F
 
     def __enter__(self):
         global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        _FP8_ACTIVATION_RECOMPUTE_ENABLED = (
-            self.activation_recompute and FP8GlobalStateManager.is_fp8_enabled()
-        )
+        # Track the checkpoint region independently of the FP8 state at entry.
+        # A checkpointed callable may open its own FP8 autocast context (for
+        # example, to select precision per layer). Delayed-scaling modules in
+        # that inner context must still save their scale and amax metadata for
+        # the recompute forward.
+        _FP8_ACTIVATION_RECOMPUTE_ENABLED = self.activation_recompute
         _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase
 
         qstate = FP8GlobalStateManager.quantization_state
@@ -275,7 +278,7 @@ def __exit__(self, *exc_details):
 
 def is_fp8_activation_recompute_enabled() -> bool:
     """Return global boolean"""
-    return _FP8_ACTIVATION_RECOMPUTE_ENABLED
+    return _FP8_ACTIVATION_RECOMPUTE_ENABLED and FP8GlobalStateManager.is_fp8_enabled()
 
 
 def in_fp8_activation_recompute_phase() -> bool: