NVFP4 and MXFP8 integrations

- includes capability to log out stats for MXFP8
and NVFP4 at the same time
- Enables layer-wise precision setting
- Includes support for layer-wise quant config
- Adds support for FSDP2 DDP and MFSDP for MXFP8 NVFP4
- note: FP32 master weights wont work with ddp,mfsdp

Signed-off-by: Jonathan Mitchell <jomitchell@nvidia.com>

Author: Jonathan Mitchell

.vscode/settings.json

@@ -26,5 +26,6 @@
     "editor.rulers": [
         120
     ],
-    "autoDocstring.docstringFormat": "google-notypes"
+    "autoDocstring.docstringFormat": "google-notypes",
+    "search.exclude": { "**/logs/**": true },
 }

bionemo-recipes/models/esm2/modeling_esm_te.py

@@ -23,11 +23,13 @@
 """
 
 import warnings
+from contextlib import nullcontext
 from typing import ClassVar, Literal, Optional, Unpack
 
 # TODO: put import guard around transformer_engine here, with an informative error message around
 # installation and the nvidia docker container.
 import torch
+import transformer_engine.common.recipe
 import transformer_engine.pytorch
 from torch import nn
 from torch.nn import CrossEntropyLoss
@@ -71,6 +73,7 @@ def __init__(
         max_seq_length: Optional[int] = None,
         padded_vocab_size: Optional[int] = 64,
         attn_mask_type: str = "padding",
+        layer_precision: list[str | None] | None = None,
         **kwargs,
     ):
         """Initialize the NVEsmConfig with additional TE-related config options.
