diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 39d3e79e62..419d4e8e83 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -62,6 +62,7 @@ if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then python3 $TE_PATH/tests/pytorch/test_checkpoint.py --save-checkpoint all || error_exit "Failed to generate checkpoint files" fi python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py || test_fail "test_checkpoint.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fp8_activation_recompute.xml $TE_PATH/tests/pytorch/test_fp8_activation_recompute.py || test_fail "test_fp8_activation_recompute.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_router.xml $TE_PATH/tests/pytorch/test_fused_router.py || test_fail "test_fused_router.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml $TE_PATH/tests/pytorch/test_partial_cast.py || test_fail "test_partial_cast.py" # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision diff --git a/tests/pytorch/test_fp8_activation_recompute.py b/tests/pytorch/test_fp8_activation_recompute.py new file mode 100644 index 0000000000..887a9de6f6 --- /dev/null +++ b/tests/pytorch/test_fp8_activation_recompute.py @@ -0,0 +1,99 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe +from transformer_engine.pytorch import Linear, autocast, checkpoint +from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) + + +def _make_input(): + return torch.randn( + 16, + 16, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + +def _assert_finite_loss_and_grads(loss, inp, *layers): + assert torch.isfinite(loss) + assert inp.grad is not None + assert torch.isfinite(inp.grad).all() + for layer in layers: + assert layer.weight.grad is not None + assert torch.isfinite(layer.weight.grad).all() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", [True, False]) +def test_fp8_checkpoint_with_inner_autocast(use_reentrant): + """Delayed-scaling metadata is preserved when FP8 starts inside the checkpoint.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + inp = _make_input() + + def checkpointed_body(value): + with autocast(enabled=True, recipe=fp8_recipe): + return layer(value) + + with torch.autocast("cuda", dtype=torch.bfloat16): + out = checkpoint(checkpointed_body, inp, use_reentrant=use_reentrant) + loss = out.float().sum() + loss.backward() + torch.cuda.synchronize() + + _assert_finite_loss_and_grads(loss, inp, layer) + assert "global_fp8_buffer_pos_fwd_recompute" in layer.fp8_meta + + +@pytest.mark.parametrize("use_reentrant", [True, False]) +def test_checkpoint_without_fp8_does_not_save_fp8_recompute_state(use_reentrant): + """A checkpointed non-FP8 module does not save FP8 recompute metadata.""" + FP8GlobalStateManager.reset() + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + inp = _make_input() + + with torch.autocast("cuda", dtype=torch.bfloat16): + out = checkpoint(layer, inp, use_reentrant=use_reentrant) + loss = out.float().sum() + loss.backward() + torch.cuda.synchronize() + + _assert_finite_loss_and_grads(loss, inp, layer) + assert "global_fp8_buffer_pos_fwd_recompute" not in layer.fp8_meta + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", [True, False]) +def test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state(use_reentrant): + """Only the inner FP8 region of a mixed checkpoint saves recompute metadata.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + non_fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + inp = _make_input() + + def checkpointed_body(value): + value = non_fp8_layer(value) + with autocast(enabled=True, recipe=fp8_recipe): + return fp8_layer(value) + + with torch.autocast("cuda", dtype=torch.bfloat16): + out = checkpoint(checkpointed_body, inp, use_reentrant=use_reentrant) + loss = out.float().sum() + loss.backward() + torch.cuda.synchronize() + + _assert_finite_loss_and_grads(loss, inp, non_fp8_layer, fp8_layer) + assert "global_fp8_buffer_pos_fwd_recompute" not in non_fp8_layer.fp8_meta + assert "global_fp8_buffer_pos_fwd_recompute" in fp8_layer.fp8_meta diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d1525b53f0..89829fe6d4 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/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: