adds NVFP4 support

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

Author: jomitchellnv

bionemo-recipes/recipes/esm2_minifold_te/fp4_debugging_stats.yaml

@@ -0,0 +1,31 @@
+example_fp4_tensor_stat_collection:
+    enabled: True
+    layers:
+        # Match MiniFold te.Linear sublayers in FP4 blocks
+        layer_types: [pi, gi, po, go, fc1, fc2]
+    transformer_engine:
+        LogNvfp4TensorStats:
+            enabled: True
+            tensors_struct:
+            - tensor: activation
+              stats: [underflows%, mse]
+              freq: 100
+            - tensor: gradient
+              stats: [underflows%, mse]
+              freq: 100
+
+example_fp8_tensor_stat_collection:
+    enabled: True
+    layers:
+        # Match MiniFold te.Linear sublayers in FP8 blocks
+        layer_types: [pi, gi, po, go, fc1, fc2]
+    transformer_engine:
+        LogFp8TensorStats:
+            enabled: True
+            tensors_struct:
+            - tensor: activation
+              stats: [mxfp8_underflows%, mxfp8_scale_inv_min, mxfp8_scale_inv_max, mxfp8_mse]
+              freq: 100
+            - tensor: gradient
+              stats: [mxfp8_underflows%, mxfp8_scale_inv_min, mxfp8_scale_inv_max, mxfp8_mse]
+              freq: 100

bionemo-recipes/recipes/esm2_minifold_te/quantization.py

@@ -331,14 +331,12 @@ def update_quant_stats_config(
         config = yaml.safe_load(f)
 
     if "example_fp4_tensor_stat_collection" in config:
-        config["example_fp4_tensor_stat_collection"]["enabled"] = False
+        fp4_regex = generate_layer_regex(fp4_layers, component_precision=component_precision)
+        config["example_fp4_tensor_stat_collection"]["layers"]["layer_name_regex_pattern"] = fp4_regex
         if fp4_layers:
-            logger.warning(
-                "NVFP4 quant stats logging is not yet supported (requires a future TransformerEngine release). "
-                f"Disabling FP4 stats collection for blocks {fp4_layers}. FP8 stats will still be collected."
-            )
+            logger.info(f"Updated FP4 block regex to match blocks: {fp4_layers}")
         else:
-            logger.info("FP4 stats section disabled (no FP4 blocks and feature not yet supported)")
+            logger.info("FP4 blocks empty - regex set to match nothing")
 
     if "example_fp8_tensor_stat_collection" in config:
         fp8_regex = generate_layer_regex(fp8_layers, component_precision=component_precision)

bionemo-recipes/recipes/esm2_minifold_te/tests/test_quantization.py

@@ -256,15 +256,19 @@ def test_none_layers_disables_matching(fp8_only_config):
     assert "DISABLED" in regex
 
 
-def test_fp4_section_disabled_fp8_still_updated(fp4_fp8_config):
+def test_fp4_and_fp8_both_updated(fp4_fp8_config):
     output_path = update_quant_stats_config(config_file=fp4_fp8_config, fp4_layers=[1, 2, 3], fp8_layers=[4, 5, 6])
     with open(output_path) as f:
         result = yaml.safe_load(f)
 
-    assert result["example_fp4_tensor_stat_collection"]["enabled"] is False
+    # FP4 section should have regex for blocks 1-3 (0-indexed 0-2)
+    fp4_regex = result["example_fp4_tensor_stat_collection"]["layers"]["layer_name_regex_pattern"]
+    assert re.search(fp4_regex, "fold.miniformer.blocks.0.transition.fc1")
+    assert re.search(fp4_regex, "fold.miniformer.blocks.2.triangular.pi")
+    assert not re.search(fp4_regex, "fold.miniformer.blocks.3.triangular.pi")
 
+    # FP8 section should have regex for blocks 4-6 (0-indexed 3-5)
     fp8_regex = result["example_fp8_tensor_stat_collection"]["layers"]["layer_name_regex_pattern"]
-    # 1-indexed [4,5,6] -> 0-indexed [3,4,5]
     assert re.search(fp8_regex, "fold.miniformer.blocks.4.triangular.pi")
     assert not re.search(fp8_regex, "fold.miniformer.blocks.1.triangular.pi")