ESM2 NVFP4 and MXFP8 support and documentation update. (#1484)

Layer-wise MXFP8/NVFP4 precision for ESM-2 TransformerEngine training
Adds support for per-layer quantization precision control, enabling
mixed FP8/FP4/BF16
configurations across transformer layers during training. This allows
users to assign
different quantization formats to different layers via Hydra config
(1-indexed fp8_layers
and fp4_layers lists), enabling convergence/performance tradeoff
exploration.

  Key changes:

- Per-layer quantization context in encoder forward: NVEsmEncoder now
maintains a
layer_number_quantized_recipe_map that selects the appropriate TE
autocast context per
layer (nullcontext for FP8 to respect outer autocast, explicit autocast
for FP4, or
  autocast(enabled=False) for BF16).
- quantization.py: New utilities for resolving layer-wise quantization
assignments
  (resolve_quantization_layers), generating debug API regex patterns
(generate_layer_regex), and initializing nvdlfw_inspect quant stats
logging
(initialize_quant_stats_logging). Handles 0-indexed (model internals)
and 1-indexed
  (user-facing) layer numbering.
- train_ddp.py / train_fsdp2.py: Integrated layer-wise quantization
setup --
resolves layer assignments from config, builds recipe map, assigns to
encoder, and
  optionally initializes quant stats logging.
- train_fsdp2_cp.py: Switched from AutoConfig/AutoModelForMaskedLM to
local
NVEsmConfig/NVEsmForMaskedLM for consistency and to avoid remote code
trust issues.
  - Hydra config (defaults.yaml): Added fp4_config, quant_stats_config,
  fp8_layers, fp4_layers, and use_fp32_master_weights settings.
- Model files: Updated esm_nv.py across all checkpoint directories
(native_te,
accelerate, peft) and the models package with layer-wise quantization
support, NVTX
annotations per encoder layer, and FP8_RECIPES/FP4_RECIPES type
constants.
  - Tests: Added comprehensive tests for resolve_quantization_layers,
generate_layer_regex, and update_quant_stats_config covering defaults,
explicit
layers, mixed assignments, overlap validation, and edge cases.###
Description



#### Usage

<!--- How does a user interact with the changed code -->

```python
TODO: Add code snippet
```

### Type of changes

<!-- Mark the relevant option with an [x] -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Refactor
- [ ] Documentation update
- [ ] Other (please describe):

### CI Pipeline Configuration

Configure CI behavior by applying the relevant labels. By default, only
basic unit tests are run.

-
[ciflow:skip](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:skip)
- Skip all CI tests for this PR
-
[ciflow:notebooks](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:notebooks)
- Run Jupyter notebooks execution tests for bionemo2
-
[ciflow:slow](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:slow)
- Run slow single GPU integration tests marked as @pytest.mark.slow for
bionemo2
-
[ciflow:all](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all)
- Run all tests (unit tests, slow tests, and notebooks) for bionemo2.
This label can be used to enforce running tests for all bionemo2.
-
[ciflow:all-recipes](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all-recipes)
- Run tests for all recipes (under bionemo-recipes). This label can be
used to enforce running tests for all recipes.

Unit tests marked as `@pytest.mark.multi_gpu` or
`@pytest.mark.distributed` are not run in the PR pipeline.

For more details, see [CONTRIBUTING](CONTRIBUTING.md)

> [!NOTE]
> By default, only basic unit tests are run. Add appropriate labels to
enable an additional test coverage.

#### Authorizing CI Runs

We use
[copy-pr-bot](https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/#automation)
to manage authorization of CI
runs on NVIDIA's compute resources.

- If a pull request is opened by a trusted user and contains only
trusted changes, the pull request's code will
automatically be copied to a pull-request/ prefixed branch in the source
repository (e.g. pull-request/123)
- If a pull request is opened by an untrusted user or contains untrusted
changes, an NVIDIA org member must leave an
`/ok to test` comment on the pull request to trigger CI. This will need
to be done for each new commit.

#### Triggering Code Rabbit AI Review

To trigger a code review from code rabbit, comment on a pull request
with one of these commands:

- @coderabbitai review - Triggers a standard review
- @coderabbitai full review - Triggers a comprehensive review

See https://docs.coderabbit.ai/reference/review-commands for a full list
of commands.

### Pre-submit Checklist

<!--- Ensure all items are completed before submitting -->

- [ ] I have tested these changes locally
- [ ] I have updated the documentation accordingly
- [ ] I have added/updated tests as needed
- [ ] All existing tests pass successfully


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Layer-wise FP8/FP4 quantization control and runtime recipe attachment
for ESM2
* TransformerEngine-optimized ESM2 model variants with per-layer
autocast and NVTX profiling

* **Documentation**
* Expanded low-precision training guide and quantization debugging
examples; added benchmarks and convergence notes

* **Configuration**
* New FP4 settings block and unified quantization statistics config;
support for specifying per-layer precision

* **Tests**
* Added comprehensive tests for quantization utilities and per-layer
behavior
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jonathan Mitchell <jomitchell@nvidia.com>
Co-authored-by: Jonathan Mitchell <jomitchell@umb-b300-dp-218.cl1u1.colossus.nvidia.com>

Author: jomitchellnv

.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