expanded esm2 fp8 tests (#1330)

These are currently passing in
`gitlab-master.nvidia.com/dl/transformerengine/transformerengine:main-pytorch-py3-base`
but may require some additional changes for the TE that ships in 25.10

Signed-off-by: Peter St. John <[email protected]>

Author: pstjohn

.devcontainer/recipes/Dockerfile

@@ -1,3 +1,5 @@
+# Uncomment to use the latest TE from the NGC registry for debugging changes with latest TE.
+# FROM gitlab-master.nvidia.com/dl/transformerengine/transformerengine:main-pytorch-py3-base
 FROM nvcr.io/nvidia/pytorch:25.10-py3
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=requirements.txt,target=/workspace/requirements.txt \

bionemo-recipes/models/esm2/tests/test_fp8.py

@@ -18,27 +18,67 @@
 import torch.distributed.checkpoint as dcp
 import transformer_engine
 from torch.distributed.checkpoint.state_dict import get_model_state_dict
-from transformer_engine.common.recipe import DelayedScaling, MXFP8BlockScaling
-from transformer_engine.pytorch.fp8 import check_fp8_support, check_mxfp8_support
-from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
+from transformer_engine.common import recipe as recipe_module
+from transformer_engine.pytorch import fp8
+from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensor
 
 from esm.collator import MLMDataCollatorWithFlattening
 from esm.modeling_esm_te import NVEsmConfig, NVEsmForMaskedLM
 
 
-def requires_fp8(func):
-    """Decorator to skip tests that require FP8 support."""
-    fp8_available, reason = check_fp8_support()
-    return pytest.mark.skipif(not fp8_available, reason=f"FP8 is not supported on this GPU: {reason}")(func)
-
-
-def requires_mxfp8(func):
-    """Decorator to skip tests that require MXFP8 support."""
-    mxfp8_available, reason = check_mxfp8_support()
-    if torch.cuda.get_device_capability() == (12, 0):
-        mxfp8_available = False
-        reason = "MXFP8 is not supported on sm120"
-    return pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8 is not supported on this GPU: {reason}")(func)
+ALL_RECIPES = [
+    recipe_module.DelayedScaling(),
+    recipe_module.Float8CurrentScaling(),
+    recipe_module.Float8BlockScaling(),
+    recipe_module.MXFP8BlockScaling(),
+    # recipe_module.NVFP4BlockScaling(disable_rht=True, disable_stochastic_rounding=True),
+]
+
+
+def _check_recipe_support(recipe: recipe_module.Recipe):
+    """Check if a recipe is supported and return (supported, reason)."""
+    if isinstance(recipe, recipe_module.DelayedScaling):
+        recipe_supported, reason = fp8.check_fp8_support()
+    elif isinstance(recipe, recipe_module.Float8CurrentScaling):
+        recipe_supported, reason = fp8.check_fp8_support()
+    elif isinstance(recipe, recipe_module.Float8BlockScaling):
+        recipe_supported, reason = fp8.check_fp8_block_scaling_support()
+    elif isinstance(recipe, recipe_module.MXFP8BlockScaling):
+        recipe_supported, reason = fp8.check_mxfp8_support()
+    elif isinstance(recipe, recipe_module.NVFP4BlockScaling):
+        recipe_supported, reason = fp8.check_nvfp4_support()
+    else:
+        recipe_supported = False
+        reason = "Unsupported recipe"
+    return recipe_supported, reason
+
+
+def requires_recipe_support(recipe: recipe_module.Recipe):
+    """Decorator to skip tests that require recipe support."""
+
+    def requires_recipe_support_inner(func):
+        recipe_supported, reason = _check_recipe_support(recipe)
+        return pytest.mark.skipif(not recipe_supported, reason=reason)(func)
+
+    return requires_recipe_support_inner
+
+
+def parametrize_recipes_with_support(recipes):
+    """Generate pytest.param objects with skip marks for unsupported recipes."""
+    parametrized_recipes = []
+    for recipe in recipes:
+        recipe_supported, reason = _check_recipe_support(recipe)
+        parametrized_recipes.append(
+            pytest.param(
+                recipe,
+                id=recipe.__class__.__name__,
+                marks=pytest.mark.skipif(
+                    not recipe_supported,
+                    reason=reason,
+                ),
+            )
+        )
+    return parametrized_recipes
 
 
 @pytest.fixture
@@ -53,24 +93,30 @@ def input_data_thd(tokenizer, tokenized_proteins):
     return data_collator(tokenized_proteins)
 
 
-@requires_fp8
-def test_fp8_forward_and_backward_pass(te_model_checkpoint, input_data):
+@pytest.mark.parametrize("fp8_recipe", parametrize_recipes_with_support(ALL_RECIPES))
+def test_fp8_forward_and_backward_pass(te_model_checkpoint, input_data, fp8_recipe):
     model_te = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, dtype=torch.bfloat16)
     model_te.to("cuda")
 
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
     outputs = model_te(**input_data)
 
