diff --git a/recipes/evo2_megatron_quant/.ruff.toml b/recipes/evo2_megatron_quant/.ruff.toml new file mode 100644 index 0000000000..7e9a31bf5d --- /dev/null +++ b/recipes/evo2_megatron_quant/.ruff.toml @@ -0,0 +1 @@ +extend = "../.ruff.toml" diff --git a/recipes/evo2_megatron_quant/README.md b/recipes/evo2_megatron_quant/README.md new file mode 100644 index 0000000000..2991717ad0 --- /dev/null +++ b/recipes/evo2_megatron_quant/README.md @@ -0,0 +1,105 @@ +# 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 + +```text +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 +├── scripts/ +│ ├── download_models.sh +│ ├── run_single_model.py # ModelOpt PTQ harness runner (GPU + NGC) +│ └── 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 + +# ModelOpt PTQ harness on one model (GPU + NGC): +python scripts/run_single_model.py --model esm2 --quant INT8_DEFAULT_CFG +``` + +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/` holds only the CPU-only L0 gate, +`tests/test_compressed_linear_roundtrip.py`: 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. It needs neither a GPU nor a +checkpoint, so it runs in normal CI. + +The GPU runners live under `scripts/` (`run_single_model.py`, +`benchmark_real_quant.py`) — they require a CUDA device and NGC checkpoints and +are meant to be run manually, so they are deliberately not pytest tests. diff --git a/recipes/evo2_megatron_quant/configs/quant_methods.yaml b/recipes/evo2_megatron_quant/configs/quant_methods.yaml new file mode 100644 index 0000000000..1a70c1c062 --- /dev/null +++ b/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." diff --git a/recipes/evo2_megatron_quant/docker/start_container.sh b/recipes/evo2_megatron_quant/docker/start_container.sh new file mode 100644 index 0000000000..0a729f76c8 --- /dev/null +++ b/recipes/evo2_megatron_quant/docker/start_container.sh @@ -0,0 +1,58 @@ +#!/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 ---------- +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +RECIPE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" # recipe root (src/, scripts/, tests/) +BIONEMO_TAG="${BIONEMO_TAG:-2.6.3}" # BioNeMo container version +LOCAL_WORKSPACE="${LOCAL_WORKSPACE:-${RECIPE_ROOT}}" # 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 + +# Host networking is opt-in (some NCCL/distributed setups need it). Enable with USE_HOST_NET=1. +NET_ARGS=() +if [ "${USE_HOST_NET:-0}" = "1" ]; then + NET_ARGS+=(--net=host) +fi + +# ---------- 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_ARGS[@]+"${NET_ARGS[@]}"} \ + -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 diff --git a/recipes/evo2_megatron_quant/requirements.txt b/recipes/evo2_megatron_quant/requirements.txt new file mode 100644 index 0000000000..a5d68ac2ac --- /dev/null +++ b/recipes/evo2_megatron_quant/requirements.txt @@ -0,0 +1,10 @@ +# 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-*, +# and nvidia-modelopt (vendored via NeMo; imported by src/quantize.py). +# +# Recipe-specific dependencies, pinned for reproducibility: +pyyaml==6.0.2 # reads configs/quant_methods.yaml +pytest==8.3.4 # CPU-only L0 sanity test (tests/test_compressed_linear_roundtrip.py) diff --git a/recipes/evo2_megatron_quant/scripts/benchmark_real_quant.py b/recipes/evo2_megatron_quant/scripts/benchmark_real_quant.py new file mode 100644 index 0000000000..640475b3a9 --- /dev/null +++ b/recipes/evo2_megatron_quant/scripts/benchmark_real_quant.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +# 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. + +"""Benchmark Evo2 with REAL INT8 weight compression. + +Unlike ModelOpt simulated quantization, this script truly stores weights in +int8, achieving real GPU-memory savings (~2x on the compressed Linear layers). + +For each config: + 1. Load Evo2 in BF16 + 2. Run BF16 inference → baseline logits + 3. Compress weights to int8 (in-place) + 4. Measure real memory savings + 5. Run compressed inference → compare quality + 6. Binary search for max sequence length + +Usage (inside BioNeMo Docker): + python scripts/benchmark_real_quant.py --precisions int8 +""" + +import argparse +import csv +import json +import os +import subprocess +import sys +import time + + +# === Subprocess worker === +WORKER_SCRIPT = r''' +import gc, json, os, sys, time +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", os.environ.get("WORKER_PORT", "29500")) +os.environ.setdefault("WORLD_SIZE", "1") +os.environ.setdefault("RANK", "0") +os.environ.setdefault("LOCAL_RANK", "0") + +import torch +import torch.nn.functional as F +import torch.distributed as dist +if not dist.is_initialized(): + dist.init_process_group(backend="nccl", world_size=1, rank=0) +from megatron.core import parallel_state +if not parallel_state.is_initialized(): + parallel_state.initialize_model_parallel(1, 1) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +model_parallel_cuda_manual_seed(42) + +import nemo.lightning as nl +from bionemo.core.data.load import load as bionemo_load +from bionemo.evo2.data.tokenizer import Evo2Tokenizer + +# Add src to path +sys.path.insert(0, "/workspace/quantization/src") +from compressed_linear import METHOD_TO_PRECISION +from compress_model import compress_model + +# Parse args +precision = sys.argv[1] # int8 +method_name = sys.argv[2] # e.g. FP8_DEFAULT_CFG +eval_seq_len = int(sys.argv[3]) # seq length for quality eval +max_seqlen_search = sys.argv[4] == "1" # whether to binary-search max seqlen +output_file = sys.argv[5] + +result = { + "method": method_name, + "precision": precision, + "eval_seq_len": eval_seq_len, + "status": "ERROR", "error": "", + # Memory + "mem_bf16_mb": 0, "mem_compressed_mb": 0, "mem_saved_mb": 0, "mem_saved_pct": 0, + # Quality + "cos_sim_avg": 0.0, "cos_sim_min": 0.0, "top1_agree": 0.0, "mse": 0.0, + # Speed + "infer_bf16_ms": 0.0, "infer_compressed_ms": 0.0, + # Compression stats + "linears_compressed": 0, "params_compressed": 0, + # Max sequence length + "max_seqlen": 0, "max_seqlen_peak_gb": 0.0, +} + +try: + # Load model + ckpt = str(bionemo_load("evo2/7b-8k:1.0", source="ngc")) + ctx = nl.io.load_context(ckpt) + config = ctx.model.config + config.initial_ckpt_path = ckpt + tokenizer = Evo2Tokenizer() + tokenizer.vocab_size = tokenizer.tokenizer.vocab_size + model = config.configure_model(tokenizer) + model = model.bfloat16().cuda().eval() + + mem_bf16 = torch.cuda.memory_allocated() / 1e6 + result["mem_bf16_mb"] = mem_bf16 + + # Generate eval input + input_ids = torch.randint(0, 512, (1, eval_seq_len), dtype=torch.long).cuda() + position_ids = torch.arange(eval_seq_len, device="cuda").unsqueeze(0) + attention_mask = torch.ones_like(input_ids) + + # BF16 baseline + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + # Warmup + warmup_out = model(input_ids=input_ids[:, :256], position_ids=position_ids[:, :256], + attention_mask=attention_mask[:, :256]) + torch.cuda.synchronize() + del warmup_out + + t0 = time.perf_counter() + bf16_out = model(input_ids=input_ids, position_ids=position_ids, + attention_mask=attention_mask) + torch.cuda.synchronize() + result["infer_bf16_ms"] = (time.perf_counter() - t0) * 1000 + + bf16_logits = bf16_out if isinstance(bf16_out, torch.Tensor) else \ + getattr(bf16_out, "token_logits", getattr(bf16_out, "logits", bf16_out)) + + # Compress model + print(f" Compressing to {precision}...") + stats = compress_model(model, precision=precision, verbose=True) + result["mem_compressed_mb"] = stats["mem_after_mb"] + result["mem_saved_mb"] = stats["mem_saved_mb"] + result["mem_saved_pct"] = stats["mem_saved_pct"] + result["linears_compressed"] = stats["compressed_linears"] + result["params_compressed"] = stats["params_compressed"] + + # Compressed inference + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + t0 = time.perf_counter() + q_out = model(input_ids=input_ids, position_ids=position_ids, + attention_mask=attention_mask) + torch.cuda.synchronize() + result["infer_compressed_ms"] = (time.perf_counter() - t0) * 1000 + + q_logits = q_out if isinstance(q_out, torch.Tensor) else \ + getattr(q_out, "token_logits", getattr(q_out, "logits", q_out)) + + # Quality metrics + bf = bf16_logits.reshape(-1, bf16_logits.shape[-1]).float() + qf = q_logits.reshape(-1, q_logits.shape[-1]).float() + cos = F.cosine_similarity(bf, qf, dim=-1) + result["cos_sim_avg"] = cos.mean().item() + result["cos_sim_min"] = cos.min().item() + top1_bf = bf.argmax(dim=-1) + top1_q = qf.argmax(dim=-1) + result["top1_agree"] = (top1_bf == top1_q).float().mean().item() + result["mse"] = F.mse_loss(bf, qf).item() + + # Free eval tensors + del bf16_logits, q_logits, bf16_out, q_out, bf, qf, input_ids, position_ids, attention_mask + gc.collect() + torch.cuda.empty_cache() + + # Max sequence length: linear sweep over candidate lengths; report the largest that fits. + if max_seqlen_search: + print(" Sweeping candidate sequence lengths...") + test_lens = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144, 524288] + best_len = 0 + best_peak = 0.0 + + for slen in test_lens: + ids = pos = mask = probe_out = None + oom = False + try: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + ids = torch.randint(0, 512, (1, slen), dtype=torch.long).cuda() + pos = torch.arange(slen, device="cuda").unsqueeze(0) + mask = torch.ones_like(ids) + + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + probe_out = model(input_ids=ids, position_ids=pos, attention_mask=mask) + torch.cuda.synchronize() + + peak = torch.cuda.max_memory_allocated() / 1e9 + print(f" seq_len={slen:>8d}: OK (peak={peak:.1f} GB)") + best_len = slen + best_peak = peak + except (torch.cuda.OutOfMemoryError, RuntimeError) as e: + if "out of memory" in str(e).lower(): + print(f" seq_len={slen:>8d}: OOM") + oom = True + else: + raise + finally: + del ids, pos, mask, probe_out + gc.collect() + torch.cuda.empty_cache() + if oom: + break + + result["max_seqlen"] = best_len + result["max_seqlen_peak_gb"] = best_peak + + result["status"] = "PASS" + +except Exception as e: + import traceback + result["error"] = traceback.format_exc()[-800:] + +with open(output_file, "w") as f: + json.dump(result, f, indent=2) +print(f" Result saved to {output_file}") +''' + + +def run_test(precision, method_name, eval_seq_len, max_seqlen_search, port, output_dir): + """Run one compression test in isolated subprocess.""" + output_file = os.path.join(output_dir, f"result_{method_name}_{precision}.json") + worker_file = os.path.join(output_dir, "_worker_real.py") + + with open(worker_file, "w") as f: + f.write(WORKER_SCRIPT) + + # Remove any stale result so a crashed worker cannot be mistaken for a fresh run. + if os.path.exists(output_file): + os.remove(output_file) + + env = os.environ.copy() + env["WORKER_PORT"] = str(port) + env["CUDA_VISIBLE_DEVICES"] = "0" + + cmd = [ + sys.executable, worker_file, + precision, method_name, str(eval_seq_len), + "1" if max_seqlen_search else "0", + output_file, + ] + + print(f"\n{'─' * 70}") + print(f" {method_name} → {precision.upper()}") + print(f"{'─' * 70}") + + t0 = time.time() + try: + subprocess.run(cmd, env=env, capture_output=False, timeout=1200) + except subprocess.TimeoutExpired: + return { + "method": method_name, "precision": precision, + "status": "TIMEOUT", "error": "Worker exceeded the 1200s timeout", + "wall_time_s": time.time() - t0, + } + elapsed = time.time() - t0 + + if os.path.exists(output_file): + try: + with open(output_file) as f: + result = json.load(f) + except json.JSONDecodeError as e: + # A worker killed mid-write leaves truncated JSON. Treat it as a + # crash for this config instead of letting the exception escape and + # abort the whole sweep before any CSV is written. + return { + "method": method_name, "precision": precision, + "status": "CRASH", "error": f"Truncated/invalid result JSON: {e}", + "wall_time_s": elapsed, + } + result["wall_time_s"] = elapsed + return result + return { + "method": method_name, "precision": precision, + "status": "CRASH", "error": "No output file produced", + "wall_time_s": elapsed, + } + + +def main(): + parser = argparse.ArgumentParser(description="Evo2 Real Compression Benchmark") + parser.add_argument("--precisions", nargs="+", default=["int8"], choices=["int8"], + help="Precision families to test") + parser.add_argument("--eval-seq-len", type=int, default=2048, + help="Sequence length for quality evaluation") + parser.add_argument("--max-seqlen", action=argparse.BooleanOptionalAction, default=True, + help="Sweep candidate sequence lengths and report the largest that fits " + "(use --no-max-seqlen to skip)") + parser.add_argument("--output-dir", default="/workspace/quantization/results") + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + # Map precision families to representative method names + precision_methods = { + "int8": "INT8_DEFAULT_CFG", + } + + configs = [(p, precision_methods[p]) for p in args.precisions if p in precision_methods] + + print("=" * 80) + print(" Evo2 REAL Weight Compression Benchmark") + print(f" Configs: {[(p, m) for p, m in configs]}") + print(f" Eval seq_len: {args.eval_seq_len}") + print(f" Max seqlen search: {args.max_seqlen}") + print("=" * 80) + + all_results = [] + port = 29700 + + for precision, method_name in configs: + port += 1 + result = run_test(precision, method_name, args.eval_seq_len, + args.max_seqlen, port, args.output_dir) + + status = result.get("status", "?") + mem_bf16 = result.get("mem_bf16_mb", 0) + mem_comp = result.get("mem_compressed_mb", 0) + saved = result.get("mem_saved_pct", 0) + cos = result.get("cos_sim_avg", 0) + top1 = result.get("top1_agree", 0) + mse = result.get("mse", 0) + max_sl = result.get("max_seqlen", 0) + + print(f"\n ► {status}: {precision.upper()}") + print(f" Memory: {mem_bf16:.0f} → {mem_comp:.0f} MB (saved {saved:.1f}%)") + print(f" Quality: CosSim={cos:.6f} Top1={top1:.4f} MSE={mse:.6f}") + if max_sl > 0: + print(f" Max SeqLen: {max_sl:,}") + + if result.get("error"): + print(f" Error: {result['error'][:300]}") + + all_results.append(result) + + # Save CSV + csv_path = os.path.join(args.output_dir, "real_quant_benchmark.csv") + if all_results: + keys = sorted(set(k for r in all_results for k in r.keys())) + with open(csv_path, "w", newline="") as f: + w = csv.DictWriter(f, fieldnames=keys) + w.writeheader() + for r in all_results: + w.writerow({k: r.get(k, "") for k in keys}) + + # Summary + print(f"\n{'=' * 90}") + print(" REAL COMPRESSION RESULTS SUMMARY") + print(f"{'=' * 90}") + print(f" {'Precision':<8s} {'Method':<25s} {'BF16 MB':<10s} {'Comp MB':<10s} {'Saved%':<8s} " + f"{'CosSim':<10s} {'Top1':<8s} {'MaxSeqLen':<12s}") + print(f" {'-' * 85}") + for r in all_results: + print(f" {r.get('precision','?'):<8s} {r.get('method','?'):<25s} " + f"{r.get('mem_bf16_mb',0):>8.0f} {r.get('mem_compressed_mb',0):>8.0f} " + f"{r.get('mem_saved_pct',0):>6.1f} {r.get('cos_sim_avg',0):>8.6f} " + f"{r.get('top1_agree',0):>6.4f} {r.get('max_seqlen',0):>10,}") + + print(f"\n✅ Results saved to: {csv_path}") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/recipes/evo2_megatron_quant/scripts/download_models.sh b/recipes/evo2_megatron_quant/scripts/download_models.sh new file mode 100644 index 0000000000..7cbd229916 --- /dev/null +++ b/recipes/evo2_megatron_quant/scripts/download_models.sh @@ -0,0 +1,81 @@ +#!/bin/bash +# ============================================================================== +# Download Pretrained BioNeMo Models from NGC +# ============================================================================== +# +# Downloads all 3 pretrained models used for quantization testing. +# Models are cached locally by bionemo.core.data.load and will not be +# re-downloaded if already present. +# +# Models: +# 1. ESM-2 8M - Protein language model (Transformer, ~30MB) +# 2. Geneformer 10M - Gene expression model (Transformer, ~40MB) +# 3. Evo2 7B - DNA foundation model (Hyena/Mamba SSM, ~14GB) +# +# Prerequisites: +# - Must run INSIDE the BioNeMo container +# - NGC authentication configured +# +# Usage: +# bash scripts/download_models.sh # Download all models +# bash scripts/download_models.sh esm2 # Download ESM-2 only +# bash scripts/download_models.sh evo2 # Download Evo2 only +# ============================================================================== + +set -euo pipefail + +echo "================================================" +echo " BioNeMo Model Downloader" +echo "================================================" + +# Model tags on NGC +declare -A MODEL_TAGS=( + ["esm2"]="esm2/8m:2.0" + ["geneformer"]="geneformer/10M_241113:2.0" + ["evo2"]="evo2/7b-8k:1.0" +) + +# Model descriptions +declare -A MODEL_DESC=( + ["esm2"]="ESM-2 8M (Protein Transformer, ~30MB)" + ["geneformer"]="Geneformer 10M (Gene Expression Transformer, ~40MB)" + ["evo2"]="Evo2 7B (DNA Hyena/Mamba SSM, ~14GB)" +) + +# Determine which models to download +if [ $# -gt 0 ]; then + MODELS=("$@") +else + MODELS=("esm2" "geneformer" "evo2") +fi + +# Validate requested names up front so a typo fails fast instead of silently +# skipping while the script still exits 0. +for model in "${MODELS[@]}"; do + if [[ ! -v "MODEL_TAGS[$model]" ]]; then + echo " ❌ Unknown model: $model (available: esm2, geneformer, evo2)" >&2 + exit 1 + fi +done + +# Download each model +for model in "${MODELS[@]}"; do + tag="${MODEL_TAGS[$model]}" + desc="${MODEL_DESC[$model]}" + + echo "" + echo " 📥 Downloading: $desc" + echo " NGC tag: $tag" + echo "" + + python -c " +from bionemo.core.data.load import load +path = load('${tag}', source='ngc') +print(f' ✅ Downloaded to: {path}') +" +done + +echo "" +echo "================================================" +echo " All downloads complete!" +echo "================================================" diff --git a/recipes/evo2_megatron_quant/scripts/run_single_model.py b/recipes/evo2_megatron_quant/scripts/run_single_model.py new file mode 100644 index 0000000000..2c3257bdd1 --- /dev/null +++ b/recipes/evo2_megatron_quant/scripts/run_single_model.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +# 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. + +"""Run the ModelOpt PTQ harness on a single BioNeMo model. + +This is a manual GPU runner (needs a CUDA device and an NGC checkpoint), not a +pytest test — the CPU-only L0 test lives in +tests/test_compressed_linear_roundtrip.py. + +Loads the model once, gets a BF16 baseline, then quantizes in-place and +compares quality metrics. For each quant method, reloads a fresh model +to ensure clean state. Exits non-zero if any method FAILs or ERRORs. + +Usage: + # Single method + python scripts/run_single_model.py --model esm2 --quant INT8_DEFAULT_CFG + + # Multiple methods + python scripts/run_single_model.py --model esm2 \ + --quant INT8_DEFAULT_CFG,INT8_SMOOTHQUANT_CFG + + # Evo2 with all-MLP quantization + python scripts/run_single_model.py --model evo2 \ + --quant INT8_DEFAULT_CFG --all-mlp +""" + +import argparse +import gc +import os +import sys +import time +import traceback + +import torch + +# Add parent directory to path so we can import src modules +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.adapters import get_adapter, list_models +from src.quantize import quantize_model +from src.metrics import compute_metrics, print_summary, save_csv + + +def run_single_method( + adapter, model_name, ckpt_path, tokenizer, input_type, + input_ids, attention_mask, bf16_logits, quant_name, + enable_all_mlp=False, +): + """Test one quantization method by reloading the model fresh. + + This ensures clean state — no leftover quantizers from previous methods. + + Args: + adapter: ModelAdapter instance. + model_name: Model name string. + ckpt_path: Checkpoint path for reloading. + tokenizer: Tokenizer from first load (reused for data generation). + input_type: Data type string. + input_ids: Test input tensor. + attention_mask: Test attention mask. + bf16_logits: Baseline logits from unquantized model. + quant_name: Quantization method name. + enable_all_mlp: Whether to quantize all MLPs (Evo2). + + Returns: + Result dict with status, metrics, and timing info. + """ + result = { + "model": model_name, + "quant_method": quant_name, + "status": "ERROR", + "error": "", + "cos_sim_avg": 0.0, "cos_sim_min": 0.0, + "top1_agree": 0.0, "top5_overlap": 0.0, "mse": 0.0, + "quant_time_s": 0.0, "infer_time_ms": 0.0, + "num_quantizers": 0, + "mode": "all-mlp" if enable_all_mlp else "default", + } + + model_copy = None + try: + # Reload fresh model for this quant method + model_copy, _, _ = adapter.load_model(ckpt_path) + result["num_params_m"] = sum(p.numel() for p in model_copy.parameters()) / 1e6 + + # Quantize in-place + seq_len = 256 if model_name == "evo2" else 128 + q_info = quantize_model( + model_copy, quant_name, adapter, tokenizer, + enable_all_mlp=enable_all_mlp, + calib_seq_length=seq_len, + ) + result.update(q_info) + + # Forward pass with quantized model + t0 = time.time() + quant_logits = adapter.run_forward(model_copy, input_ids, attention_mask) + result["infer_time_ms"] = (time.time() - t0) * 1000 + + if quant_logits is None: + result["error"] = "Forward returned None" + return result + + # Compare with baseline + metrics = compute_metrics(bf16_logits, quant_logits) + result.update(metrics) + result["status"] = ( + "PASS" if metrics["cos_sim_avg"] >= 0.90 and metrics["top1_agree"] >= 0.50 + else "FAIL" + ) + + except Exception as e: + result["error"] = traceback.format_exc()[-200:] + result["status"] = "ERROR" + + finally: + if model_copy is not None: + del model_copy + gc.collect() + torch.cuda.empty_cache() + + return result + + +def main(): + parser = argparse.ArgumentParser( + description="Test quantization methods on a BioNeMo model", + ) + parser.add_argument( + "--model", type=str, required=True, choices=list_models(), + help="Model to test", + ) + parser.add_argument( + "--quant", type=str, default="INT8_DEFAULT_CFG", + help="Comma-separated list of quant methods (default: INT8_DEFAULT_CFG)", + ) + parser.add_argument( + "--all-mlp", action="store_true", + help="Quantize all MLP layers (relevant for Evo2 Hyena layers)", + ) + parser.add_argument( + "--output", type=str, default=None, + help="CSV output path (default: results/_quant.csv)", + ) + args = parser.parse_args() + + methods = [m.strip() for m in args.quant.split(",") if m.strip()] + if not methods: + parser.error("--quant did not contain any method names") + adapter = get_adapter(args.model) + + print(f"\n{'=' * 80}") + print(f" {adapter.description}") + print(f" Testing {len(methods)} quantization method(s)") + print(f" Mode: {'all-mlp' if args.all_mlp else 'default'}") + print(f"{'=' * 80}") + + # 1. Download checkpoint + print(f"\n Downloading {args.model} checkpoint...") + ckpt_path = adapter.download_checkpoint() + + # 2. Load model for BF16 baseline + print(f" Loading model for BF16 baseline...") + t0 = time.time() + model, tokenizer, input_type = adapter.load_model(ckpt_path) + num_params = sum(p.numel() for p in model.parameters()) / 1e6 + print(f" ✓ Loaded: {num_params:.1f}M params in {time.time() - t0:.1f}s") + + # 3. Generate test data + seq_length = 256 if args.model == "evo2" else 128 + input_ids, attention_mask = adapter.generate_test_data( + tokenizer, batch_size=4, seq_length=seq_length, + ) + print(f" ✓ Test data: {input_ids.shape}") + + # 4. BF16 baseline + print(f" Running BF16 baseline...") + bf16_logits = adapter.run_forward(model, input_ids, attention_mask) + print(f" ✓ BF16 output: {bf16_logits.shape}") + + # Free baseline model + del model + gc.collect() + torch.cuda.empty_cache() + + # 5. Test each method + results = [] + for i, qname in enumerate(methods): + print(f"\n [{i+1}/{len(methods)}] {qname}...", end=" ", flush=True) + result = run_single_method( + adapter, args.model, ckpt_path, tokenizer, input_type, + input_ids, attention_mask, bf16_logits, qname, + enable_all_mlp=args.all_mlp, + ) + icon = {"PASS": "✅", "FAIL": "❌"}.get(result["status"], "⚠️") + if result["status"] in ("PASS", "FAIL"): + print(f"{icon} cos={result['cos_sim_avg']:.4f} " + f"top1={result['top1_agree']*100:.1f}% " + f"t={result['quant_time_s']:.1f}s") + else: + print(f"{icon} {result.get('error', '')[:80]}") + results.append(result) + + # 6. Summary + print_summary(results, f"{args.model.upper()} — Quantization Results") + + # 7. Save CSV + output_path = args.output or f"results/{args.model}_quant.csv" + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + save_csv(results, output_path) + + # 8. Reflect outcome in the exit code so CI/automation can detect failures. + n_not_pass = sum(1 for r in results if r["status"] != "PASS") + return 1 if n_not_pass else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/evo2_megatron_quant/src/__init__.py b/recipes/evo2_megatron_quant/src/__init__.py new file mode 100644 index 0000000000..f839efe671 --- /dev/null +++ b/recipes/evo2_megatron_quant/src/__init__.py @@ -0,0 +1,20 @@ +# 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. + +"""BioNeMo-ModelOpt Quantization Bridge. + +This package provides utilities to quantize BioNeMo models using NVIDIA ModelOpt. +Supported models: ESM-2, Geneformer, Evo2. +""" diff --git a/recipes/evo2_megatron_quant/src/adapters.py b/recipes/evo2_megatron_quant/src/adapters.py new file mode 100644 index 0000000000..e98d9e1953 --- /dev/null +++ b/recipes/evo2_megatron_quant/src/adapters.py @@ -0,0 +1,493 @@ +#!/usr/bin/env python3 +# 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. + +"""Model adapters for BioNeMo models — Adapter Pattern. + +Each model has a dedicated adapter class that knows how to: + 1. Download its pretrained checkpoint + 2. Create/load the model on CUDA + 3. Build synthetic calibration data + 4. Prepare forward() arguments + 5. Parse the model output + +Supported models (3): + - ESM-2: Transformer encoder for protein sequences (NGC checkpoint) + - Geneformer: Transformer encoder for gene expression (NGC checkpoint) + - Evo2: Hyena/Mamba SSM for DNA sequences (NGC checkpoint) + +Usage: + from src.adapters import get_adapter, ADAPTER_REGISTRY + + adapter = get_adapter("esm2") + ckpt = adapter.download_checkpoint() + model = adapter.load_model(ckpt) + data = adapter.build_calibration_data(tokenizer, num_samples=8) +""" + +from __future__ import annotations + +import os +import random +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.utils.data + + +# ============================================================================ +# Constants +# ============================================================================ + +# NGC checkpoint tags for each model +MODEL_TAGS = { + "esm2": "esm2/8m:2.0", + "geneformer": "geneformer/10M_241113:2.0", + "evo2": "evo2/7b-8k:1.0", +} + + +# ============================================================================ +# Distributed environment initialization (required by Megatron-Core) +# ============================================================================ + +_DIST_INITIALIZED = False + + +def init_distributed(): + """Initialize single-GPU distributed environment for Megatron-Core. + + Sets up torch.distributed (NCCL backend) and Megatron's parallel state + with tensor/pipeline parallelism = 1. BioNeMo models require this even + for single-GPU inference because they use Megatron-Core internals + (ColumnParallelLinear, etc.). + + Safe to call multiple times — only the first call has effect. + """ + global _DIST_INITIALIZED + if _DIST_INITIALIZED: + return + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29500") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + + import torch.distributed as dist + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", world_size=1, rank=0) + + from megatron.core import parallel_state + if not parallel_state.is_initialized(): + parallel_state.initialize_model_parallel(1, 1) + + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + model_parallel_cuda_manual_seed(42) + + _DIST_INITIALIZED = True + + +# ============================================================================ +# Abstract Model Adapter +# ============================================================================ + +class ModelAdapter(ABC): + """Abstract base class for BioNeMo model adapters. + + Each concrete adapter encapsulates all model-specific logic: + loading, tokenization, data preparation, and forward pass handling. + + Attributes: + name: Short identifier (e.g., "esm2") + input_type: Data domain ("protein", "dna", "gene_ids") + architecture: Model architecture type ("transformer", "hyena_ssm") + description: Human-readable model description + """ + + name: str = "base" + input_type: str = "unknown" + architecture: str = "unknown" + description: str = "" + + # ------------------------------------------------------------------ + # Required methods (subclasses must implement) + # ------------------------------------------------------------------ + + @abstractmethod + def download_checkpoint(self) -> Optional[str]: + """Download pretrained checkpoint from NGC. + + Returns: + Local filesystem path to the checkpoint, or None if N/A. + """ + + @abstractmethod + def load_model(self, ckpt_path: Optional[str] = None) -> Tuple[torch.nn.Module, Any, str]: + """Load the model onto CUDA in eval mode. + + Args: + ckpt_path: Path to the pretrained checkpoint. + + Returns: + (model, tokenizer, input_type): + - model: nn.Module, bf16, CUDA, eval mode + - tokenizer: tokenizer instance (or None) + - input_type: "protein", "dna", or "gene_ids" + """ + + @abstractmethod + def generate_test_data( + self, tokenizer: Any, batch_size: int = 4, seq_length: int = 128, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Generate synthetic test data for this model. + + Args: + tokenizer: Tokenizer instance from load_model(). + batch_size: Number of sequences to generate. + seq_length: Maximum sequence length. + + Returns: + (input_ids, attention_mask): Both on CUDA, shape [B, L]. + """ + + @abstractmethod + def run_forward( + self, model: torch.nn.Module, input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + """Run model forward pass and extract logits. + + Args: + model: The BioNeMo model. + input_ids: Token IDs, shape [B, L]. + attention_mask: Attention mask, shape [B, L]. + + Returns: + Logits tensor, typically shape [B, L, V]. + """ + + @abstractmethod + def build_calibration_forward_args( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """Build keyword arguments for model.forward() during calibration. + + Args: + input_ids: Token IDs tensor. + attention_mask: Attention mask tensor. + + Returns: + Dict of kwargs to pass to model(**kwargs). + """ + + # ------------------------------------------------------------------ + # Shared helper: parse model output into logits + # ------------------------------------------------------------------ + + @staticmethod + def extract_logits(output) -> torch.Tensor: + """Extract logits from various model output formats. + + BioNeMo models return different output types: + - Raw Tensor + - Dict with 'token_logits' or 'logits' + - Named object with .token_logits or .logits attributes + """ + if isinstance(output, torch.Tensor): + return output + if isinstance(output, dict): + return output.get("token_logits", output.get("logits", None)) + for attr in ["token_logits", "logits"]: + if hasattr(output, attr): + return getattr(output, attr) + return output + + +# ============================================================================ +# ESM-2 Adapter (Protein Transformer) +# ============================================================================ + +class ESM2Adapter(ModelAdapter): + """Adapter for ESM-2 protein language model. + + ESM-2 is a masked language model for protein sequences based on the + Transformer encoder architecture. It uses the standard ESM tokenizer + with a 20-amino-acid vocabulary plus special tokens. + + Architecture: Transformer encoder (6 layers, 320 hidden, 20 heads for 8M) + Input: Amino acid sequences (e.g., "MKTLLILAVVAA...") + Checkpoint: NGC esm2/8m:2.0 + """ + + name = "esm2" + input_type = "protein" + architecture = "transformer" + description = "ESM-2 8M protein language model" + + def download_checkpoint(self) -> str: + from bionemo.core.data.load import load + return str(load(MODEL_TAGS["esm2"], source="ngc")) + + def load_model(self, ckpt_path: str) -> Tuple[torch.nn.Module, Any, str]: + init_distributed() + import nemo.lightning as nl + from bionemo.esm2.data.tokenizer import BioNeMoESMTokenizer + + ctx = nl.io.load_context(ckpt_path) + config = ctx.model.config + config.initial_ckpt_path = ckpt_path + + tokenizer = BioNeMoESMTokenizer() + model = config.configure_model(tokenizer) + model = model.bfloat16().cuda().eval() + + return model, tokenizer, self.input_type + + def generate_test_data(self, tokenizer, batch_size=4, seq_length=128): + rng = random.Random(42) + aa = "ACDEFGHIKLMNPQRSTVWY" + seqs = ["".join(rng.choices(aa, k=seq_length - 2)) for _ in range(batch_size)] + enc = tokenizer( + seqs, padding="max_length", max_length=seq_length, + truncation=True, return_tensors="pt", + ) + return enc["input_ids"].cuda(), enc["attention_mask"].cuda() + + def run_forward(self, model, input_ids, attention_mask): + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = model(input_ids=input_ids, attention_mask=attention_mask) + return self.extract_logits(output) + + def build_calibration_forward_args(self, input_ids, attention_mask): + return {"input_ids": input_ids, "attention_mask": attention_mask} + + +# ============================================================================ +# Geneformer Adapter (Gene Expression Transformer) +# ============================================================================ + +class GeneformerAdapter(ModelAdapter): + """Adapter for Geneformer single-cell gene expression model. + + Geneformer is a Transformer encoder pretrained on single-cell + transcriptomic data. It takes gene expression rank tokens as input + and predicts masked gene tokens. + + Special handling: + - Requires patching `load_settings_from_checkpoint` to load + outside of NeMo's full training pipeline. + - Uses a mock tokenizer with vocab_size=25429 for configure_model(). + - Test data is random integer gene IDs (not text sequences). + + Architecture: Transformer encoder (6 layers, 256 hidden, 4 heads for 10M) + Input: Gene expression rank tokens (integers 0-25428) + Checkpoint: NGC geneformer/10M_241113:2.0 + """ + + name = "geneformer" + input_type = "gene_ids" + architecture = "transformer" + description = "Geneformer 10M single-cell gene expression model" + + def download_checkpoint(self) -> str: + from bionemo.core.data.load import load + return str(load(MODEL_TAGS["geneformer"], source="ngc")) + + def load_model(self, ckpt_path: str) -> Tuple[torch.nn.Module, Any, str]: + init_distributed() + from bionemo.geneformer.api import GeneformerConfig + import bionemo.llm.model.config as bionemo_config + + # Patch: skip checkpoint settings loading (only needed for training) + original = ( + bionemo_config.MegatronBioNeMoTrainableModelConfig + .load_settings_from_checkpoint + ) + bionemo_config.MegatronBioNeMoTrainableModelConfig.load_settings_from_checkpoint = ( + lambda self, path: None + ) + + try: + class _GeneformerTokenizer: + """Minimal tokenizer stub for Geneformer's configure_model().""" + vocab_size = 25429 + def __len__(self): + return self.vocab_size + + config = GeneformerConfig( + num_layers=6, hidden_size=256, num_attention_heads=4, + ffn_hidden_size=512, initial_ckpt_path=ckpt_path, + ) + tokenizer = _GeneformerTokenizer() + model = config.configure_model(tokenizer) + model = model.bfloat16().cuda().eval() + finally: + bionemo_config.MegatronBioNeMoTrainableModelConfig.load_settings_from_checkpoint = original + + return model, None, self.input_type + + def generate_test_data(self, tokenizer, batch_size=4, seq_length=128): + # Geneformer uses integer gene IDs, not text + input_ids = torch.randint(0, 25426, (batch_size, seq_length), dtype=torch.long).cuda() + return input_ids, torch.ones_like(input_ids).cuda() + + def run_forward(self, model, input_ids, attention_mask): + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = model(input_ids=input_ids, attention_mask=attention_mask) + return self.extract_logits(output) + + def build_calibration_forward_args(self, input_ids, attention_mask): + return {"input_ids": input_ids, "attention_mask": attention_mask} + + +# ============================================================================ +# Evo2 Adapter (Hyena/Mamba SSM for DNA) +# ============================================================================ + +class Evo2Adapter(ModelAdapter): + """Adapter for Evo2 DNA foundation model. + + Evo2 uses a Hyena/Mamba state-space model (SSM) architecture — NOT a + standard Transformer. This has several important implications: + + Key differences from Transformer models: + 1. Requires position_ids in forward() call + 2. Uses Evo2Tokenizer with byte-level DNA vocabulary (A/C/G/T) + 3. Has NON-DETERMINISTIC initialization: the `configure_model()` + function initializes some weights randomly, and subsequent + `load_state_dict` doesn't fully override them due to CUDA RNG + state differences. This means loading the same checkpoint twice + in the same process can produce different models. + 4. For quantization testing, we use SUBPROCESS ISOLATION: each + quant method runs in its own process to ensure a clean load. + + Architecture: Hyena/Mamba SSM (32 layers, mix of Hyena + Attention) + Input: DNA nucleotide sequences (A/C/G/T) + Checkpoint: NGC evo2/7b-8k:1.0 + + Quantization notes: + - Default ModelOpt configs only quantize TE Linear layers in the + 5 attention layers (out of 32 total layers). + - Removing 'default: enable=False' quantizes all 32 MLP layers + (QuantMLP), including the 27 HyenaLayer MLPs. + - HyenaFilter's internal nn.Linear layers (inside nn.Sequential) + are NOT wrapped by ModelOpt because it doesn't traverse Sequential. + - ImplicitModalFilter (gamma, R, p) and ExplicitSingleDecayFilter (h) + are nn.Parameter, not Linear layers, and cannot be quantized. + """ + + name = "evo2" + input_type = "dna" + architecture = "hyena_ssm" + description = "Evo2 7B DNA foundation model (Hyena/Mamba SSM)" + + def download_checkpoint(self) -> str: + from bionemo.core.data.load import load + return str(load(MODEL_TAGS["evo2"], source="ngc")) + + def load_model(self, ckpt_path: str) -> Tuple[torch.nn.Module, Any, str]: + init_distributed() + import nemo.lightning as nl + from bionemo.evo2.data.tokenizer import Evo2Tokenizer + + ctx = nl.io.load_context(ckpt_path) + config = ctx.model.config + config.initial_ckpt_path = ckpt_path + + tokenizer = Evo2Tokenizer() + # Evo2 tokenizer requires explicit vocab_size setting + tokenizer.vocab_size = tokenizer.tokenizer.vocab_size + + model = config.configure_model(tokenizer) + model = model.bfloat16().cuda().eval() + + return model, tokenizer, self.input_type + + def generate_test_data(self, tokenizer, batch_size=4, seq_length=256): + rng = random.Random(42) + seqs = ["".join(rng.choices("ACGT", k=seq_length - 2)) for _ in range(batch_size)] + all_ids = [] + for s in seqs: + ids = tokenizer.tokenize(s) + # Evo2 tokenizer returns list-of-lists + if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list): + ids = ids[0] + ids = ids[:seq_length] + ids = ids + [0] * (seq_length - len(ids)) + all_ids.append(ids) + input_ids = torch.tensor(all_ids, dtype=torch.long).cuda() + return input_ids, (input_ids != 0).long().cuda() + + def run_forward(self, model, input_ids, attention_mask): + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .unsqueeze(0).expand_as(input_ids) + ) + output = model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + ) + return self.extract_logits(output) + + def build_calibration_forward_args(self, input_ids, attention_mask): + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .unsqueeze(0).expand_as(input_ids) + ) + return { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": attention_mask, + } + + +# ============================================================================ +# Adapter Registry +# ============================================================================ + +ADAPTER_REGISTRY: Dict[str, type] = { + "esm2": ESM2Adapter, + "geneformer": GeneformerAdapter, + "evo2": Evo2Adapter, +} + + +def get_adapter(model_name: str) -> ModelAdapter: + """Get the adapter instance for a given model name. + + Args: + model_name: One of "esm2", "geneformer", "evo2". + + Returns: + Instantiated ModelAdapter subclass. + + Raises: + ValueError: If model_name is not recognized. + """ + if model_name not in ADAPTER_REGISTRY: + raise ValueError( + f"Unknown model: '{model_name}'. " + f"Available: {list(ADAPTER_REGISTRY.keys())}" + ) + return ADAPTER_REGISTRY[model_name]() + + +def list_models() -> List[str]: + """Return list of supported model names.""" + return list(ADAPTER_REGISTRY.keys()) diff --git a/recipes/evo2_megatron_quant/src/compress_model.py b/recipes/evo2_megatron_quant/src/compress_model.py new file mode 100644 index 0000000000..7b7a9c41ff --- /dev/null +++ b/recipes/evo2_megatron_quant/src/compress_model.py @@ -0,0 +1,184 @@ +# 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. + +"""Compress an Evo2 model by replacing nn.Linear layers with real low-precision storage. + +Pipeline: + 1. Load Evo2 in BF16 + 2. (Optional) Run ModelOpt calibration to get optimal scales + 3. Walk all nn.Linear / TE Linear layers + 4. Replace with INT8Linear (real int8 weight storage) + 5. Return compressed model with real memory savings +""" + +import gc + +import torch +import torch.nn as nn + +try: + from .compressed_linear import METHOD_TO_PRECISION, compress_linear +except ImportError: # allow running with src/ placed directly on sys.path + from compressed_linear import METHOD_TO_PRECISION, compress_linear + + +def get_memory_mb(): + """Current GPU memory allocated in MB.""" + return torch.cuda.memory_allocated() / 1e6 + + +def compress_model(model: nn.Module, precision: str = "int8", + skip_patterns: tuple = ("embedding", "layernorm", "ln_", "norm"), + verbose: bool = True, **kwargs) -> dict: + """Replace all nn.Linear layers with compressed versions. + + Args: + model: PyTorch model (Evo2) + precision: Target precision (currently "int8") + skip_patterns: Module name patterns to skip (embeddings, norms) + verbose: Print compression stats + **kwargs: Extra args forwarded to the compressed linear constructor + + Returns: + dict with compression statistics + """ + stats = { + "precision": precision, + "total_linears": 0, + "compressed_linears": 0, + "skipped_linears": 0, + "mem_before_mb": get_memory_mb(), + "params_compressed": 0, + "params_skipped": 0, + } + + # Collect all (parent, name, linear) triples + replacements = [] + for name, module in model.named_modules(): + # Check for nn.Linear or TE Linear + is_linear = isinstance(module, nn.Linear) + if not is_linear: + # Check for TransformerEngine Linear + try: + import transformer_engine.pytorch as te + is_linear = isinstance(module, te.Linear) + except ImportError: + pass + + if not is_linear: + continue + + stats["total_linears"] += 1 + + # Skip patterns + name_lower = name.lower() + if any(p in name_lower for p in skip_patterns): + stats["skipped_linears"] += 1 + param_count = sum(p.numel() for p in module.parameters()) + stats["params_skipped"] += param_count + continue + + replacements.append((name, module)) + + if verbose: + print(f" Found {stats['total_linears']} Linear layers, " + f"compressing {len(replacements)}, " + f"skipping {stats['skipped_linears']}") + + # Do replacements — free old weights aggressively + for idx, (full_name, linear) in enumerate(replacements): + parts = full_name.split(".") + parent = model + for p in parts[:-1]: + parent = getattr(parent, p) + attr_name = parts[-1] + + param_count = linear.weight.numel() + + # Extract weight and bias from any Linear type + weight_data = linear.weight.data + bias_data = None + if hasattr(linear, 'bias') and linear.bias is not None: + b = linear.bias + if hasattr(b, 'data') and b.data.numel() > 0: + bias_data = b.data + param_count += b.data.numel() + + in_f = weight_data.shape[1] + out_f = weight_data.shape[0] + + # Create a wrapper nn.Linear on meta device, then assign weight/bias + has_bias = bias_data is not None + temp = nn.Linear(in_f, out_f, bias=has_bias, device='meta') + temp.weight = nn.Parameter(weight_data, requires_grad=False) + if has_bias: + temp.bias = nn.Parameter(bias_data, requires_grad=False) + + # Compress + compressed = compress_linear(temp, precision, **kwargs) + + # Free old weight memory + linear.weight.data = torch.empty(0, device='cpu') + if hasattr(linear, 'bias') and linear.bias is not None: + try: + linear.bias.data = torch.empty(0, device='cpu') + except Exception: + pass + del temp + + # Replace in model + setattr(parent, attr_name, compressed) + del linear, compressed + + stats["compressed_linears"] += 1 + stats["params_compressed"] += param_count + + # Periodic GC + if (idx + 1) % 8 == 0: + gc.collect() + torch.cuda.empty_cache() + + gc.collect() + torch.cuda.empty_cache() + + stats["mem_after_mb"] = get_memory_mb() + stats["mem_saved_mb"] = stats["mem_before_mb"] - stats["mem_after_mb"] + stats["mem_saved_pct"] = (stats["mem_saved_mb"] / stats["mem_before_mb"] * 100 + if stats["mem_before_mb"] > 0 else 0) + + if verbose: + print(f" Compressed {stats['compressed_linears']} layers ({stats['params_compressed']:,} params)") + print(f" Memory: {stats['mem_before_mb']:.0f} → {stats['mem_after_mb']:.0f} MB " + f"(saved {stats['mem_saved_mb']:.0f} MB, {stats['mem_saved_pct']:.1f}%)") + + return stats + + +def compress_from_method(model: nn.Module, method_name: str, + verbose: bool = True, **kwargs) -> dict: + """Compress model using precision determined by ModelOpt method name. + + Args: + model: PyTorch model + method_name: One of the 22 ModelOpt method names (e.g. "FP8_DEFAULT_CFG") + verbose: Print stats + """ + precision = METHOD_TO_PRECISION.get(method_name) + if precision is None: + raise ValueError(f"Unknown method: {method_name}. " + f"Known: {list(METHOD_TO_PRECISION.keys())}") + if verbose: + print(f" Method {method_name} → precision: {precision}") + return compress_model(model, precision, verbose=verbose, **kwargs) diff --git a/recipes/evo2_megatron_quant/src/compressed_linear.py b/recipes/evo2_megatron_quant/src/compressed_linear.py new file mode 100644 index 0000000000..22be7a5bf9 --- /dev/null +++ b/recipes/evo2_megatron_quant/src/compressed_linear.py @@ -0,0 +1,107 @@ +# 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. + +"""Compressed Linear modules: real low-precision weight storage. + +Provides a drop-in replacement for nn.Linear that stores weights in INT8, +achieving real memory savings (~2x on the compressed layers) while keeping +inference quality near-lossless via per-channel (per-output-row) scaling. + + INT8 -> torch.int8, dequant-on-the-fly to the activation dtype +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class INT8Linear(nn.Module): + """Linear layer with weights stored as int8. + + Dequantizes to the activation dtype on the fly for the matmul. + Per-channel (per-output-row) symmetric quantization preserves accuracy. + """ + + def __init__(self, weight_int8: torch.Tensor, scale: torch.Tensor, + bias: torch.Tensor = None, out_features: int = 0, in_features: int = 0, + return_bias: bool = True): + super().__init__() + self.out_features = out_features or weight_int8.shape[0] + self.in_features = in_features or weight_int8.shape[1] + self.return_bias = return_bias + self.register_buffer("weight_int8", weight_int8) # torch.int8 + self.register_buffer("scale", scale) # per-channel: (out_features, 1) + if bias is not None: + self.register_buffer("bias", bias) + else: + self.bias = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Dequantize: int8 -> activation dtype (per-output-row scale broadcasts over columns) + w = self.weight_int8.to(x.dtype) * self.scale.to(x.dtype) + out = F.linear(x, w, self.bias) + return (out, None) if self.return_bias else out + + def extra_repr(self): + return (f"in_features={self.in_features}, out_features={self.out_features}, " + f"bias={self.bias is not None}, dtype=int8") + + @staticmethod + def from_linear(linear: nn.Linear, scale: torch.Tensor = None, per_channel: bool = True, + return_bias: bool = True): + """Convert nn.Linear to INT8Linear. Quantizes on CPU to save GPU memory.""" + device = linear.weight.data.device + w = linear.weight.data.cpu().float() + if scale is None: + if per_channel: + amax = w.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) + else: + amax = w.abs().amax().unsqueeze(0).unsqueeze(0).clamp(min=1e-12) + scale = amax / 127.0 + w_int8 = (w / scale).round().clamp(-128, 127).to(torch.int8).to(device) + scale = scale.to(device) + bias = linear.bias.data.to(device) if linear.bias is not None else None + return INT8Linear(w_int8, scale, bias, linear.out_features, linear.in_features, + return_bias=return_bias) + + +# === Factory === + +PRECISION_MAP = { + "int8": INT8Linear, +} + +# Map the ModelOpt INT8 method names to the real-weight-compression precision. +METHOD_TO_PRECISION = { + "INT8_DEFAULT_CFG": "int8", + "INT8_SMOOTHQUANT_CFG": "int8", + "INT8_KV_CFG": "int8", + "MXINT8_DEFAULT_CFG": "int8", + "MX_INT8_CFG": "int8", +} + + +def compress_linear(linear: nn.Linear, precision: str = "int8", **kwargs): + """Convert nn.Linear to a compressed format. + + Args: + linear: Original linear layer. + precision: Currently only "int8" is supported. + **kwargs: Extra args passed to the compressed linear constructor. + """ + cls = PRECISION_MAP.get(precision) + if cls is None: + raise ValueError(f"Unknown precision: {precision}. Choose from: {list(PRECISION_MAP.keys())}") + return cls.from_linear(linear, **kwargs) diff --git a/recipes/evo2_megatron_quant/src/metrics.py b/recipes/evo2_megatron_quant/src/metrics.py new file mode 100644 index 0000000000..e0ed639062 --- /dev/null +++ b/recipes/evo2_megatron_quant/src/metrics.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +# 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. + +"""Quality metrics for comparing BF16 baseline vs. quantized model outputs. + +Metrics computed: + - Cosine Similarity (avg & min): measures directional alignment of logit vectors. + A value of 1.0 means the logits point in the exact same direction. + - Top-1 Agreement: fraction of positions where the predicted token matches. + Directly measures whether quantization changes the model's predictions. + - Top-5 Overlap: fraction of top-5 predicted tokens that overlap. + More lenient than top-1 — captures whether quantization preserves ranking. + - MSE: mean squared error between logit values. + Measures raw numerical difference (sensitive to scale). + +Usage: + from src.metrics import compute_metrics, print_summary, save_csv + + metrics = compute_metrics(bf16_logits, quant_logits) + print_summary(results) + save_csv(results, "results.csv") +""" + +import csv +from typing import Any, Dict, List + +import torch +import torch.nn.functional as F + + +def compute_metrics( + bf16_logits: torch.Tensor, + quant_logits: torch.Tensor, +) -> Dict[str, float]: + """Compare BF16 baseline logits against quantized model logits. + + Args: + bf16_logits: Baseline logits tensor, shape [B, L, V] or [B*L, V]. + quant_logits: Quantized model logits, same shape as bf16_logits. + + Returns: + Dictionary with keys: cos_sim_avg, cos_sim_min, top1_agree, + top5_overlap, mse. + """ + # Flatten to [N, vocab_size] for per-position comparison + bf = bf16_logits.reshape(-1, bf16_logits.shape[-1]).float() + qf = quant_logits.reshape(-1, quant_logits.shape[-1]).float() + + # Cosine similarity: per-position, then aggregate + cos = F.cosine_similarity(bf, qf, dim=-1) + + # Top-1 agreement: does the argmax match? + top1_bf = bf.argmax(dim=-1) + top1_q = qf.argmax(dim=-1) + + # Top-5 overlap: what fraction of top-5 tokens overlap? + top5_bf = bf.topk(5, dim=-1).indices + top5_q = qf.topk(5, dim=-1).indices + overlap = sum( + len(set(top5_bf[i].tolist()) & set(top5_q[i].tolist())) / 5.0 + for i in range(top5_bf.shape[0]) + ) / top5_bf.shape[0] + + return { + "cos_sim_avg": cos.mean().item(), + "cos_sim_min": cos.min().item(), + "top1_agree": (top1_bf == top1_q).float().mean().item(), + "top5_overlap": overlap, + "mse": F.mse_loss(bf, qf).item(), + } + + +def print_summary( + results: List[Dict[str, Any]], + title: str = "Quantization Results", +) -> None: + """Print a formatted summary table of quantization results. + + Args: + results: List of result dicts, each containing 'quant_method', + 'status', 'cos_sim_avg', 'top1_agree', etc. + title: Title for the summary table. + """ + print(f"\n{'=' * 100}") + print(f" {title}") + print(f"{'=' * 100}") + print( + f" {'Method':<35} {'Status':>7} {'CosSim':>8} {'Top1%':>7} " + f"{'Top5%':>7} {'MSE':>10} {'#Q':>4} {'Time':>7}" + ) + print( + f" {'-'*35} {'-'*7} {'-'*8} {'-'*7} " + f"{'-'*7} {'-'*10} {'-'*4} {'-'*7}" + ) + + for r in results: + if r["status"] in ("PASS", "FAIL"): + icon = "✅" if r["status"] == "PASS" else "❌" + nq = r.get("num_quantizers", r.get("num_quant_modules", 0)) + print( + f" {r['quant_method']:<35} {icon:>7} " + f"{r['cos_sim_avg']:>8.4f} " + f"{r['top1_agree']*100:>6.1f}% " + f"{r.get('top5_overlap', 0)*100:>6.1f}% " + f"{r['mse']:>10.6f} " + f"{nq:>4} " + f"{r.get('quant_time_s', 0):>6.1f}s" + ) + else: + err = r.get("error", "")[:40] + print( + f" {r['quant_method']:<35} {'⚠️':>7} " + f"{'—':>8} {'—':>7} {'—':>7} {'—':>10} {'—':>4} {err}" + ) + + print(f"{'=' * 100}") + n_pass = sum(1 for r in results if r["status"] == "PASS") + n_fail = sum(1 for r in results if r["status"] == "FAIL") + n_err = sum(1 for r in results if r["status"] == "ERROR") + print(f" Total: {n_pass} PASS, {n_fail} FAIL, {n_err} ERROR out of {len(results)}") + + +def save_csv( + results: List[Dict[str, Any]], + output_path: str, +) -> None: + """Save quantization results to a CSV file. + + Args: + results: List of result dictionaries. + output_path: Path to the output CSV file. + """ + if not results: + return + + fieldnames = [ + "model", "quant_method", "status", + "cos_sim_avg", "cos_sim_min", "top1_agree", "top5_overlap", "mse", + "num_quantizers", "num_quant_modules", + "quant_time_s", "infer_time_ms", "num_params_m", + "mode", "error", + ] + + with open(output_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + writer.writerows(results) + + print(f"\n 📄 Results saved to: {output_path}") diff --git a/recipes/evo2_megatron_quant/src/quantize.py b/recipes/evo2_megatron_quant/src/quantize.py new file mode 100644 index 0000000000..443d321536 --- /dev/null +++ b/recipes/evo2_megatron_quant/src/quantize.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +# 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. + +"""ModelOpt quantization wrapper for BioNeMo models. + +Provides a clean interface to NVIDIA ModelOpt's quantization APIs, including: + - All 22 built-in quantization configurations + - Custom config support (e.g., quantizing all MLPs for Evo2) + - Calibration forward loop construction from adapters + - In-place quantization with error handling + +Usage: + from src.quantize import quantize_model, ALL_QUANT_METHODS, get_quant_config + from src.adapters import get_adapter + + adapter = get_adapter("esm2") + model, tok, _ = adapter.load_model(ckpt) + quantize_model(model, "FP8_DEFAULT_CFG", adapter, tok) +""" + +from __future__ import annotations + +import copy +import time +from typing import Any, Callable, Dict, List + +import torch + +try: + import modelopt.torch.quantization as mtq +except ImportError as e: + raise ImportError( + f"NVIDIA ModelOpt not found: {e}\n" + "Install with: pip install nvidia-modelopt[all]" + ) + + +# ============================================================================ +# All 22 ModelOpt quantization configurations +# ============================================================================ + +ALL_QUANT_METHODS: List[str] = [ + # --- FP8 family (E4M3 / E5M2) --- + "FP8_DEFAULT_CFG", # FP8 E4M3 per-tensor + "FP8_PER_CHANNEL_PER_TOKEN_CFG", # FP8 per-channel weights, per-token activations + "FP8_KV_CFG", # FP8 with KV cache quantization + "FP8_AFFINE_KV_CFG", # FP8 with affine KV quantization + "FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG", # FP8 blockwise weight-only + + # --- INT8 family --- + "INT8_DEFAULT_CFG", # INT8 per-tensor symmetric + "INT8_SMOOTHQUANT_CFG", # INT8 SmoothQuant (migrates difficulty) + + # --- INT4 family --- + "INT4_AWQ_CFG", # INT4 Activation-aware Weight Quantization + "INT4_BLOCKWISE_WEIGHT_ONLY_CFG", # INT4 blockwise weight-only + + # --- NVFP4 family (NVIDIA FP4) --- + "NVFP4_DEFAULT_CFG", # NVFP4 default + "NVFP4_AWQ_LITE_CFG", # NVFP4 with AWQ lite calibration + "NVFP4_AWQ_CLIP_CFG", # NVFP4 with AWQ clip calibration + "NVFP4_AWQ_FULL_CFG", # NVFP4 with full AWQ calibration + "NVFP4_KV_CFG", # NVFP4 with KV cache quantization + "NVFP4_KV_ROTATE_CFG", # NVFP4 with rotated KV quantization + "NVFP4_AFFINE_KV_CFG", # NVFP4 with affine KV quantization + "NVFP4_SVDQUANT_DEFAULT_CFG", # NVFP4 with SVD quantization + + # --- MX (Microscaling) family --- + "MXFP4_DEFAULT_CFG", # Microscaling FP4 + "MXFP6_DEFAULT_CFG", # Microscaling FP6 + "MXFP8_DEFAULT_CFG", # Microscaling FP8 + "MXINT8_DEFAULT_CFG", # Microscaling INT8 + + # --- Mixed precision --- + "W4A8_AWQ_BETA_CFG", # 4-bit weights + 8-bit activations +] + +# Friendly name → ModelOpt constant name mapping +FRIENDLY_NAMES: Dict[str, str] = { + "fp8": "FP8_DEFAULT_CFG", + "fp8_per_channel": "FP8_PER_CHANNEL_PER_TOKEN_CFG", + "int8": "INT8_DEFAULT_CFG", + "int8_smoothquant": "INT8_SMOOTHQUANT_CFG", + "int4_awq": "INT4_AWQ_CFG", + "int4_blockwise": "INT4_BLOCKWISE_WEIGHT_ONLY_CFG", + "nvfp4": "NVFP4_DEFAULT_CFG", + "w4a8": "W4A8_AWQ_BETA_CFG", +} + + +def get_quant_config( + method_name: str, + enable_all_mlp: bool = False, +) -> dict: + """Get a ModelOpt quantization config by name. + + Args: + method_name: Either a ModelOpt constant name (e.g., "FP8_DEFAULT_CFG") + or a friendly name (e.g., "fp8"). Surrounding whitespace + is ignored. + enable_all_mlp: If True, removes the 'default: enable=False' entry + from the config to quantize ALL MLP layers + (relevant for Evo2's Hyena layers). + + Returns: + Deep copy of the quantization config dict. + + Raises: + ValueError: If the method name is not recognized. + """ + # Normalize, then resolve a friendly name if one was given. + method_name = method_name.strip() + resolved = FRIENDLY_NAMES.get(method_name, method_name) + + # Only whitelisted configs may reach ModelOpt: getattr() on an arbitrary + # string could otherwise return an unrelated mtq attribute. + if resolved not in ALL_QUANT_METHODS: + raise ValueError( + f"Unknown quantization method: '{method_name}'. " + f"Available: {ALL_QUANT_METHODS}" + ) + + cfg = getattr(mtq, resolved, None) + if cfg is None: + raise ValueError( + f"'{resolved}' is not available in this ModelOpt build. " + f"Available: {ALL_QUANT_METHODS}" + ) + + cfg = copy.deepcopy(cfg) + + if enable_all_mlp: + # ModelOpt configs carry a 'default: {enable: False}' entry that stops + # modules not matched by an explicit wildcard from being quantized. + # Dropping it lets all MLP layers be quantized, including the Hyena + # layers' TE ColumnParallelLinear / RowParallelLinear. + # + # This assumes quant_cfg is a mapping. If a ModelOpt version uses a + # different layout we fail loudly rather than silently doing nothing: + # a silent no-op would report "all-MLP" numbers for a default-scope run. + quant_cfg = cfg.get("quant_cfg") + if isinstance(quant_cfg, dict) and "default" in quant_cfg: + del quant_cfg["default"] + else: + raise RuntimeError( + f"enable_all_mlp: no 'default' entry found in quant_cfg " + f"(type={type(quant_cfg).__name__}) for '{resolved}'. This " + "ModelOpt version likely uses a different config layout, so " + "--all-mlp would have no effect; refusing to report misleading " + "all-MLP results." + ) + + return cfg + + +def build_calibration_loop( + adapter: Any, + tokenizer: Any, + model_name: str, + num_batches: int = 4, + batch_size: int = 2, + seq_length: int = 128, +) -> Callable: + """Build a calibration forward loop function for mtq.quantize(). + + ModelOpt's mtq.quantize() requires a callable that drives the model + through representative inputs so it can collect activation statistics + for calibration (needed by methods like SmoothQuant, AWQ, etc.). + + Args: + adapter: ModelAdapter instance. + tokenizer: Tokenizer from adapter.load_model(). + model_name: Model name string. + num_batches: Number of calibration batches. + batch_size: Samples per batch. + seq_length: Sequence length per sample. + + Returns: + Callable that takes a model and runs calibration forward passes. + """ + # Pre-generate calibration data + calib_batches = [] + for _ in range(num_batches): + ids, mask = adapter.generate_test_data(tokenizer, batch_size, seq_length) + fwd_args = adapter.build_calibration_forward_args(ids, mask) + calib_batches.append(fwd_args) + + def forward_loop(model): + """Drive model through calibration data for quantizer statistics.""" + failures = 0 + for kwargs in calib_batches: + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + try: + model(**kwargs) + except Exception as e: + failures += 1 + print(f" [calib] forward pass failed: {type(e).__name__}: {e}") + if failures == len(calib_batches): + print( + f" [calib] WARNING: all {failures} calibration passes failed; " + "quantization scales are uncalibrated." + ) + + return forward_loop + + +def quantize_model( + model: torch.nn.Module, + method_name: str, + adapter: Any, + tokenizer: Any, + enable_all_mlp: bool = False, + num_calib_batches: int = 4, + calib_batch_size: int = 2, + calib_seq_length: int = 128, +) -> Dict[str, Any]: + """Quantize a model in-place using the specified ModelOpt method. + + Args: + model: BioNeMo model (will be modified in-place). + method_name: Quantization method (e.g., "FP8_DEFAULT_CFG"). + adapter: ModelAdapter instance for data generation. + tokenizer: Tokenizer for generating calibration data. + enable_all_mlp: If True, quantize all MLP layers (for Evo2). + num_calib_batches: Number of calibration batches. + calib_batch_size: Batch size for calibration. + calib_seq_length: Sequence length for calibration. + + Returns: + Dict with 'quant_time_s'. + """ + cfg = get_quant_config(method_name, enable_all_mlp=enable_all_mlp) + + forward_loop = build_calibration_loop( + adapter=adapter, + tokenizer=tokenizer, + model_name=adapter.name, + num_batches=num_calib_batches, + batch_size=calib_batch_size, + seq_length=calib_seq_length, + ) + + t0 = time.time() + mtq.quantize(model, cfg, forward_loop=forward_loop) + quant_time = time.time() - t0 + + return { + "quant_time_s": quant_time, + } diff --git a/recipes/evo2_megatron_quant/tests/test_compressed_linear_roundtrip.py b/recipes/evo2_megatron_quant/tests/test_compressed_linear_roundtrip.py new file mode 100644 index 0000000000..aaf2301803 --- /dev/null +++ b/recipes/evo2_megatron_quant/tests/test_compressed_linear_roundtrip.py @@ -0,0 +1,74 @@ +# 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. + +"""CPU-only round-trip tests for INT8Linear (L0 sanity). + +These require neither a GPU nor a model checkpoint, so they run in normal CI. +They exercise the quantize -> store -> dequantize path of INT8Linear on tiny +random Linear layers and assert that reconstruction is near-lossless and that +the stored buffer is genuinely smaller than BF16. +""" + +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +# Import the module under test (recipe layout: ../src) +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) +from compressed_linear import INT8Linear # noqa: E402 + + +def _reference_linear(out_features: int, in_features: int, bias: bool = True) -> nn.Linear: + torch.manual_seed(0) + lin = nn.Linear(in_features, out_features, bias=bias) + # Small, realistic weight magnitudes (like a trained transformer projection). + with torch.no_grad(): + lin.weight.mul_(0.03) + return lin.eval() + + +@pytest.mark.parametrize("out_f,in_f", [(16, 256), (8, 512), (32, 129)]) +def test_int8_near_lossless(out_f, in_f): + lin = _reference_linear(out_f, in_f) + q = INT8Linear.from_linear(lin, return_bias=False) + x = torch.randn(4, in_f) + y_ref = lin(x) + y_q = q(x) + cos = torch.cosine_similarity(y_ref.flatten(), y_q.flatten(), dim=0).item() + assert cos > 0.999, f"INT8 cosine too low: {cos}" + + +def test_int8_no_systematic_bias(): + """Per-channel symmetric INT8 should not introduce a directional weight bias.""" + lin = _reference_linear(16, 256) + w_rec = (q := INT8Linear.from_linear(lin)).weight_int8.float() * q.scale.float() + mean_signed_err = (w_rec - lin.weight.data).mean().item() + signal = lin.weight.data.abs().mean().item() + assert abs(mean_signed_err) < 0.05 * signal, ( + f"INT8 has an unexpected systematic bias of {mean_signed_err:.5f} (signal {signal:.5f})" + ) + + +def test_int8_memory_footprint_shrinks(): + """The stored int8 weight buffer must be ~half of BF16 (plus a tiny per-row scale).""" + lin = _reference_linear(64, 512) + bf16_bytes = lin.weight.numel() * 2 # bf16 = 2 bytes/param + q = INT8Linear.from_linear(lin) + int8_bytes = q.weight_int8.numel() * q.weight_int8.element_size() + scale_bytes = q.scale.numel() * q.scale.element_size() + assert int8_bytes + scale_bytes < 0.55 * bf16_bytes