Adds fp8 stats logging to ESM2 / LLAMA3 DDP/FSDP2 training (#1413)

### Description

This MR enables the option to log FP8 stats during training.

#### Usage
By adjusting the hydraconfig or by adding the following commandline
arguments to your training script you can receive FP8 statistics.


```python
python train_fsdp2.py \
fp8_stats_config.enabled=True \ # whether to log stats or not
fp8_stats_config.fp8_log_dir=./logs/fp8_stats_logs_dummy \ # where to store the logs
fp8_stats_config.fp8_stats_file=./fp8_stats.yaml \ # specifies what stats you want to run. Currently this is saved in this yaml file.
fp8_config.enabled=True # set this to use FP8 otherwise stats logging wont work
```

### Type of changes

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

- [ ] Bug fix (non-breaking change which fixes an issue)
- [X] 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.

### 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

## Release Notes

* **New Features**
* FP8 tensor statistics logging during training with configurable stats
collection and debugging outputs.
* FP8 analysis tool for automated metric parsing and visualization with
publication-quality heatmaps.
* Configuration options for enabling FP8 statistics logging and
specifying output directories.

* **Documentation**
  * Added FP8 debugging guides with setup and usage instructions.
* Added FP8 analysis tool user guide with examples and interpretation
guidance.

* **Tests**
* Added validation tests for FP8 statistics logging in training
workflows.
  * Added comprehensive tests for FP8 analysis tool functionality.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

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

Author: jomitchellnv

.devcontainer/recipes/requirements.txt

@@ -16,3 +16,5 @@ transformers
 typer
 wandb
 zstandard
+nvdlfw_inspect @ git+https://github.com/NVIDIA/nvidia-dlfw-inspect
+seaborn

.gitignore

@@ -74,7 +74,6 @@ MNISTCustom/
 *.pot
 
 # Django stuff:
-*.log
 local_settings.py
 db.sqlite3
 
@@ -170,6 +169,10 @@ local/
 
 # Logs
 *.log
+!bionemo-recipes/recipes/fp8_analysis/dummy_logs_esm2/rank_0/nvdlfw_inspect_logs/nvdlfw_inspect_globalrank-0.log
+!bionemo-recipes/recipes/fp8_analysis/dummy_logs_esm2/rank_0/nvdlfw_inspect_statistics_logs/nvdlfw_inspect_globalrank-0.log
+!bionemo-recipes/recipes/fp8_analysis/dummy_logs_llama3/rank_0/nvdlfw_inspect_logs/nvdlfw_inspect_globalrank-0.log
+!bionemo-recipes/recipes/fp8_analysis/dummy_logs_llama3/rank_0/nvdlfw_inspect_statistics_logs/nvdlfw_inspect_globalrank-0.log
 
 # Tests
 tests/__pycache__/

bionemo-recipes/recipes/esm2_native_te/README.md

@@ -106,6 +106,27 @@ configuration parameters, including switching to `MXFP8BlockScaling`, can be set
 python train_fsdp2.py --config-name L0_sanity fp8_config.enabled=true
 ```
 
+#### FP8 Debugging
+
+We also provide a mechanism to receive tensor data related to FP8 layers during training which may include activations, weights and gradients.
+
+To enable this please select the following config options.
+
+```python
+python train_fsdp2.py \
+fp8_stats_config.enabled=True # whether to log stats or not
+fp8_stats_config.fp8_log_dir=./logs/fp8_stats_logs_dummy # where to store the logs
+fp8_stats_config.fp8_stats_file=./fp8_debugging_stats.yaml # specifies what stats you want to run. Currently this is saved in this yaml file.
+fp8_config.enabled=True # set this to use FP8 otherwise stats logging won't work
+```
+
+Note: This feature is available for the `train_ddp` and the `train_fsdp2` scripts. It is not yet available for `train_mfsdp`.
+
+The config file structure [fp8_debugging_stats.yaml](fp8_debugging_stats.yaml) is explained in the [NVIDIA Transformer Engine config file documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/debug/2_config_file_structure.html) in more detail. Below we will cover some very basic elements of the file structure.
+
+This comes as a performance cost that is dependent on the `freq` parameter mentioned above. `freq=1` collects stats on every step which in our
+experiments caused a ~29% decrease in throughput (executed on a single RTX 5090). We recommend using `freq>=10` to reduce this performance hit.
+
 ### Sequence Packing (THD input format)
 
 Sequence packing is handled via a padding-free collator (in `collator.py`) that provides input arguments (e.g.

bionemo-recipes/recipes/esm2_native_te/fp8_debugging_stats.yaml

@@ -0,0 +1,18 @@
+example_fp8_tensor_stat_collection:
+    enabled: True
+    layers:
+        # Match the actual linear layers within attention that support FP8 stats
+        layer_types: [layernorm_qkv]
+    transformer_engine:
+        LogFp8TensorStats:
+            enabled: True
+            tensors_struct:
+            - tensor: activation
+              stats: [underflows%, scale_inv_min, scale_inv_max, mse]
+              freq: 10
+            - tensor: gradient
+              stats: [underflows%, scale_inv_min, scale_inv_max, mse]
+              freq: 10
+            - tensor: weight
+              stats: [underflows%, scale_inv_min, scale_inv_max, mse]
+              freq: 10

bionemo-recipes/recipes/esm2_native_te/hydra_config/defaults.yaml

@@ -75,3 +75,8 @@ checkpoint:
 
 logger:
   frequency: 100
+
+fp8_stats_config:
+  enabled: false
+  fp8_stats_file: ./fp8_debugging_stats.yaml
+  fp8_log_dir: ./log_fp8_stats

bionemo-recipes/recipes/esm2_native_te/requirements.txt

@@ -10,3 +10,4 @@ tqdm
 transformer_engine[pytorch]
 transformers
 wandb
+nvdlfw_inspect @ git+https://github.com/NVIDIA/nvidia-dlfw-inspect

bionemo-recipes/recipes/esm2_native_te/tests/conftest.py

@@ -34,6 +34,17 @@ def recipe_path() -> Path:
     return Path(__file__).parent.parent
 
 
+def pytest_collection_modifyitems(items):
+    """Run FP8 stats logging tests first to avoid late debug initialization."""
+    stats_test_names = {
+        "test_sanity_ddp_fp8_stats_logging",
+        "test_sanity_fsdp2_fp8_stats_logging",
+    }
+    stats_tests = [item for item in items if item.name in stats_test_names]
+    other_tests = [item for item in items if item.name not in stats_test_names]
+    items[:] = stats_tests + other_tests
+
+
 @pytest.fixture(scope="session", autouse=True)
 def device_mesh():
     """Create a re-usable device mesh for testing.

bionemo-recipes/recipes/esm2_native_te/tests/test_train.py

@@ -142,6 +142,46 @@ def test_sanity_ddp_fp8(tmp_path, recipe_path):
     main_ddp(sanity_config)
 
 
+@requires_fp8
+def test_sanity_ddp_fp8_stats_logging(tmp_path, recipe_path):
+    """Test that FP8 stats logging creates the expected log files."""
+    fp8_log_dir = tmp_path / "fp8_stats_logs"
+
+    with initialize_config_dir(config_dir=str(recipe_path / "hydra_config"), version_base="1.2"):
+        sanity_config = compose(
+            config_name="L0_sanity",
+            overrides=[
+                f"+wandb_init_args.dir={tmp_path}",
+                f"checkpoint.ckpt_dir={tmp_path}",
+                "fp8_config.enabled=true",
+                "fp8_stats_config.enabled=true",
+                f"fp8_stats_config.fp8_log_dir={fp8_log_dir}",
+                "num_train_steps=4",
+            ],
+        )
+
+    main_ddp(sanity_config)
+
+    # Verify the log directory structure was created
+    assert fp8_log_dir.exists(), "FP8 log directory was not created"
+    assert (fp8_log_dir / "rank_0").exists(), "rank_0 directory was not created"
+    assert (fp8_log_dir / "rank_0" / "nvdlfw_inspect_logs").exists(), "nvdlfw_inspect_logs directory was not created"
+    assert (fp8_log_dir / "rank_0" / "nvdlfw_inspect_statistics_logs").exists(), (
+        "nvdlfw_inspect_statistics_logs directory was not created"
+    )
+
+    # Verify the log files exist
+    metadata_log = fp8_log_dir / "rank_0" / "nvdlfw_inspect_logs" / "nvdlfw_inspect_globalrank-0.log"
+    stats_log = fp8_log_dir / "rank_0" / "nvdlfw_inspect_statistics_logs" / "nvdlfw_inspect_globalrank-0.log"
+
+    assert metadata_log.exists(), "Metadata log file was not created"
+    assert stats_log.exists(), "Statistics log file was not created"
+
+    # Verify files are non-empty
+    assert metadata_log.stat().st_size > 0, "Metadata log file is empty"
+    assert stats_log.stat().st_size > 0, "Statistics log file is empty"
+
+
 @requires_fp8
 def test_sanity_convergence_fsdp2_fp8(tmp_path, recipe_path):
     """For FSDP2, we check that the script can run successfully with FP8 and check convergence."""
@@ -159,6 +199,32 @@ def test_sanity_convergence_fsdp2_fp8(tmp_path, recipe_path):
     assert final_loss < 3.0, f"Final loss {final_loss} is too high"
 
 
+@requires_fp8
+def test_sanity_fsdp2_fp8_stats_logging(tmp_path, recipe_path):
+    """Test that FP8 stats logging works with FSDP2."""
+    fp8_log_dir = tmp_path / "fp8_stats_logs"
+
+    with initialize_config_dir(config_dir=str(recipe_path / "hydra_config"), version_base="1.2"):
+        sanity_config = compose(
+            config_name="L0_sanity",
+            overrides=[
+                f"+wandb_init_args.dir={tmp_path}",
+                f"checkpoint.ckpt_dir={tmp_path}",
+                "fp8_config.enabled=true",
+                "fp8_stats_config.enabled=true",
+                f"fp8_stats_config.fp8_log_dir={fp8_log_dir}",
+                "num_train_steps=4",
+            ],
+        )
+
+    main_fsdp2(sanity_config)
+
+    # Verify log structure (same assertions as above)
+    assert fp8_log_dir.exists()
+    assert (fp8_log_dir / "rank_0" / "nvdlfw_inspect_logs" / "nvdlfw_inspect_globalrank-0.log").exists()
+    assert (fp8_log_dir / "rank_0" / "nvdlfw_inspect_statistics_logs" / "nvdlfw_inspect_globalrank-0.log").exists()
+
+
 @requires_fp8
 @pytest.mark.xfail(reason="MFSDP doesn't seem to support fp8_model_init (BIONEMO-3012)")
 def test_sanity_mfsdp_fp8_and_model_init(tmp_path, recipe_path):

bionemo-recipes/recipes/esm2_native_te/train_ddp.py

@@ -17,7 +17,9 @@
 from pathlib import Path
 
 import hydra
+import nvdlfw_inspect.api as debug_api
 import torch
+import transformer_engine
 import transformer_engine.pytorch
 from omegaconf import DictConfig
 from torch.distributed.device_mesh import init_device_mesh
@@ -50,6 +52,24 @@ def main(args: DictConfig) -> float | None:
     torch.distributed.init_process_group(backend="nccl", device_id=device)
     torch.cuda.set_device(dist_config.local_rank)
 
+    # TE Debug feature logging
+    if args.fp8_stats_config.enabled and not args.fp8_config.enabled:
+        raise ValueError(
+            "fp8_stats_config.enabled is true but fp8_config.enabled is false, please set fp8_config.enabled to true in the config if you wish to collect FP8 stats"
+        )
+
+    if args.fp8_stats_config.enabled:
+        fp8_stats_file = args.fp8_stats_config.fp8_stats_file
+        fp8_log_dir = Path(args.fp8_stats_config.fp8_log_dir) / f"rank_{dist_config.rank}"
+        fp8_log_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Logging FP8 stats to {fp8_log_dir}")
+        te_features_dir = str(Path(transformer_engine.__file__).parent / "debug" / "features")
+        debug_api.initialize(
+            config_file=fp8_stats_file,
+            feature_dirs=[te_features_dir],
+            log_dir=fp8_log_dir,
+            default_logging_enabled=True,
+        )
     # Create a device mesh for DDP. While this isn't strictly necessary, it mirrors the device mesh we create for FSDP2
     # and MFSDP.
     device_mesh = init_device_mesh("cuda", mesh_shape=(dist_config.world_size,), mesh_dim_names=("ddp",))
@@ -84,6 +104,9 @@ def main(args: DictConfig) -> float | None:
     optimizer = AdamW(model.parameters(), **args.adamw_kwargs)
     scheduler = get_linear_schedule_with_warmup(optimizer, **args.lr_scheduler_kwargs)
 
+    if args.fp8_stats_config.enabled:
+        debug_api.infer_and_assign_layer_names(model)
+
     model = model.to(device=device)
     model = torch.nn.parallel.DistributedDataParallel(
         model,
@@ -134,6 +157,8 @@ def main(args: DictConfig) -> float | None:
             loss = outputs.loss
             loss.backward()
 
+            if args.fp8_stats_config.enabled:
+                debug_api.step()
             # Compute and clip gradient norms.
             total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0).item()
 
@@ -181,6 +206,8 @@ def main(args: DictConfig) -> float | None:
 
     # Clean up distributed training
     perf_logger.finish()
+    if args.fp8_stats_config.enabled:
+        debug_api.end_debug()
     torch.distributed.destroy_process_group()
 
     return perf_logger.min_loss

bionemo-recipes/recipes/esm2_native_te/train_fsdp2.py

@@ -18,7 +18,9 @@
 from pathlib import Path
 
 import hydra
+import nvdlfw_inspect.api as debug_api
 import torch
+import transformer_engine
 import transformer_engine.pytorch
 from omegaconf import DictConfig, OmegaConf
 from torch.distributed.device_mesh import init_device_mesh
@@ -55,6 +57,25 @@ def main(args: DictConfig) -> float | None:
     torch.distributed.init_process_group(backend="nccl", device_id=device)
     torch.cuda.set_device(dist_config.local_rank)
 
+    # TE Debug feature logging - MUST be done BEFORE FSDP wrapping
+    if args.fp8_stats_config.enabled and not args.fp8_config.enabled:
+        raise ValueError(
+            "fp8_stats_config.enabled is true but fp8_config.enabled is false, please set fp8_config.enabled to true in the config if you wish to collect FP8 stats"
+        )
+
+    if args.fp8_stats_config.enabled:
+        fp8_stats_file = args.fp8_stats_config.fp8_stats_file
+        fp8_log_dir = Path(args.fp8_stats_config.fp8_log_dir) / f"rank_{dist_config.rank}"
+        fp8_log_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Logging FP8 stats to {fp8_log_dir}")
+        te_features_dir = str(Path(transformer_engine.__file__).parent / "debug" / "features")
+        debug_api.initialize(
+            config_file=fp8_stats_file,
+            feature_dirs=[te_features_dir],
+            log_dir=fp8_log_dir,
+            default_logging_enabled=True,
+        )
+
     # Create a device mesh for FSDP.
     device_mesh = init_device_mesh(
         "cuda",
@@ -86,6 +107,7 @@ def main(args: DictConfig) -> float | None:
 
     # We call the transformer stack "layers" in our TE models, but it's called "layer" in the original ESM-2 models.
     transformer_stack = model.esm.encoder.layers if hasattr(model.esm.encoder, "layers") else model.esm.encoder.layer
+
     for layer in transformer_stack:
         fully_shard(layer, mesh=device_mesh["dp"])
     fully_shard(model, mesh=device_mesh["dp"])
@@ -100,6 +122,10 @@ def main(args: DictConfig) -> float | None:
             model.to_empty(device=device)
             model.apply(model._init_weights)
 
+    # Assign names to layers so debug API can identify them
+    if args.fp8_stats_config.enabled:
+        debug_api.infer_and_assign_layer_names(model)
+
     # Create optimizer. Convert OmegaConf to regular dict to avoid serialization issues (BIONEMO-2873).
     optimizer = AdamW(model.parameters(), **OmegaConf.to_container(args.adamw_kwargs, resolve=True))  # type: ignore
     scheduler = get_linear_schedule_with_warmup(optimizer, **args.lr_scheduler_kwargs)
@@ -152,6 +178,10 @@ def main(args: DictConfig) -> float | None:
             # Step optimizer.
             optimizer.step()
             scheduler.step()
+
+            if args.fp8_stats_config.enabled:
+                debug_api.step()
+
             optimizer.zero_grad()
 
             perf_logger.log_step(
@@ -193,6 +223,8 @@ def main(args: DictConfig) -> float | None:
 
     # Clean up distributed training
     perf_logger.finish()
+    if args.fp8_stats_config.enabled:
+        debug_api.end_debug()
     torch.distributed.destroy_process_group()
 
     return perf_logger.min_loss