feat(recipes): add evo2_megatron_quant — INT8 weight compression + ModelOpt PTQ harness#1682
feat(recipes): add evo2_megatron_quant — INT8 weight compression + ModelOpt PTQ harness#1682huthvincent wants to merge 3 commits into
Conversation
…lOpt PTQ harness Adds an inference-side recipe: real INT8 per-channel weight compression (INT8Linear, ~15% GPU-mem saving on Evo2 7B, cosine sim ~0.998) plus a ModelOpt mtq.quantize PTQ analysis harness for ESM-2/Geneformer/Evo2. Includes a CPU-only L0 sanity test (no GPU/checkpoint needed). Follows up on NVIDIA-BioNeMo#1681. Signed-off-by: Rui Zhu <huthruiz@gmail.com>
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds a BioNeMo quantization recipe with model adapters, ModelOpt PTQ support, real INT8 linear compression, evaluation metrics, benchmarking, model download utilities, and container setup for ESM-2, Geneformer, and Evo2. ChangesBioNeMo quantization recipe
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI as test_single_model
participant Adapter as ModelAdapter
participant Model as BioNeMo model
participant Quantizer as quantize_model
participant Metrics as compute_metrics
CLI->>Adapter: load model and generate test data
Adapter->>Model: run BF16 inference
CLI->>Quantizer: quantize a fresh model
Quantizer->>Model: apply ModelOpt configuration
CLI->>Model: run quantized inference
CLI->>Metrics: compare baseline and quantized logits
Metrics-->>CLI: return quality metrics
Possibly related issues
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 16
🧹 Nitpick comments (3)
recipes/evo2_megatron_quant/src/adapters.py (2)
260-276: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
run_forward/build_calibration_forward_argsbetween ESM2Adapter and GeneformerAdapter.Both classes have byte-for-byte identical
run_forward(Lines 270-273, 348-351) andbuild_calibration_forward_args(Lines 275-276, 353-354) bodies. Consider hoisting default implementations intoModelAdapter(onlyEvo2Adapterneeds to override them forposition_ids), removing the duplication.♻️ Proposed refactor sketch
class ModelAdapter(ABC): ... + def run_forward(self, model, input_ids, attention_mask): + """Default forward pass shared by transformer-style adapters.""" + with torch.no_grad(), torch.cuda.amp.autocast(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}Then remove the now-redundant overrides from
ESM2AdapterandGeneformerAdapter, keepingEvo2Adapter's overrides (it needsposition_ids).Also applies to: 343-354
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@recipes/evo2_megatron_quant/src/adapters.py` around lines 260 - 276, Hoist the shared run_forward and build_calibration_forward_args implementations into ModelAdapter, preserving their current behavior and signatures. Remove the redundant overrides from ESM2Adapter and GeneformerAdapter, while retaining Evo2Adapter’s overrides because it requires position_ids.
271-271: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace deprecated
torch.cuda.amp.autocastcallsPyTorch 2.4 deprecates
torch.cuda.amp.autocast; switch these four sites inrecipes/evo2_megatron_quant/src/adapters.pyandrecipes/evo2_megatron_quant/src/quantize.pytotorch.amp.autocast("cuda", dtype=torch.bfloat16)in one pass.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@recipes/evo2_megatron_quant/src/adapters.py` at line 271, Replace all four deprecated torch.cuda.amp.autocast usages in the relevant adapter and quantization code with torch.amp.autocast("cuda", dtype=torch.bfloat16), preserving the surrounding torch.no_grad contexts and autocast behavior.recipes/evo2_megatron_quant/src/compressed_linear.py (1)
86-93: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueAlign the INT8 aliases with the public method list
INT8_KV_CFGandMX_INT8_CFGonly exist inMETHOD_TO_PRECISION; they’re not listed inquant_methods.yamlorALL_QUANT_METHODS. If they’re intentional aliases, add them to the catalog/docs; otherwise drop them or switch to the canonicalMXINT8_DEFAULT_CFGname to keep the supported set consistent.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@recipes/evo2_megatron_quant/src/compressed_linear.py` around lines 86 - 93, Update METHOD_TO_PRECISION so its INT8 aliases match the supported methods in quant_methods.yaml and ALL_QUANT_METHODS: remove INT8_KV_CFG and MX_INT8_CFG unless they are formally added to those catalogs, and use the canonical MXINT8_DEFAULT_CFG name where applicable.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@recipes/evo2_megatron_quant/docker/start_container.sh`:
- Line 22: Update the LOCAL_WORKSPACE default in start_container.sh to mount the
current recipe root rather than appending /quantization to the working
directory. Derive it from the script’s parent directory or use the current
directory when the script is run from the recipe root, while preserving any
explicitly provided LOCAL_WORKSPACE value.
- Line 44: Update the start_container.sh Docker launch options so --net=host is
applied only when an explicit networking flag is enabled, while preserving host
networking for NCCL/distributed runs that require it. Keep the default
invocation on isolated container networking.
In `@recipes/evo2_megatron_quant/README.md`:
- Line 44: Update the fenced layout block in the README to specify the text
language identifier, preserving its existing tree-style contents so markdownlint
MD040 passes.
In `@recipes/evo2_megatron_quant/requirements.txt`:
- Around line 9-14: Pin the dependencies in
recipes/evo2_megatron_quant/requirements.txt, specifically
nvidia-modelopt[torch], pyyaml, and pytest, to exact versions matching the
target BioNeMo image; alternatively, add and reference a matching constraints
file so installations remain reproducible.
In `@recipes/evo2_megatron_quant/scripts/benchmark_real_quant.py`:
- Around line 167-203: Replace the fixed candidate sweep in the
max_seqlen_search block with a bounded binary search that records the highest
successful length and the first failing upper bound, then probes between them
until the maximum supported length is identified. Preserve the existing result
fields and peak-memory reporting, or rename the option/output consistently if
retaining a candidate sweep instead.
- Around line 263-264: Update the argparse definition for --max-seqlen so the
expensive max sequence-length search can be explicitly disabled, using a paired
opt-out flag or equivalent boolean configuration while preserving the current
enabled-by-default behavior and the existing args.max_seqlen consumer.
- Around line 118-120: Update the sequence-probe flow around the warmup call and
its try/finally cleanup: initialize ids, pos, and mask before entering try so
cleanup remains safe when allocation fails, explicitly delete the warmup output
and each probe output instead of retaining them in _, and perform all
probe-variable cleanup in finally before empty_cache() to preserve the original
allocation error and release memory.
- Around line 217-254: Update run_test() to delete any existing output_file
before launching the worker, preventing stale JSON from being reused when the
subprocess fails. Wrap subprocess.run() to catch subprocess.TimeoutExpired and
return the existing CRASH-style fallback result with an appropriate timeout
error and elapsed wall time, while preserving normal result loading and failure
handling.
In `@recipes/evo2_megatron_quant/scripts/download_models.sh`:
- Around line 54-57: Update the unknown-model branch in the model validation
loop to fail with a nonzero status instead of continuing successfully. Preserve
the existing warning, and ensure an invalid requested model prevents the script
from reporting success or proceeding with downloads.
In `@recipes/evo2_megatron_quant/src/__init__.py`:
- Around line 16-20: Update the package docstring in __init__.py to remove
AMPLIFY from the supported-model list, matching the models registered by
ADAPTER_REGISTRY and the adapters.py documentation; do not add adapter support.
In `@recipes/evo2_megatron_quant/src/compress_model.py`:
- Line 31: Update the import in compress_model.py to use a package-relative
import from the sibling compressed_linear module, while preserving the existing
METHOD_TO_PRECISION and compress_linear names.
In `@recipes/evo2_megatron_quant/src/quantize.py`:
- Around line 179-186: Update forward_loop to log each exception raised by
model(**kwargs), including useful error details and batch context, instead of
silently passing. Preserve the existing behavior of continuing calibration after
individual batch failures.
- Around line 227-234: Update the result construction after mtq.quantize in the
quantization flow so num_quantizers is not derived from the returned model via
the isinstance check. Remove the num_quantizers field unless an existing,
authoritative quantizer-count source is available, in which case populate it
from that source.
In `@recipes/evo2_megatron_quant/tests/test_compressed_linear_roundtrip.py`:
- Around line 48-52: Update the cosine-similarity test around
INT8Linear.from_linear so both forward paths use equivalent bias behavior:
either enable the quantized bias by removing return_bias=False or construct a
bias-free reference matching the quantized layer. Keep y_ref, y_q, and the
cosine comparison otherwise unchanged.
In `@recipes/evo2_megatron_quant/tests/test_single_model.py`:
- Around line 211-221: Update main() to return a nonzero status when the
quantization results contain any FAIL or ERROR outcomes, while retaining zero
for successful runs. Propagate that return value through the __main__ entry
point using the process exit mechanism so CI observes failures.
- Around line 53-56: Rename the helper function test_single_method to
run_single_method so pytest does not collect it, and update its call site in the
same module to use the new name.
---
Nitpick comments:
In `@recipes/evo2_megatron_quant/src/adapters.py`:
- Around line 260-276: Hoist the shared run_forward and
build_calibration_forward_args implementations into ModelAdapter, preserving
their current behavior and signatures. Remove the redundant overrides from
ESM2Adapter and GeneformerAdapter, while retaining Evo2Adapter’s overrides
because it requires position_ids.
- Line 271: Replace all four deprecated torch.cuda.amp.autocast usages in the
relevant adapter and quantization code with torch.amp.autocast("cuda",
dtype=torch.bfloat16), preserving the surrounding torch.no_grad contexts and
autocast behavior.
In `@recipes/evo2_megatron_quant/src/compressed_linear.py`:
- Around line 86-93: Update METHOD_TO_PRECISION so its INT8 aliases match the
supported methods in quant_methods.yaml and ALL_QUANT_METHODS: remove
INT8_KV_CFG and MX_INT8_CFG unless they are formally added to those catalogs,
and use the canonical MXINT8_DEFAULT_CFG name where applicable.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 145379d8-da58-44a0-8994-e3656515d582
📒 Files selected for processing (15)
recipes/evo2_megatron_quant/.ruff.tomlrecipes/evo2_megatron_quant/README.mdrecipes/evo2_megatron_quant/configs/quant_methods.yamlrecipes/evo2_megatron_quant/docker/start_container.shrecipes/evo2_megatron_quant/requirements.txtrecipes/evo2_megatron_quant/scripts/benchmark_real_quant.pyrecipes/evo2_megatron_quant/scripts/download_models.shrecipes/evo2_megatron_quant/src/__init__.pyrecipes/evo2_megatron_quant/src/adapters.pyrecipes/evo2_megatron_quant/src/compress_model.pyrecipes/evo2_megatron_quant/src/compressed_linear.pyrecipes/evo2_megatron_quant/src/metrics.pyrecipes/evo2_megatron_quant/src/quantize.pyrecipes/evo2_megatron_quant/tests/test_compressed_linear_roundtrip.pyrecipes/evo2_megatron_quant/tests/test_single_model.py
| 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 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant area first.
git ls-files recipes/evo2_megatron_quant
echo
echo "== requirements.txt =="
cat -n recipes/evo2_megatron_quant/requirements.txt
echo
echo "== nearby files =="
git ls-files recipes/evo2_megatron_quant | sed 's#^`#-` #'
echo
echo "== search for constraints / pins / matching version files =="
rg -n --hidden --glob '!**/.git/**' \
-e 'constraints' \
-e 'requirement' \
-e 'nvidia-modelopt\[torch\]' \
-e '^pytest([<=>~!].*)?$' \
-e '^pyyaml([<=>~!].*)?$' \
recipes . || trueRepository: NVIDIA-BioNeMo/bionemo-recipes
Length of output: 27633
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== recipes/evo2_megatron_quant/README.md (install guidance) =="
sed -n '1,140p' recipes/evo2_megatron_quant/README.md
echo
echo "== recipes/evo2_megatron_quant/docker/start_container.sh =="
cat -n recipes/evo2_megatron_quant/docker/start_container.sh
echo
echo "== recipes/README.md (requirements guidance) =="
sed -n '90,115p' recipes/README.md
echo
echo "== .github/workflows/unit-tests-recipes.yml install snippet =="
sed -n '160,185p' .github/workflows/unit-tests-recipes.ymlRepository: NVIDIA-BioNeMo/bionemo-recipes
Length of output: 6588
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== exact paths for recipe guidance files =="
git ls-files | rg '(^|/)recipes/README\.md$|(^|/)evo2_megatron_quant/requirements\.txt$|(^|/)requirements\.txt$' || true
echo
echo "== top-level recipe guidance =="
sed -n '95,115p' bionemo-recipes/recipes/README.md
echo
echo "== recipe install paths in CI =="
sed -n '165,182p' .github/workflows/unit-tests-recipes.yml
echo
echo "== comparable recipe requirements pins =="
for f in \
bionemo-recipes/recipes/esm2_native_te/requirements.txt \
bionemo-recipes/recipes/codonfm_ptl_te/requirements.txt \
bionemo-recipes/recipes/evo2_megatron/requirements.txt
do
echo "--- $f ---"
sed -n '1,40p' "$f"
echo
doneRepository: NVIDIA-BioNeMo/bionemo-recipes
Length of output: 3667
Pin the recipe dependencies recipes/evo2_megatron_quant/requirements.txt still leaves nvidia-modelopt[torch], pyyaml, and pytest unconstrained, so pip install -r requirements.txt can drift from the target BioNeMo image and hurt reproducibility. Pin exact versions here or add a matching constraints file.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@recipes/evo2_megatron_quant/requirements.txt` around lines 9 - 14, Pin the
dependencies in recipes/evo2_megatron_quant/requirements.txt, specifically
nvidia-modelopt[torch], pyyaml, and pytest, to exact versions matching the
target BioNeMo image; alternatively, add and reference a matching constraints
file so installations remain reproducible.
| 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() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Compare equivalent forward paths in the cosine test.
return_bias=False removes the quantized bias, while lin(x) includes it. The test therefore compares different computations; keep the bias enabled or make the reference bias-free.
Proposed fix
- q = INT8Linear.from_linear(lin, return_bias=False)
+ q = INT8Linear.from_linear(lin)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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() | |
| q = INT8Linear.from_linear(lin) | |
| 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() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@recipes/evo2_megatron_quant/tests/test_compressed_linear_roundtrip.py` around
lines 48 - 52, Update the cosine-similarity test around INT8Linear.from_linear
so both forward paths use equivalent bias behavior: either enable the quantized
bias by removing return_bias=False or construct a bias-free reference matching
the quantized layer. Keep y_ref, y_q, and the cosine comparison otherwise
unchanged.
- test_single_model: rename helper so pytest doesn't collect it; propagate FAIL/ERROR via exit code; default to INT8 methods - compress_model: package-relative import with top-level fallback - quantize: surface (not swallow) calibration forward failures; drop the always-zero num_quantizers field - benchmark_real_quant: exception-safe seq-len sweep, free warmup output, make --max-seqlen disableable, drop stale result file, handle timeout, stop calling the linear sweep a binary search - start_container: mount the recipe root (not empty ./quantization), make --net=host opt-in via USE_HOST_NET - download_models: fail fast (nonzero) on an unknown model name - __init__: drop AMPLIFY from the supported-models docstring - README: language-tag the layout fence (MD040); requirements: pin pyyaml/pytest Signed-off-by: Rui Zhu <huthruiz@gmail.com>
|
Thanks for the thorough review! Addressed the findings in Fixed
One I believe is a false positive
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
recipes/evo2_megatron_quant/src/quantize.py (2)
135-142: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHandle ModelOpt’s list-based
quant_cfghere.
ModelOpt 0.45+ storesquant_cfgas an ordered list of entries, soif "default" in quant_cfgnever matches andenable_all_mlpbecomes a no-op. Remove the wildcard disable entry (quantizer_name: "*") or add an explicit MLP override instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@recipes/evo2_megatron_quant/src/quantize.py` around lines 135 - 142, Update the enable_all_mlp quant_cfg handling to support ModelOpt’s ordered list format in addition to the existing mapping format. Locate and remove the list entry whose quantizer_name is "*" and whose configuration disables matching modules, so enabling all MLP layers remains effective; preserve unrelated quant_cfg entries and existing behavior for mapping-based configurations.
123-131: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winTrim and whitelist quant method tokens before lookup.
--quantis split on commas without stripping, so values likeINT8_DEFAULT_CFG, INT8_SMOOTHQUANT_CFGfail on the leading space. Gate the resolved name againstALL_QUANT_METHODSbefore callinggetattr()so only supported configs reach ModelOpt.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@recipes/evo2_megatron_quant/src/quantize.py` around lines 123 - 131, Update the quantization-method resolution flow around FRIENDLY_NAMES, ALL_QUANT_METHODS, and getattr(mtq, resolved, None) to trim whitespace from each method token before lookup, then validate the resolved name against ALL_QUANT_METHODS before calling getattr. Reject unsupported names using the existing ValueError path, and allow only whitelisted configs to reach ModelOpt.
🧹 Nitpick comments (1)
recipes/evo2_megatron_quant/src/quantize.py (1)
183-183: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the non-deprecated AMP API.
torch.cuda.amp.autocastis deprecated; usetorch.autocast(device_type="cuda", dtype=torch.bfloat16)to avoid deprecation warnings and future compatibility issues. (docs.pytorch.org)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@recipes/evo2_megatron_quant/src/quantize.py` at line 183, Replace the deprecated torch.cuda.amp.autocast context in the quantization flow with torch.autocast using device_type="cuda" and dtype=torch.bfloat16, while preserving the surrounding torch.no_grad() behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@recipes/evo2_megatron_quant/scripts/benchmark_real_quant.py`:
- Around line 261-269: Update the output-file loading path in run_test to catch
JSONDecodeError from json.load when output_file is truncated or partially
written. Treat that case as a failed test using the existing CRASH result
structure, including method, precision, error details, and wall_time_s, so the
exception does not escape into main’s benchmark loop.
In `@recipes/evo2_megatron_quant/tests/test_single_model.py`:
- Around line 219-225: Add focused tests around main() covering the exit-code
contract: assert it returns 0 when all result statuses are PASS and 1 when any
method reports FAIL or ERROR. Mock model loading and quantization dependencies
so the tests remain independent of CPUs and checkpoints.
---
Outside diff comments:
In `@recipes/evo2_megatron_quant/src/quantize.py`:
- Around line 135-142: Update the enable_all_mlp quant_cfg handling to support
ModelOpt’s ordered list format in addition to the existing mapping format.
Locate and remove the list entry whose quantizer_name is "*" and whose
configuration disables matching modules, so enabling all MLP layers remains
effective; preserve unrelated quant_cfg entries and existing behavior for
mapping-based configurations.
- Around line 123-131: Update the quantization-method resolution flow around
FRIENDLY_NAMES, ALL_QUANT_METHODS, and getattr(mtq, resolved, None) to trim
whitespace from each method token before lookup, then validate the resolved name
against ALL_QUANT_METHODS before calling getattr. Reject unsupported names using
the existing ValueError path, and allow only whitelisted configs to reach
ModelOpt.
---
Nitpick comments:
In `@recipes/evo2_megatron_quant/src/quantize.py`:
- Line 183: Replace the deprecated torch.cuda.amp.autocast context in the
quantization flow with torch.autocast using device_type="cuda" and
dtype=torch.bfloat16, while preserving the surrounding torch.no_grad() behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 91951a2e-392b-49d3-91dc-b6f3bbcc8567
📒 Files selected for processing (9)
recipes/evo2_megatron_quant/README.mdrecipes/evo2_megatron_quant/docker/start_container.shrecipes/evo2_megatron_quant/requirements.txtrecipes/evo2_megatron_quant/scripts/benchmark_real_quant.pyrecipes/evo2_megatron_quant/scripts/download_models.shrecipes/evo2_megatron_quant/src/__init__.pyrecipes/evo2_megatron_quant/src/compress_model.pyrecipes/evo2_megatron_quant/src/quantize.pyrecipes/evo2_megatron_quant/tests/test_single_model.py
🚧 Files skipped from review as they are similar to previous changes (5)
- recipes/evo2_megatron_quant/docker/start_container.sh
- recipes/evo2_megatron_quant/src/init.py
- recipes/evo2_megatron_quant/scripts/download_models.sh
- recipes/evo2_megatron_quant/README.md
- recipes/evo2_megatron_quant/src/compress_model.py
| if os.path.exists(output_file): | ||
| with open(output_file) as f: | ||
| result = json.load(f) | ||
| result["wall_time_s"] = elapsed | ||
| return result | ||
| return { | ||
| "method": method_name, "precision": precision, | ||
| "status": "CRASH", "error": "No output file produced", | ||
| "wall_time_s": elapsed, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Guard json.load against a partially-written result file.
If the worker is killed while writing output_file (line 217 in the worker script), it can be left truncated. json.load will then raise JSONDecodeError uncaught here, propagating out of run_test() into main()'s loop and aborting the entire benchmark run before the CSV (written only after all configs complete) is ever saved — losing every previously collected result.
🐛 Proposed fix
if os.path.exists(output_file):
- with open(output_file) as f:
- result = json.load(f)
+ try:
+ with open(output_file) as f:
+ result = json.load(f)
+ except (json.JSONDecodeError, OSError) as e:
+ return {
+ "method": method_name, "precision": precision,
+ "status": "CRASH", "error": f"Malformed result file: {e}",
+ "wall_time_s": elapsed,
+ }
result["wall_time_s"] = elapsed
return result📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if os.path.exists(output_file): | |
| with open(output_file) as f: | |
| result = json.load(f) | |
| result["wall_time_s"] = elapsed | |
| return result | |
| return { | |
| "method": method_name, "precision": precision, | |
| "status": "CRASH", "error": "No output file produced", | |
| "wall_time_s": elapsed, | |
| if os.path.exists(output_file): | |
| try: | |
| with open(output_file) as f: | |
| result = json.load(f) | |
| except (json.JSONDecodeError, OSError) as e: | |
| return { | |
| "method": method_name, "precision": precision, | |
| "status": "CRASH", "error": f"Malformed result file: {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, |
🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 261-261: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(output_file)
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@recipes/evo2_megatron_quant/scripts/benchmark_real_quant.py` around lines 261
- 269, Update the output-file loading path in run_test to catch JSONDecodeError
from json.load when output_file is truncated or partially written. Treat that
case as a failed test using the existing CRASH result structure, including
method, precision, error details, and wall_time_s, so the exception does not
escape into main’s benchmark loop.
| # 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()) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add focused tests for the new exit-code contract.
Cover main() returning 0 when every method passes and 1 when any method is FAIL or ERROR; mock model loading and quantization so the tests remain CPU/checkpoint independent. As per coding guidelines, **/*test*.py requires focused tests when changing CI checks.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@recipes/evo2_megatron_quant/tests/test_single_model.py` around lines 219 -
225, Add focused tests around main() covering the exit-code contract: assert it
returns 0 when all result statuses are PASS and 1 when any method reports FAIL
or ERROR. Mock model loading and quantization dependencies so the tests remain
independent of CPUs and checkpoints.
Source: Coding guidelines
…locate GPU runner - quantize: strip + whitelist method names against ALL_QUANT_METHODS before getattr(); fail loudly instead of silently no-op'ing when --all-mlp cannot find the 'default' quant_cfg entry (a silent no-op would report all-MLP numbers for a default-scope run); drop deprecated torch.cuda.amp.autocast - benchmark_real_quant: guard json.load against a truncated worker result so one crashed config cannot abort the sweep before any CSV is written - adapters: drop deprecated torch.cuda.amp.autocast - move tests/test_single_model.py -> scripts/run_single_model.py: it is a manual GPU runner, not a pytest test. tests/ now holds only the CPU-only L0 gate, which also removes the pytest-collection hazard at the root - run_single_model: strip whitespace around comma-separated --quant tokens Signed-off-by: Rui Zhu <huthruiz@gmail.com>
|
Second pass addressed in Fixed as suggested
Two I solved differently — reasoning below, happy to revisit
|
1 participant
Description
Adds a new inference-side recipe,
recipes/evo2_megatron_quant/, for BioNeMo biological foundation models (ESM-2, Geneformer, Evo2). Upstream already covers low-precision training (TE FP8/MXFP8/NVFP4); this adds an inference path:nn.Linearreplacement (INT8Linear, per-channel symmetric) that stores weights in int8 and dequantizes on the fly. Measured on Evo2 7B: ~15% GPU-memory reduction, cosine sim ≈ 0.998, top-1 100% vs BF16.mtq.quantize()with an automated quality report (cosine sim / top-k agreement / MSE).tests/test_compressed_linear_roundtrip.py) — no GPU or checkpoint required, so it runs in fast CI.This is a Draft opened to get maintainer direction on the questions in #1681:
recipes/evo2_megatron_quant/the right home, or should this fold intorecipes/evo2_megatron/?Follows up on #1681.
Usage
Type of changes
CI Pipeline Configuration
Only the CPU-only L0 sanity test is intended for the default PR pipeline. The GPU benchmark and
tests/test_single_model.pyneed NGC checkpoints and are meant to be run manually.Pre-submit Checklist
/ok to test)Summary by CodeRabbit
New Features
Documentation
Tests