-    fp8_recipe = DelayedScaling()
     with transformer_engine.pytorch.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
         outputs_fp8 = model_te(**input_data)
     outputs_fp8.loss.backward()
 
-    torch.testing.assert_close(outputs_fp8.loss, outputs.loss)
+    if isinstance(fp8_recipe, recipe_module.NVFP4BlockScaling):
+        atol = 0.2
+        rtol = 0.05
+    else:
+        atol = None
+        rtol = None
+
+    torch.testing.assert_close(outputs_fp8.loss, outputs.loss, atol=atol, rtol=rtol)
 
 
-@requires_fp8
-def test_fp8_forward_and_backward_pass_thd(te_model_checkpoint, input_data_thd, monkeypatch):
+@pytest.mark.parametrize("fp8_recipe", parametrize_recipes_with_support(ALL_RECIPES))
+def test_fp8_forward_and_backward_pass_thd(te_model_checkpoint, input_data_thd, fp8_recipe, monkeypatch):
     if torch.cuda.get_device_capability() == (12, 0):
         # TODO(BIONEMO-2840): On sm120, we need to set NVTE_FUSED_ATTN to 0 since TE will choose fused attn by default,
         # but it's missing this THD implementation.
@@ -82,54 +128,27 @@ def test_fp8_forward_and_backward_pass_thd(te_model_checkpoint, input_data_thd,
     input_data = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in input_data_thd.items()}
     outputs = model_te(**input_data)
 
-    fp8_recipe = DelayedScaling()
     with transformer_engine.pytorch.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
         outputs_fp8 = model_te(**input_data)
     outputs_fp8.loss.backward()
 
-    torch.testing.assert_close(outputs_fp8.loss, outputs.loss)
-
-
-@requires_mxfp8
-def test_mxfp8_forward_and_backward_pass(te_model_checkpoint, input_data):
-    model_te = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, dtype=torch.bfloat16)
-    model_te.to("cuda")
-
-    input_data = {k: v.to("cuda") for k, v in input_data.items()}
-    outputs = model_te(**input_data)
-
-    mxfp8_recipe = MXFP8BlockScaling()
-    with transformer_engine.pytorch.fp8_autocast(enabled=True, fp8_recipe=mxfp8_recipe):
-        outputs_fp8 = model_te(**input_data)
-    outputs_fp8.loss.backward()
-
-    torch.testing.assert_close(outputs_fp8.loss, outputs.loss)
+    if isinstance(fp8_recipe, recipe_module.NVFP4BlockScaling):
+        atol = 0.2
+        rtol = 0.05
+    else:
+        atol = None
+        rtol = None
 
+    torch.testing.assert_close(outputs_fp8.loss, outputs.loss, atol=atol, rtol=rtol)
 
-@requires_mxfp8
-def test_mxfp8_forward_and_backward_pass_thd(te_model_checkpoint, input_data_thd):
-    model_te = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, attn_input_format="thd", dtype=torch.bfloat16)
-    model_te.to("cuda")
-
-    input_data = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in input_data_thd.items()}
-    outputs = model_te(**input_data)
 
-    mxfp8_recipe = MXFP8BlockScaling()
-    with transformer_engine.pytorch.fp8_autocast(enabled=True, fp8_recipe=mxfp8_recipe):
-        outputs_fp8 = model_te(**input_data)
-    outputs_fp8.loss.backward()
-
-    torch.testing.assert_close(outputs_fp8.loss, outputs.loss)
-
-
-@requires_fp8
-def test_fp8_model_init_forward_and_backward(te_model_checkpoint, input_data):
-    fp8_recipe = DelayedScaling()
+@pytest.mark.parametrize("fp8_recipe", parametrize_recipes_with_support(ALL_RECIPES))
+def test_fp8_model_init_forward_and_backward(te_model_checkpoint, input_data, fp8_recipe):
     config = NVEsmConfig.from_pretrained(te_model_checkpoint, dtype=torch.bfloat16)
     with transformer_engine.pytorch.fp8_model_init(enabled=True, recipe=fp8_recipe):
         model_te = NVEsmForMaskedLM(config)
 
