[Feature] Inference-side PTQ / weight-compression recipe for bio models (ESM-2 / Geneformer / Evo2) — wanted? where should it live?

[huthvincent]

Problem & Motivation

bionemo-recipes already has solid training-time low-precision support (TransformerEngine FP8 / MXFP8 / NVFP4 — e.g. #1484, #1500, recipes/fp8_analysis, recipes/geneformer_native_te_mfsdp_fp8). What I couldn't find anywhere in the repo is an inference-side post-training quantization (PTQ) / weight-compression example:

• nothing wiring the bio models to NVIDIA ModelOpt (mtq.quantize) for PTQ,
• no real low-precision weight storage (actual memory reduction) or KV-cache compression for inference,
• recipes/vllm_inference/esm2 runs inference but does no quantization.

I've been prototyping exactly this on a fork and would like to contribute it back. Before I invest in making it fully recipe-compliant, I'd rather ask whether it's wanted and where you'd want it to live than open a large PR blind.

BioNeMo Recipes Version

Forked off main around commit 4e54aaa; I'll rebase onto current main before opening any PR.

Category

Inference

Proposed Solution

A self-contained recipe that provides:

1. ModelOpt PTQ bridge (simulated Q/DQ) — an adapter connecting BioNeMo model loading to mtq.quantize(), plus an automated quality harness comparing BF16 baseline vs. quantized logits (cosine similarity / top-k agreement / MSE). Covers ESM-2, Geneformer, and Evo2 (Evo2 needs subprocess isolation because of its Hyena/SSM non-deterministic init).
2. Real low-precision weight storage — drop-in nn.Linear replacements that store weights at INT8/FP8/INT4 for genuine memory savings (dequant-on-the-fly). INT8 per-channel is the part that currently holds up; FP8/INT4 are experimental (honest results below).

Guidance I'm looking for:

1. Do you want an inference PTQ / weight-compression recipe in this repo at all, or is low-precision inference intended to go through the NeMo + TensorRT-LLM export path outside bionemo-recipes?
2. Where should it live? A new recipes/{model}_{framework}_quant/ (e.g. recipes/evo2_megatron_quant/), or folded into an existing recipe such as recipes/evo2_megatron/ / recipes/esm2_native_te/?
3. Preferred scope — the ModelOpt PTQ bridge, the real INT8 weight compression, or both?

Expected Benefits

• INT8 per-channel real weight compression on Evo2 7B: ~15% model-memory reduction, near-lossless in my tests (cosine sim ≈ 0.998, top-1 agreement 100%).
• A reusable, model-agnostic PTQ quality harness for the bio models.

To be upfront about the negatives (I'd document these rather than hide them, and am happy to ship only the parts that hold up): a naive per-tensor FP8 weight path and my INT4 path were low-quality in my experiments, and external KV-cache compression didn't help Evo2 (only ~5/32 layers use attention; the memory bottleneck is Hyena conv intermediates, and FlashAttention manages K/V internally). So a realistic first contribution might be just INT8 weight compression + the ModelOpt PTQ analysis harness.

Code Example

# 1) ModelOpt PTQ bridge (simulated) + automated quality eval
model  = adapter.load_model("esm2/8m:2.0")
qmodel = quantize(model, "INT8_DEFAULT_CFG", calib_data=adapter.calib())
metrics = compare(model, qmodel)   # -> {cosine_sim, top1_agree, top5_overlap, mse}

# 2) Real INT8 weight compression (actual GPU memory saving)
stats = compress_model(evo2_model, precision="int8")   # per-channel; swaps Linear -> INT8Linear
# -> ~15% GPU memory reduction on Evo2 7B, cosine_sim ~0.998 vs BF16

Happy to align with your recipe conventions (self-contained dir, Dockerfile, pinned deps, hydra config, L0 sanity test, SPDX headers, DCO sign-off) once you point me at the preferred home and scope. Thanks!