@@ -100,6 +103,9 @@ def __init__(
             padded_vocab_size: The padded vocabulary size to support FP8. If not provided, defaults
                 to vocab_size. Must be greater than or equal to vocab_size.
             attn_mask_type: The type of attention mask to use.
+            layer_precision: Per-layer quantization precision, a list of length ``num_hidden_layers``
+                where each element is ``"fp8"``, ``"fp4"``, or ``None`` (BF16 fallback). ``None``
+                (the default) means no quantization is configured.
             **kwargs: Additional config options to pass to EsmConfig.
         """
         super().__init__(**kwargs)
@@ -111,6 +117,7 @@ def __init__(
         self.micro_batch_size = micro_batch_size
         self.max_seq_length = max_seq_length
         self.attn_mask_type = attn_mask_type
+        self.layer_precision = layer_precision
 
         # Set padded_vocab_size with default fallback to vocab_size
         self.padded_vocab_size = padded_vocab_size if padded_vocab_size is not None else self.vocab_size
@@ -165,6 +172,8 @@ def _init_method(x):
                 for i in range(config.num_hidden_layers)
             ]
         )
+        self._fp8_recipe: transformer_engine.common.recipe.Recipe | None = None
+        self._fp4_recipe: transformer_engine.common.recipe.Recipe | None = None
         self.emb_layer_norm_after = transformer_engine.pytorch.LayerNorm(
             config.hidden_size,
             eps=config.layer_norm_eps,
@@ -174,6 +183,49 @@ def _init_method(x):
         if config.position_embedding_type == "rotary":
             self.rotary_embeddings = RotaryPositionEmbedding(config.hidden_size // config.num_attention_heads)
 
+    def set_recipes(
+        self,
+        fp8_recipe: transformer_engine.common.recipe.Recipe | None = None,
+        fp4_recipe: transformer_engine.common.recipe.Recipe | None = None,
+    ) -> None:
+        """Attach quantization recipe objects for per-layer autocast.
+
+        Recipes are not serializable and must be set at runtime after model creation
+        and sharding (FSDP/DDP/mFSDP) but before training. The per-layer precision
+        assignments are read from ``self.config.layer_precision``.
+
+        These recipes are also hardware specific, so we should not store them as
+        attributes of the model and attach them at runtime.
+
+        Args:
+            fp8_recipe: The FP8 recipe instance (e.g., MXFP8BlockScaling), or None.
+            fp4_recipe: The FP4 recipe instance (e.g., NVFP4BlockScaling), or None.
+        """
+        self._fp8_recipe = fp8_recipe
+        self._fp4_recipe = fp4_recipe
+
+    def get_layer_autocast(self, layer_number: int):
+        """Return the appropriate TE autocast context manager for a given layer.
+
+        The context interacts with the outer FP8 autocast in the training script:
+        - FP8 layer: nullcontext() -- lets the outer FP8 autocast take effect.
+        - FP4 layer: te.pytorch.autocast(enabled=True, recipe=fp4_recipe) -- overrides to FP4.
+        - BF16 layer: te.pytorch.autocast(enabled=False) -- disables quantized compute.
+
+        Args:
+            layer_number: The 0-indexed layer number.
+
+        Returns:
+            A context manager for the layer's quantization mode.
+        """
+        precision = self.config.layer_precision[layer_number] if self.config.layer_precision is not None else None
+        if precision == "fp8":
+            return nullcontext()
+        elif precision == "fp4":
+            return transformer_engine.pytorch.autocast(enabled=True, recipe=self._fp4_recipe)
+        else:
+            return transformer_engine.pytorch.autocast(enabled=False)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -201,22 +253,26 @@ def forward(
             if te_rope_emb.dtype == torch.float32:
                 warnings.warn("Rotary embeddings should be in float32 for optimal performance.", UserWarning)
 
-        for layer_module in self.layers:
-            if kwargs.get("output_hidden_states", False):
-                all_hidden_states = (*all_hidden_states, hidden_states)
-
-            hidden_states = layer_module(
-                hidden_states,
-                attention_mask,
-                rotary_pos_emb=te_rope_emb,
-                cu_seqlens_q=kwargs.get("cu_seq_lens_q", None),
-                cu_seqlens_kv=kwargs.get("cu_seq_lens_k", None),
-                cu_seqlens_q_padded=kwargs.get("cu_seq_lens_q_padded", None),
-                cu_seqlens_kv_padded=kwargs.get("cu_seq_lens_k_padded", None),
-                max_seqlen_q=kwargs.get("max_length_q", None),
-                max_seqlen_kv=kwargs.get("max_length_k", None),
-                pad_between_seqs=kwargs.get("pad_between_seqs", None),
-            )
+        # Outer FP8 autocast enables FP8 compute for the encoder stack. Per-layer overrides (FP4, BF16) are handled
+        # by get_layer_autocast(), which nests inside this context.
+        with transformer_engine.pytorch.autocast(enabled=self._fp8_recipe is not None, recipe=self._fp8_recipe):
+            for layer_number, layer_module in enumerate(self.layers):
+                if kwargs.get("output_hidden_states", False):
+                    all_hidden_states = (*all_hidden_states, hidden_states)
+
+                with self.get_layer_autocast(layer_number):
+                    hidden_states = layer_module(
+                        hidden_states,
+                        attention_mask,
+                        rotary_pos_emb=te_rope_emb,
+                        cu_seqlens_q=kwargs.get("cu_seq_lens_q", None),
+                        cu_seqlens_kv=kwargs.get("cu_seq_lens_k", None),
+                        cu_seqlens_q_padded=kwargs.get("cu_seq_lens_q_padded", None),
+                        cu_seqlens_kv_padded=kwargs.get("cu_seq_lens_k_padded", None),
+                        max_seqlen_q=kwargs.get("max_length_q", None),
+                        max_seqlen_kv=kwargs.get("max_length_k", None),
+                        pad_between_seqs=kwargs.get("pad_between_seqs", None),
+                    )
 
         hidden_states = self.emb_layer_norm_after(hidden_states)
 

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

@@ -161,7 +161,10 @@ def is_main_process(self) -> bool:
     )
     device = f"cuda:{dist_config.local_rank}"
 
+    fp8_recipe = DelayedScaling(fp8_format=Format.HYBRID, amax_compute_algo="max", amax_history_len=10)
+
     config = NVEsmConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", dtype=torch.bfloat16, revision="c731040f")
+    config.layer_precision = ["fp8"] * config.num_hidden_layers
     model = NVEsmForMaskedLM(config)
 
     if args.strategy is Strategy.FSDP2:
@@ -195,13 +198,15 @@ def is_main_process(self) -> bool:
             tp_dim="tp",
         )
 
+    # Attach FP8 recipes to the encoder (layer precision is already on config).
+    encoder = model.module.esm.encoder if args.strategy is Strategy.DDP else model.esm.encoder
+    encoder.set_recipes(fp8_recipe=fp8_recipe, fp4_recipe=None)
+
     model.train()
 
     generator = torch.Generator()
     generator.manual_seed(torch.distributed.get_rank())
 
-    fp8_recipe = DelayedScaling(fp8_format=Format.HYBRID, amax_compute_algo="max", amax_history_len=10)
-
     for _ in range(3):
         input_data = {
             "input_ids": torch.randint(0, config.vocab_size, (1, 32), generator=generator),
@@ -211,8 +216,7 @@ def is_main_process(self) -> bool:
         input_data = {k: v.to(torch.cuda.current_device()) for k, v in input_data.items()}
 
         with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-            with transformer_engine.pytorch.autocast(enabled=True, recipe=fp8_recipe):
-                outputs = model(**input_data)
+            outputs = model(**input_data)
 
         outputs.loss.backward()
 

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

@@ -0,0 +1,179 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-Apache2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for NVEsmEncoder.set_recipes and get_layer_autocast."""
+
+from contextlib import nullcontext
+from unittest.mock import patch
+
+import pytest
+import transformer_engine.common.recipe
+import transformer_engine.pytorch
+
+from modeling_esm_te import NVEsmConfig, NVEsmEncoder
+
+
+@pytest.fixture
+def encoder():
+    """Create a small NVEsmEncoder on CUDA for testing."""
+    config = NVEsmConfig(
+        hidden_size=320,
+        intermediate_size=1280,
+        num_hidden_layers=6,
+        num_attention_heads=20,
+        max_position_embeddings=1026,
+    )
+    return NVEsmEncoder(config)
+
+
+# -- set_recipes --
+
+
+def test_all_fp8(encoder):
+    encoder.config.layer_precision = ["fp8"] * 6
+    fp8_recipe = transformer_engine.common.recipe.DelayedScaling()
+    encoder.set_recipes(fp8_recipe=fp8_recipe, fp4_recipe=None)
+    assert encoder._fp8_recipe is fp8_recipe
+    assert encoder._fp4_recipe is None
+    assert all(p == "fp8" for p in encoder.config.layer_precision)
+
+
+def test_all_fp4(encoder):
+    encoder.config.layer_precision = ["fp4"] * 6
+    fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+    encoder.set_recipes(fp8_recipe=None, fp4_recipe=fp4_recipe)
+    assert encoder._fp8_recipe is None
+    assert encoder._fp4_recipe is fp4_recipe
+    assert all(p == "fp4" for p in encoder.config.layer_precision)
+
+
+def test_all_bf16(encoder):
+    encoder.config.layer_precision = [None] * 6
+    encoder.set_recipes(fp8_recipe=None, fp4_recipe=None)
+    assert all(p is None for p in encoder.config.layer_precision)
+
+
+def test_mixed_fp8_fp4(encoder):
+    encoder.config.layer_precision = ["fp8", "fp8", "fp8", "fp4", "fp4", "fp4"]
+    fp8_recipe = transformer_engine.common.recipe.DelayedScaling()
+    fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+    encoder.set_recipes(fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe)
+    assert encoder.config.layer_precision == ["fp8", "fp8", "fp8", "fp4", "fp4", "fp4"]
+
+
+def test_mixed_fp8_bf16(encoder):
+    encoder.config.layer_precision = ["fp8", None, "fp8", None, "fp8", None]
+    fp8_recipe = transformer_engine.common.recipe.DelayedScaling()
+    encoder.set_recipes(fp8_recipe=fp8_recipe, fp4_recipe=None)
+    assert encoder.config.layer_precision == ["fp8", None, "fp8", None, "fp8", None]
+
+
+def test_mixed_all_three(encoder):
+    encoder.config.layer_precision = ["fp8", "fp8", None, None, "fp4", "fp4"]
+    fp8_recipe = transformer_engine.common.recipe.DelayedScaling()
+    fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+    encoder.set_recipes(fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe)
+    assert encoder.config.layer_precision == ["fp8", "fp8", None, None, "fp4", "fp4"]
+
+
+def test_covers_all_layers(encoder):
+    encoder.config.layer_precision = ["fp8"] + [None] * 5
+    encoder.set_recipes(fp8_recipe=transformer_engine.common.recipe.DelayedScaling(), fp4_recipe=None)
+    assert len(encoder.config.layer_precision) == 6
+
+
+def test_recipes_stored_as_attributes(encoder):
+    encoder.config.layer_precision = ["fp8", "fp4", None, None, None, None]
+    fp8_recipe = transformer_engine.common.recipe.DelayedScaling()
+    fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+    encoder.set_recipes(fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe)
+    assert encoder._fp8_recipe is fp8_recipe
+    assert encoder._fp4_recipe is fp4_recipe
+    # The precision list only contains strings/None, not recipe objects.
+    for v in encoder.config.layer_precision:
+        assert v is None or isinstance(v, str)
+
+
+# -- get_layer_autocast --
+
+
+def test_fp8_layer_returns_nullcontext(encoder):
+    encoder.config.layer_precision = ["fp8"] + [None] * 5
+    encoder.set_recipes(fp8_recipe=transformer_engine.common.recipe.DelayedScaling(), fp4_recipe=None)
+    ctx = encoder.get_layer_autocast(0)
+    assert isinstance(ctx, nullcontext)
+
+
+def test_fp4_layer_returns_te_autocast(encoder):
+    fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+    encoder.config.layer_precision = ["fp4"] + [None] * 5
+    encoder.set_recipes(fp8_recipe=None, fp4_recipe=fp4_recipe)
+    with patch.object(transformer_engine.pytorch, "autocast") as mock_autocast:
+        mock_autocast.return_value = "fp4_context"
+        ctx = encoder.get_layer_autocast(0)
+        mock_autocast.assert_called_once_with(enabled=True, recipe=fp4_recipe)
+        assert ctx == "fp4_context"
+
+
+def test_bf16_layer_returns_te_autocast_disabled(encoder):
+    encoder.config.layer_precision = [None] * 6
+    encoder.set_recipes(fp8_recipe=None, fp4_recipe=None)
+    with patch.object(transformer_engine.pytorch, "autocast") as mock_autocast:
+        mock_autocast.return_value = "bf16_context"
+        ctx = encoder.get_layer_autocast(0)
+        mock_autocast.assert_called_once_with(enabled=False)
+        assert ctx == "bf16_context"
+
+
+def test_uninitialized_defaults_to_bf16(encoder):
+    """When layer_precision is None (default), all layers default to BF16."""
+    assert encoder.config.layer_precision is None
+    with patch.object(transformer_engine.pytorch, "autocast") as mock_autocast:
+        mock_autocast.return_value = "bf16_context"
+        ctx = encoder.get_layer_autocast(0)
+        mock_autocast.assert_called_once_with(enabled=False)
+        assert ctx == "bf16_context"
+
+
+def test_mixed_layers_return_correct_contexts(encoder):
+    fp8_recipe = transformer_engine.common.recipe.DelayedScaling()
+    fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling()
+    encoder.config.layer_precision = ["fp8", "fp8", "fp4", "fp4", None, None]
+    encoder.set_recipes(fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe)
+
+    # FP8 layers -> nullcontext
+    assert isinstance(encoder.get_layer_autocast(0), nullcontext)
+    assert isinstance(encoder.get_layer_autocast(1), nullcontext)
+
+    # FP4 layers -> te.pytorch.autocast
+    with patch.object(transformer_engine.pytorch, "autocast") as mock_autocast:
+        mock_autocast.return_value = "fp4_context"
+        encoder.get_layer_autocast(2)
+        mock_autocast.assert_called_with(enabled=True, recipe=fp4_recipe)
+
+    # BF16 layers -> te.pytorch.autocast(enabled=False)
+    with patch.object(transformer_engine.pytorch, "autocast") as mock_autocast:
+        mock_autocast.return_value = "bf16_context"
+        encoder.get_layer_autocast(4)
+        mock_autocast.assert_called_with(enabled=False)
+
+
+def test_layer_precision_is_pickleable(encoder):
+    """The config.layer_precision list should be trivially pickleable."""
+    import pickle
+
+    encoder.config.layer_precision = ["fp8", "fp8", "fp4", "fp4", None, None]
+    roundtripped = pickle.loads(pickle.dumps(encoder.config.layer_precision))
+    assert roundtripped == encoder.config.layer_precision