-    assert isinstance(model_te.lm_head.dense.weight, Float8Tensor)
+    assert isinstance(model_te.lm_head.dense.weight, QuantizedTensor)
 
     model_te.to("cuda")
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
@@ -140,39 +159,34 @@ def test_fp8_model_init_forward_and_backward(te_model_checkpoint, input_data):
     outputs_fp8.loss.backward()
 
 
-@requires_fp8
 @pytest.mark.xfail(reason="BIONEMO-3055: fp8 model init and pretrained loading is not currently supported.")
-def test_fp8_model_init_from_pretrained(te_model_checkpoint, input_data):
-    fp8_recipe = DelayedScaling()
-
+@pytest.mark.parametrize("fp8_recipe", parametrize_recipes_with_support(ALL_RECIPES))
+def test_fp8_model_init_from_pretrained(te_model_checkpoint, fp8_recipe):
     # TODO: this will be renamed to quantized_model_init in the future, fp8_model_init will be removed in 3.0
     with transformer_engine.pytorch.fp8_model_init(enabled=True, recipe=fp8_recipe):
-        model_te = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint)
+        model_te = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, dtype=torch.bfloat16)
 
-    assert isinstance(model_te.esm.encoder.layers[0].layernorm_mlp.fc2_weight, Float8Tensor)
-    assert isinstance(model_te.lm_head.dense.weight, Float8Tensor)
+    assert isinstance(model_te.esm.encoder.layers[0].layernorm_mlp.fc2_weight, QuantizedTensor)
+    assert isinstance(model_te.lm_head.dense.weight, QuantizedTensor)
 
 
-@requires_fp8
 @pytest.mark.xfail(reason="BIONEMO-3055: fp8 model init and pretrained saving is not currently supported.")
-def test_fp8_model_init_save_pretrained(te_model_checkpoint, tmp_path):
-    fp8_recipe = DelayedScaling()
+@pytest.mark.parametrize("fp8_recipe", parametrize_recipes_with_support(ALL_RECIPES))
+def test_fp8_model_init_save_pretrained(te_model_checkpoint, tmp_path, fp8_recipe):
     config = NVEsmConfig.from_pretrained(te_model_checkpoint, dtype=torch.bfloat16)
     with transformer_engine.pytorch.fp8_model_init(enabled=True, recipe=fp8_recipe):
         model_fp8 = NVEsmForMaskedLM(config)
 
-    assert isinstance(model_fp8.esm.encoder.layers[0].layernorm_mlp.fc2_weight, Float8Tensor)
-    assert isinstance(model_fp8.lm_head.dense.weight, Float8Tensor)
+    assert isinstance(model_fp8.esm.encoder.layers[0].layernorm_mlp.fc2_weight, QuantizedTensor)
+    assert isinstance(model_fp8.lm_head.dense.weight, QuantizedTensor)
 
     model_fp8.save_pretrained(tmp_path / "fp8_checkpoint")
     del model_fp8
     NVEsmForMaskedLM.from_pretrained(tmp_path / "fp8_checkpoint", dtype=torch.bfloat16)
 
 
-@requires_fp8
-@pytest.mark.xfail(reason="BIONEMO-3055: fp8 model init and distributed checkpointing is not currently supported.")
-def test_fp8_model_distributed_checkpointing_save_and_load(te_model_checkpoint, tmp_path, input_data):
-    fp8_recipe = DelayedScaling()
+@pytest.mark.parametrize("fp8_recipe", parametrize_recipes_with_support(ALL_RECIPES))
+def test_fp8_model_distributed_checkpointing_save_and_load(te_model_checkpoint, tmp_path, input_data, fp8_recipe):
     config = NVEsmConfig.from_pretrained(te_model_checkpoint, dtype=torch.bfloat16)
     with transformer_engine.pytorch.fp8_model_init(enabled=True, recipe=fp8_recipe):
         model_fp8 = NVEsmForMaskedLM(config)
@@ -184,6 +198,7 @@ def test_fp8_model_distributed_checkpointing_save_and_load(te_model_checkpoint,
     outputs.loss.backward()
 
     state_dict = get_model_state_dict(model_fp8)
+    state_dict = {key: val for key, val in state_dict.items() if not key.endswith("_extra_state")}
     dcp.save(state_dict, checkpoint_id=tmp_path / "fp8_checkpoint")
 
     del model_fp8, state_dict
@@ -192,4 +207,5 @@ def test_fp8_model_distributed_checkpointing_save_and_load(te_model_checkpoint,
         model_fp8 = NVEsmForMaskedLM(config)
 
     state_dict = model_fp8.state_dict()
+    state_dict = {key: val for key, val in state_dict.items() if not key.endswith("_extra_state")}
     dcp.load(state_dict, checkpoint_id=tmp_path / "fp8_checkpoint")