feat(recipes): add evo2_megatron_quant — INT8 weight compression + ModelOpt PTQ harness

recipes/evo2_megatron_quant/.ruff.toml

@@ -0,0 +1 @@
+extend = "../.ruff.toml"

recipes/evo2_megatron_quant/README.md

@@ -0,0 +1,98 @@
+# evo2_megatron_quant
+
+Inference-side **INT8 real weight compression** and a **ModelOpt PTQ analysis
+harness** for NVIDIA BioNeMo biological foundation models (ESM-2, Geneformer,
+Evo2), on top of
+[NVIDIA ModelOpt](https://github.com/NVIDIA/TensorRT-Model-Optimizer).
+
+Upstream `bionemo-recipes` covers low-precision *training* (TransformerEngine
+FP8/MXFP8/NVFP4); this recipe adds an *inference*-side path:
+
+1. **Real INT8 weight compression** — a drop-in `nn.Linear` replacement
+   (`INT8Linear`) that stores weights in int8 (per-channel symmetric scaling)
+   and dequantizes on the fly, for genuine GPU-memory savings.
+2. **ModelOpt PTQ harness** — wires BioNeMo model loading to `mtq.quantize()`
+   (simulated Q/DQ) with an automated quality report (cosine similarity, top-k
+   agreement, MSE) vs. the BF16 baseline.
+
+> This is the first, deliberately narrow slice. FP8/INT4 real weight compression
+> and KV-cache experiments are intentionally left out of this PR and will follow
+> once their GPU quality is re-measured.
+
+## Results (INT8 real weight compression)
+
+Measured on **Evo2 7B** (H200), compressing the TE `Linear` layers:
+
+| Precision | Model memory | Saved | Cosine sim | Top-1 |
+|-----------|:------------:|:-----:|:----------:|:-----:|
+| BF16 (baseline) | 13,035 MB | — | 1.000 | 100% |
+| **INT8** (per-channel) | **11,057 MB** | **~15%** | **0.998** | **100%** |
+
+Most of Evo2's parameters live in non-`Linear` Hyena/SSM components that this
+pass does not touch, so the saving reflects compressing the Linear layers only.
+
+## Supported models
+
+| Model | Architecture | Domain | Params | Checkpoint (NGC) |
+|-------|-------------|--------|:------:|-------------------|
+| ESM-2 | Transformer encoder | Protein | 8M | `esm2/8m:2.0` |
+| Geneformer | Transformer encoder | Gene expression | 10M | `geneformer/10M_241113:2.0` |
+| Evo2 | Hyena/Mamba SSM | DNA | 7B | `evo2/7b-8k:1.0` |
+
+## Layout
+
+```
+evo2_megatron_quant/
+├── src/
+│   ├── adapters.py            # per-model load / tokenize / forward adapters
+│   ├── quantize.py            # ModelOpt mtq.quantize() wrapper
+│   ├── metrics.py             # cosine sim / top-k / MSE
+│   ├── compressed_linear.py   # INT8Linear (real int8 weight storage)
+│   └── compress_model.py      # walk model, swap Linear -> INT8Linear
+├── tests/
+│   ├── test_compressed_linear_roundtrip.py   # L0 sanity: CPU-only, no checkpoint
+│   └── test_single_model.py                  # PTQ harness demo (GPU + NGC, manual)
+├── scripts/
+│   ├── download_models.sh
+│   └── benchmark_real_quant.py               # INT8 memory + quality benchmark (GPU)
+├── configs/quant_methods.yaml
+├── docker/start_container.sh
+├── requirements.txt
+└── README.md
+```
+
+## Quick start
+
+```bash
+# L0 sanity tests (CPU-only, run in CI — no GPU or checkpoint needed):
+pytest tests/test_compressed_linear_roundtrip.py -v
+
+# Real INT8 weight compression benchmark (inside the BioNeMo container, GPU + NGC):
+python scripts/benchmark_real_quant.py --precisions int8
+```
+
+Programmatic use:
+
+```python
+from src.compress_model import compress_model
+
+stats = compress_model(evo2_model, precision="int8")   # swaps Linear -> INT8Linear
+# -> ~15% GPU memory reduction on Evo2 7B, cosine sim ~0.998 vs BF16
+```
+
+## How INT8 weight compression works
+
+1. Extract weight/bias from each `nn.Linear` / TE `Linear`.
+2. Quantize to int8 on CPU (avoids a GPU memory spike), per output channel:
+   `scale = amax(|W|, dim=1) / 127`, `W_int8 = round(W / scale)`.
+3. Replace the layer with an `INT8Linear` that dequantizes (`W_int8 * scale`) on
+   the fly for the matmul.
+4. Free the original BF16 weights.
+
+## Tests
+
+`tests/test_compressed_linear_roundtrip.py` is the CPU-only L0 gate: it
+round-trips INT8 on tiny random layers and asserts near-lossless reconstruction,
+no systematic weight bias, and a real reduction in the stored buffer size.
+`test_single_model.py` and the benchmark require a GPU and NGC checkpoints and
+are meant to be run manually, not in fast CI.

recipes/evo2_megatron_quant/configs/quant_methods.yaml

@@ -0,0 +1,222 @@
+# BioNeMo Quantization Methods Configuration
+#
+# All 22 NVIDIA ModelOpt quantization configurations, organized by family.
+# Each entry maps to a constant in modelopt.torch.quantization (e.g., mtq.FP8_DEFAULT_CFG).
+#
+# Fields:
+#   name:        ModelOpt constant name (pass to --quant CLI argument)
+#   family:      Quantization family (FP8, INT8, INT4, NVFP4, MX, Mixed)
+#   weight_bits: Weight precision in bits
+#   act_bits:    Activation precision in bits (null = not quantized)
+#   description: Brief description of the method
+#   calibration: Whether method requires calibration data
+#   notes:       Additional notes or warnings
+
+methods:
+
+  # ================================
+  # FP8 Family (E4M3 / E5M2)
+  # ================================
+  # Highest precision among quantized formats. Minimal quality loss.
+  # Uses 4-bit exponent + 3-bit mantissa (E4M3) format.
+
+  - name: FP8_DEFAULT_CFG
+    family: FP8
+    weight_bits: 8
+    act_bits: 8
+    description: "FP8 E4M3, per-tensor quantization"
+    calibration: false
+    notes: "Best starting point. Works well with all 3 models."
+
+  - name: FP8_PER_CHANNEL_PER_TOKEN_CFG
+    family: FP8
+    weight_bits: 8
+    act_bits: 8
+    description: "FP8 per-channel weights, per-token activations"
+    calibration: false
+    notes: "Finer-grained scaling for better accuracy."
+
+  - name: FP8_KV_CFG
+    family: FP8
+    weight_bits: 8
+    act_bits: 8
+    description: "FP8 with KV cache quantization"
+    calibration: false
+    notes: "Also quantizes KV cache for attention layers. No effect on Hyena layers."
+
+  - name: FP8_AFFINE_KV_CFG
+    family: FP8
+    weight_bits: 8
+    act_bits: 8
+    description: "FP8 with affine KV cache quantization"
+    calibration: false
+    notes: "Affine (zero-point + scale) KV quantization for better range coverage."
+
+  - name: FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG
+    family: FP8
+    weight_bits: 8
+    act_bits: null
+    description: "FP8 2D blockwise, weight-only quantization"
+    calibration: false
+    notes: "Only quantizes weights. Can improve inference latency with no activation overhead."
+
+  # ================================
+  # INT8 Family
+  # ================================
+  # Good balance of compression and quality.
+
+  - name: INT8_DEFAULT_CFG
+    family: INT8
+    weight_bits: 8
+    act_bits: 8
+    description: "INT8 symmetric, per-tensor quantization"
+    calibration: false
+    notes: "Straightforward integer quantization."
+
+  - name: INT8_SMOOTHQUANT_CFG
+    family: INT8
+    weight_bits: 8
+    act_bits: 8
+    description: "INT8 SmoothQuant — migrates quantization difficulty from activations to weights"
+    calibration: true
+    notes: "Requires calibration data. Better handles activation outliers."
+
+  # ================================
+  # INT4 Family
+  # ================================
+  # Higher compression. May show quality degradation on sensitive layers.
+
+  - name: INT4_AWQ_CFG
+    family: INT4
+    weight_bits: 4
+    act_bits: 16
+    description: "INT4 Activation-aware Weight Quantization"
+    calibration: true
+    notes: "Identifies salient weights via activation patterns. Keeps activations in FP16."
+
+  - name: INT4_BLOCKWISE_WEIGHT_ONLY_CFG
+    family: INT4
+    weight_bits: 4
+    act_bits: null
+    description: "INT4 blockwise weight-only quantization"
+    calibration: false
+    notes: "Block-level scaling for weight-only INT4."
+
+  # ================================
+  # NVFP4 Family (NVIDIA FP4)
+  # ================================
+  # NVIDIA's custom 4-bit floating point format.
+  # Designed for Hopper/Blackwell GPU architectures.
+
+  - name: NVFP4_DEFAULT_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 default configuration"
+    calibration: false
+    notes: "NVIDIA's FP4 format. Requires Hopper+ GPU."
+
+  - name: NVFP4_AWQ_LITE_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 with AWQ lite calibration"
+    calibration: true
+    notes: "Lightweight AWQ calibration for NVFP4."
+
+  - name: NVFP4_AWQ_CLIP_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 with AWQ clip calibration"
+    calibration: true
+    notes: "Clips outlier weights before quantization."
+
+  - name: NVFP4_AWQ_FULL_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 with full AWQ calibration"
+    calibration: true
+    notes: "Most thorough AWQ calibration, but slower."
+
+  - name: NVFP4_KV_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 with KV cache quantization"
+    calibration: false
+    notes: "Also quantizes attention KV cache."
+
+  - name: NVFP4_KV_ROTATE_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 with rotated KV cache quantization"
+    calibration: false
+    notes: "Applies rotation to KV before quantization for better quality."
+
+  - name: NVFP4_AFFINE_KV_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 with affine KV cache quantization"
+    calibration: false
+    notes: "Affine (zero-point + scale) KV for NVFP4."
+
+  - name: NVFP4_SVDQUANT_DEFAULT_CFG
+    family: NVFP4
+    weight_bits: 4
+    act_bits: 8
+    description: "NVFP4 with SVD quantization"
+    calibration: true
+    notes: "Uses SVD decomposition before quantization."
+
+  # ================================
+  # MX (Microscaling) Family
+  # ================================
+  # Block-scaled formats using the OCP Microscaling specification.
+
+  - name: MXFP4_DEFAULT_CFG
+    family: MX
+    weight_bits: 4
+    act_bits: 4
+    description: "Microscaling FP4 — 4-bit weights and activations"
+    calibration: false
+    notes: "Maximum compression among MX formats. May cause quality loss."
+
+  - name: MXFP6_DEFAULT_CFG
+    family: MX
+    weight_bits: 6
+    act_bits: 6
+    description: "Microscaling FP6 — 6-bit weights and activations"
+    calibration: false
+    notes: "Good balance between compression and quality."
+
+  - name: MXFP8_DEFAULT_CFG
+    family: MX
+    weight_bits: 8
+    act_bits: 8
+    description: "Microscaling FP8 — 8-bit with block scaling"
+    calibration: false
+    notes: "Similar to FP8 but uses OCP block scaling."
+
+  - name: MXINT8_DEFAULT_CFG
+    family: MX
+    weight_bits: 8
+    act_bits: 8
+    description: "Microscaling INT8 — 8-bit integer with block scaling"
+    calibration: false
+    notes: "INT8 variant of MX format."
+
+  # ================================
+  # Mixed Precision
+  # ================================
+
+  - name: W4A8_AWQ_BETA_CFG
+    family: Mixed
+    weight_bits: 4
+    act_bits: 8
+    description: "4-bit weights + 8-bit activations (AWQ beta)"
+    calibration: true
+    notes: "Aggressive weight compression with higher-precision activations."

recipes/evo2_megatron_quant/docker/start_container.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+# ==============================================================================
+# BioNeMo Framework Container Launcher
+# ==============================================================================
+#
+# Starts the NVIDIA BioNeMo Framework Docker container with GPU support.
+# All quantization work runs inside this container.
+#
+# Prerequisites:
+#   - Docker with NVIDIA Container Toolkit
+#   - NGC authentication (nvcr.io access)
+#   - At least one NVIDIA GPU
+#
+# Usage:
+#   bash docker/start_container.sh
+# ==============================================================================
+
+set -euo pipefail
+
+# ---------- Configuration ----------
+BIONEMO_TAG="${BIONEMO_TAG:-2.6.3}"                     # BioNeMo container version
+LOCAL_WORKSPACE="${LOCAL_WORKSPACE:-$(pwd)/quantization}" # Mount to /workspace/bionemo_dev
+DATA_DIR="${DATA_DIR:-$(pwd)/data}"                       # Mount to /data (model cache)
+CONTAINER_NAME="${CONTAINER_NAME:-bionemo-quant}"         # Docker container name
+
+# ---------- Create directories ----------
+mkdir -p "${LOCAL_WORKSPACE}"
+mkdir -p "${DATA_DIR}"
+
+echo "================================================"
+echo "  BioNeMo Framework Container"
+echo "  Image:     nvcr.io/nvidia/clara/bionemo-framework:${BIONEMO_TAG}"
+echo "  Workspace: ${LOCAL_WORKSPACE} → /workspace/bionemo_dev"
+echo "  Data:      ${DATA_DIR} → /data"
+echo "  Container: ${CONTAINER_NAME}"
+echo "================================================"
+
+# ---------- Launch container ----------
+docker run --rm -it \
+    --gpus all \
+    --shm-size=32g \
+    --ulimit memlock=-1 \
+    --ulimit stack=67108864 \
+    --net=host \
+    -v "${LOCAL_WORKSPACE}":/workspace/bionemo_dev \
+    -v "${DATA_DIR}":/data \
+    -w /workspace/bionemo_dev \
+    --name "${CONTAINER_NAME}" \
+    nvcr.io/nvidia/clara/bionemo-framework:"${BIONEMO_TAG}" \
+    /bin/bash

recipes/evo2_megatron_quant/requirements.txt

@@ -0,0 +1,14 @@
+# Dependencies for the evo2_megatron_quant recipe.
+#
+# Provided by the BioNeMo / nvcr.io PyTorch base image — do NOT re-pin these here,
+# re-pinning risks clashing with the container build:
+#   torch, transformer-engine, megatron-core, nemo-toolkit, bionemo-*
+#
+# Recipe-specific dependencies (pin to the versions shipped in the BioNeMo
+# container you target before opening the PR):
+nvidia-modelopt[torch]   # ModelOpt PTQ bridge used by src/quantize.py
+pyyaml                   # reads configs/quant_methods.yaml
+
+# Test-only (the L0 sanity test tests/test_compressed_linear_roundtrip.py is
+# CPU-only and needs just torch + pytest):
+pytest