Background

While quantizing a large-scale LLM (70B+ parameters), the process crashed midway through with an out-of-memory error after running for several hours. There was no way to know in advance whether the available GPU VRAM was sufficient — the failure only surfaced deep into the quantization loop, wasting significant compute time.

This PR introduces a --check-env preflight flag that detects OOM risk before quantization starts, based on the physical characteristics of the execution environment.

Summary

• Adds --check-env CLI flag (also available as Runner.auto_run(check_env=True) in the Python API)
• Loads the model architecture on a meta device (zero GPU/CPU memory) to count parameters
• Collects hardware info: GPU VRAM (total & free), CPU RAM (via optional psutil), disk space
• Estimates memory requirements at 2-bit, 4-bit, and 8-bit quantization using existing weight_memory_gb(), plus a calibration overhead factor
• Classifies OOM risk as safe / warning / danger and prints a human-readable report
• On danger: exits with code 1 (CLI) or raises RuntimeError (library API), stopping quantization before it wastes GPU time
• On safe / warning: prints the report and continues with quantization as normal (preflight behavior)

Example output

============================================================
  OneComp Environment Check
============================================================

Hardware
  GPU count              : 1
  GPU name               : NVIDIA A100 80GB PCIe
  GPU VRAM (total)       : 80.0 GB
  GPU VRAM (free)        : 78.3 GB
  CPU RAM (total)        : 251.6 GB
  CPU RAM (avail)        : 230.1 GB
  Disk (avail)           : 320.4 GB  [/home/user/output]

Model: meta-llama/Llama-2-7b-hf
  Parameters             : 6,738,415,616
  FP16 footprint         : 12.54 GB

Memory Estimates
  2-bit quantized        :  1.96 GB
  4-bit quantized        :  3.77 GB
  8-bit quantized        :  7.28 GB
  Calib. overhead        :  1.88 GB  (15% of FP16)
  4-bit + overhead       :  5.65 GB

OOM Risk Assessment
  Risk level             : WARNING
  Detail                 : Free VRAM (78.3 GB) fits 4-bit quantized
                           weights but is tight (calibration overhead included).

  Recommended wbits      :  3.84  (VRAM-estimated)
============================================================

Risk thresholds

Changed files

Test plan

• onecomp <model_id> --check-env --no-eval — verify report is printed and quantization continues
• onecomp <model_id> --check-env --total-vram-gb 1.0 — verify danger triggers exit code 1
• onecomp <model_id> --check-env --total-vram-gb 6.0 --no-eval — verify warning continues
• CPU-only environment with --total-vram-gb — verify graceful handling without CUDA
• Without psutil installed — verify n/a fallback message appears
• Runner.auto_run(..., check_env=True) — verify RuntimeError raised on danger

🤖 Generated with Claude Code