From 4bc188e27d0e50d3d0ffbe21c800a0707e539ad7 Mon Sep 17 00:00:00 2001 From: codewithdark-git Date: Sun, 1 Nov 2026 20:21:32 +0500 Subject: [PATCH 1/2] fix: harden real-GPU paths and add an end-to-end GPU test suite Audit-driven fixes across loading, generation, fine-tuning, export, serving, Hub and packaging, verified on real GPUs (Kaggle 2x T4, torch 2.11 / transformers 5.16 / bitsandbytes 0.50): 23 passed, 1 skipped. Highlights: - Export no longer discards fine-tuned weights (in-memory dequantize, correct 8-bit SCB scaling, LoRA merged before ONNX/MLX staging) - trust_remote_code defaults to False everywhere; CLI --trust-remote-code - Load errors re-raised with the real cause instead of blaming architecture - Generation: no double BOS, multi-EOS respected, native stop strings, streaming errors propagate instead of hanging - Fine-tuning: transformers 5 compatible (no warmup_ratio), kwargs override defaults, correct LoRA targets, EOS/pad handling, adapter saved, final_loss - Server: non-blocking endpoints, model lock, OpenAI-compatible streaming - GGUF: converter deadlock, Windows .exe paths, real bf16/f32 outtype, no temp files in cwd, imatrix support - HQQ: real half-quadratic zero-point optimization; non-divisible groups - Hub: private repos honored, existing model cards preserved, hub 1.x registry - Packaging: optimum-onnx extra, transformers>=4.46, no git-URL extra GPU testing: - tests/gpu: behavioural E2E suite on a real model (skips without CUDA) - scripts/kaggle_gpu_tests.py: run the working tree on Kaggle's free GPUs - scripts/modal_gpu_tests.py, scripts/run_gpu_tests.sh, GPU E2E workflow Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 12 +- .github/workflows/gpu.yml | 45 ++ .gitignore | 4 + .pre-commit-config.yaml | 1 + CLAUDE.md | 1 + mypy.ini | 2 +- pyproject.toml | 37 +- quantllm/benchmark/auto_bench.py | 14 +- quantllm/benchmark/compare.py | 7 +- quantllm/benchmark/runner.py | 73 ++- quantllm/cli/__init__.py | 49 +- quantllm/core/compilation.py | 10 +- quantllm/core/export.py | 17 +- quantllm/core/export_router.py | 218 +++++--- quantllm/core/finetune_mixin.py | 188 ++++--- quantllm/core/flash_attention.py | 72 ++- quantllm/core/generation.py | 289 +++++++---- quantllm/core/hardware.py | 29 +- quantllm/core/memory.py | 64 ++- quantllm/core/model_analyzer.py | 46 +- quantllm/core/model_gguf.py | 5 +- quantllm/core/model_loader.py | 218 ++++++-- quantllm/core/push.py | 10 +- quantllm/core/smart_config.py | 76 ++- quantllm/core/training.py | 162 +++++- quantllm/core/turbo_model.py | 202 +++++--- quantllm/hub/hub_manager.py | 60 ++- quantllm/hub/model_card.py | 50 +- quantllm/kernels/cpu/__init__.py | 23 +- quantllm/kernels/triton/quantized_linear.py | 22 +- quantllm/quant/__init__.py | 10 +- quantllm/quant/hqq.py | 89 +++- quantllm/quant/llama_cpp.py | 216 ++++---- quantllm/registry.py | 119 ++++- quantllm/server/app.py | 111 +++-- quantllm/server/engine.py | 287 ++++++++--- quantllm/utils/hf_compat.py | 22 + quantllm/utils/logging.py | 4 +- quantllm/utils/progress.py | 23 + scripts/kaggle_gpu_tests.py | 163 ++++++ scripts/modal_gpu_tests.py | 115 +++++ scripts/run_gpu_tests.sh | 23 + tests/gpu/conftest.py | 75 +++ tests/gpu/test_gpu_e2e.py | 266 ++++++++++ tests/test_architecture_fallback.py | 6 +- tests/test_benchmark.py | 57 +++ tests/test_cli.py | 102 ++++ tests/test_export_fixes.py | 317 ++++++++++++ tests/test_finetune_fixes.py | 276 ++++++++++ tests/test_generation.py | 177 +++++++ tests/test_hub_push.py | 159 ++++++ tests/test_loading_fixes.py | 526 ++++++++++++++++++++ tests/test_registry.py | 129 ++++- tests/test_server.py | 236 +++++++++ 54 files changed, 4674 insertions(+), 840 deletions(-) create mode 100644 .github/workflows/gpu.yml create mode 100644 quantllm/utils/hf_compat.py create mode 100644 scripts/kaggle_gpu_tests.py create mode 100644 scripts/modal_gpu_tests.py create mode 100644 scripts/run_gpu_tests.sh create mode 100644 tests/gpu/conftest.py create mode 100644 tests/gpu/test_gpu_e2e.py create mode 100644 tests/test_export_fixes.py create mode 100644 tests/test_finetune_fixes.py create mode 100644 tests/test_hub_push.py create mode 100644 tests/test_loading_fixes.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0c77d44..7016073 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -55,7 +55,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.10", "3.11", "3.12", "3.13"] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -67,19 +67,23 @@ jobs: python -m pip install --upgrade pip pip install --index-url https://download.pytorch.org/whl/cpu "torch>=2.0.0" pip install \ - "transformers>=4.36.0" \ + "transformers>=4.46.0" \ "datasets>=2.14.0" \ "accelerate>=0.24.0" \ - "peft>=0.6.0" \ + "peft>=0.12.0" \ "scipy>=1.10.0" \ "scikit-learn>=1.3.0" \ "tqdm>=4.65.0" \ "rich>=13.0.0" \ "huggingface_hub>=0.20.0" \ - "bitsandbytes>=0.41.0" \ + "bitsandbytes>=0.43.0" \ "psutil" \ "gguf" \ "py-cpuinfo" \ + "fastapi>=0.100.0" \ + "uvicorn>=0.23.0" \ + "pydantic>=2.0.0" \ + "httpx>=0.24.0" \ "pytest>=7.4.0" - name: Install QuantLLM (no deps) run: pip install --no-deps -e . diff --git a/.github/workflows/gpu.yml b/.github/workflows/gpu.yml new file mode 100644 index 0000000..f9deaa9 --- /dev/null +++ b/.github/workflows/gpu.yml @@ -0,0 +1,45 @@ +name: GPU E2E + +# Runs tests/gpu on a real GPU through Modal. Manual, or on PRs labelled +# "run-gpu-tests" so normal pushes don't spend GPU credit. +# Needs repo secrets MODAL_TOKEN_ID and MODAL_TOKEN_SECRET. + +on: + workflow_dispatch: + inputs: + gpu: + description: "GPU type (T4, L4, A10G, H100)" + default: "T4" + slow: + description: "Include slow tests (GGUF via llama.cpp, ONNX)" + type: boolean + default: false + pull_request: + types: [labeled, synchronize] + +permissions: + contents: read + +concurrency: + group: gpu-${{ github.ref }} + cancel-in-progress: true + +jobs: + gpu-tests: + if: github.event_name == 'workflow_dispatch' || contains(github.event.pull_request.labels.*.name, 'run-gpu-tests') + runs-on: ubuntu-latest + timeout-minutes: 90 + env: + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - run: pip install "modal>=1.0" + - name: Run GPU suite on Modal + run: | + ARGS="--gpu ${{ inputs.gpu || 'T4' }}" + if [ "${{ inputs.slow }}" = "true" ]; then ARGS="$ARGS --slow"; fi + modal run scripts/modal_gpu_tests.py $ARGS diff --git a/.gitignore b/.gitignore index 1d3bc1c..77754c2 100644 --- a/.gitignore +++ b/.gitignore @@ -190,3 +190,7 @@ venv upcoming.md post.md modal_custom_workflow.py +hub_staging/ +gpu-test-results/ +gpu-probe-*/ +gpu-test-results-*/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index aa4c47e..6f314c9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -29,3 +29,4 @@ repos: hooks: - id: ruff args: ["--fix"] + - id: ruff-format diff --git a/CLAUDE.md b/CLAUDE.md index f91855c..0bea544 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -157,6 +157,7 @@ quantllm/ - Tests use CPU-only PyTorch (installed via `--index-url https://download.pytorch.org/whl/cpu` in CI) - Test `QUANTLLM_BANNER=0` to suppress banner +- **GPU end-to-end suite** lives in `tests/gpu/` (real model, default `Qwen/Qwen2.5-0.5B-Instruct`, override with `QUANTLLM_GPU_TEST_MODEL`). It skips itself without CUDA. Free option: `python scripts/kaggle_gpu_tests.py [--slow] [-k expr]` packs the working tree, runs it on Kaggle "GPU T4 x2", and downloads results to `gpu-test-results/` (needs `~/.kaggle/access_token` and a phone-verified account). T4 has no native bf16/FlashAttention-2, so those paths need an Ampere+ GPU. ONNX tests skip on transformers 5.x (optimum-onnx caps transformers<4.58). Paid/credit option: `modal run scripts/modal_gpu_tests.py [--gpu T4|L4|A10G|H100] [--slow] [--k Finetune]`, on any NVIDIA box/Colab/Kaggle with `bash scripts/run_gpu_tests.sh`, or in GitHub Actions via the `GPU E2E` workflow (manual or `run-gpu-tests` PR label; needs `MODAL_TOKEN_ID`/`MODAL_TOKEN_SECRET` secrets) - Test files mirror core modules: `test_hqq.py`, `test_cli.py`, `test_benchmark.py`, `test_server.py`, `test_registry.py`, `test_cpu_kernels.py`, `test_generation.py`, `test_quantization_state.py`, `test_resolve_model_type.py`, `test_export_push_config.py`, `test_architecture_fallback.py` ## Common Tasks diff --git a/mypy.ini b/mypy.ini index 3191665..4ae6f3d 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1,5 +1,5 @@ [mypy] -python_version = 3.10 +python_version = 3.12 ignore_missing_imports = true disallow_untyped_defs = false warn_unused_ignores = false diff --git a/pyproject.toml b/pyproject.toml index a80bb08..3c2428f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,11 +33,11 @@ keywords = [ ] dependencies = [ "torch>=2.0.0", - "transformers>=4.36.0", + "transformers>=4.46.0", # Trainer(processing_class=...) "datasets>=2.14.0", "accelerate>=0.24.0", - "peft>=0.6.0", - "bitsandbytes>=0.41.0", + "peft>=0.12.0", + "bitsandbytes>=0.43.0; platform_system != 'Darwin'", "scipy>=1.10.0", "scikit-learn>=1.3.0", "tqdm>=4.65.0", @@ -47,7 +47,6 @@ dependencies = [ "einops>=0.7.0", "psutil>=5.9.0", "huggingface-hub>=0.19.0", - "rich>=13.0.0", "gguf>=0.10.0", ] @@ -56,30 +55,35 @@ gguf = ["gguf>=0.10.0"] onnx = [ "onnx>=1.14.0", "onnxruntime>=1.16.0", - "optimum[onnxruntime]>=1.14.0", + # ONNX Runtime export lives in optimum-onnx; it caps transformers<4.58 for now + "optimum-onnx[onnxruntime]>=0.1.0", "onnxscript>=0.1.0", ] mlx = ["mlx>=0.1.0", "mlx-lm>=0.1.0"] triton = ["triton>=2.1.0"] flash = ["flash-attn>=2.3.0"] -hub = ["hf-lifecycle @ git+https://github.com/codewithdark-git/huggingface-lifecycle.git"] +# NOTE: the former `hub` extra pointed at a git URL (hf-lifecycle), which PyPI +# rejects on upload. Install it manually if needed: +# pip install "hf-lifecycle @ git+https://github.com/codewithdark-git/huggingface-lifecycle.git" full = [ "gguf>=0.10.0", "onnx>=1.14.0", "onnxruntime>=1.16.0", - "optimum[onnxruntime]>=1.14.0", + "optimum-onnx[onnxruntime]>=0.1.0", "onnxscript>=0.1.0", "triton>=2.1.0", "wandb>=0.15.0", "tensorboard>=2.14.0", "evaluate>=0.4.0", "flash-attn>=2.3.0", + "rich>=13.0.0", ] server = ["fastapi>=0.100.0", "uvicorn>=0.23.0", "pydantic>=2.0.0"] bench = ["pandas>=1.5.0"] dev = [ "pytest>=7.4.0", "pytest-cov>=4.1.0", + "httpx>=0.24.0", "black>=23.0.0", "isort>=5.12.0", "flake8>=6.1.0", @@ -138,26 +142,12 @@ ignore = [ "E741", # ambiguous variable names ] -[tool.myft.lint.per-file-ignores] +[tool.ruff.lint.per-file-ignores] "tests/*" = ["F401", "F811", "F841"] "quantllm/__init__.py" = ["F401", "E402"] "quantllm/core/__init__.py" = ["F401"] -[tool.mypy] -python_version = "3.12" -ignore_missing_imports = true -disallow_untyped_defs = false -warn_unused_ignores = false - -[[tool.mypy.overrides]] -module = "numpy" -ignore_errors = true -follow_imports = "skip" - -[[tool.mypy.overrides]] -module = "numpy.*" -ignore_errors = true -follow_imports = "skip" +# mypy is configured in mypy.ini (single source of truth). [tool.pytest.ini_options] minversion = "7.0" @@ -165,6 +155,7 @@ testpaths = ["tests"] addopts = "-ra --strict-markers" markers = [ "integration: marks tests as integration tests requiring GPU", + "gpu: end-to-end tests under tests/gpu (need CUDA; run via scripts/modal_gpu_tests.py)", "slow: marks tests as slow", "large_model: marks tests requiring large models", "stress: marks tests as stress/performance tests", diff --git a/quantllm/benchmark/auto_bench.py b/quantllm/benchmark/auto_bench.py index 8fea285..642f1d2 100644 --- a/quantllm/benchmark/auto_bench.py +++ b/quantllm/benchmark/auto_bench.py @@ -86,19 +86,25 @@ def _micro_bench(self) -> Dict[str, float]: times = [] for n in [1024, 2048]: a, b = torch.randn(n, n, device="cuda"), torch.randn(n, n, device="cuda") - start = time.time() + for _ in range(3): # warmup (kernel selection / lazy init) + a @ b + torch.cuda.synchronize() + start = time.perf_counter() for _ in range(10): a @ b torch.cuda.synchronize() - times.append((time.time() - start) / 10) + times.append((time.perf_counter() - start) / 10) results["matmul_ms"] = sum(times) / len(times) * 1000 x = torch.randn(1024, 1024, device="cuda") - start = time.time() + for _ in range(3): # warmup + x.clone() + torch.cuda.synchronize() + start = time.perf_counter() for _ in range(100): x.clone() torch.cuda.synchronize() - results["mem_bw_gb_s"] = (1024 * 1024 * 4 * 2 * 100) / (time.time() - start) / (1024**3) + results["mem_bw_gb_s"] = (1024 * 1024 * 4 * 2 * 100) / (time.perf_counter() - start) / (1024**3) return results def _score_quant(self, micro: Dict[str, float]) -> List[QuantScore]: diff --git a/quantllm/benchmark/compare.py b/quantllm/benchmark/compare.py index 40048ac..a916808 100644 --- a/quantllm/benchmark/compare.py +++ b/quantllm/benchmark/compare.py @@ -6,7 +6,7 @@ from typing import Any, Dict -from .runner import BenchmarkRunner, BenchmarkResult +from .runner import BenchmarkRunner, BenchmarkResult, release_memory class BenchmarkComparer: @@ -27,6 +27,9 @@ def run_comparison(self, prompt: str = "The meaning of life is", max_new_tokens: runner = BenchmarkRunner(verbose=self.verbose) for name, bits, _ in configs: + # Drop the previous config's model before loading the next one (otherwise VRAM/peak stats accumulate). + m = None + release_memory() try: m = turbo(self.model_name, bits=bits, verbose=False) r = runner.benchmark_model(m.model, m.tokenizer, m.config, prompt=prompt, max_new_tokens=max_new_tokens) @@ -37,6 +40,8 @@ def run_comparison(self, prompt: str = "The meaning of life is", max_new_tokens: from ..utils import print_warning print_warning(f"{name} failed: {e}") + m = None + release_memory() return self.results def generate_comparison_report(self) -> str: diff --git a/quantllm/benchmark/runner.py b/quantllm/benchmark/runner.py index 7f0cc24..fb12a34 100644 --- a/quantllm/benchmark/runner.py +++ b/quantllm/benchmark/runner.py @@ -49,6 +49,36 @@ def get(self, key: str, default=None): return getattr(self, key, default) +def _sync() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def release_memory() -> None: + """Free cached CUDA memory and reset the peak-memory counter (call after dropping the previous model).""" + import gc + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + +def _timed_forward(model, input_ids, num_runs: int) -> float: + """Average wall time (s) of a synchronized forward pass, after one warmup pass.""" + runs = max(num_runs, 1) + times = [] + with torch.inference_mode(): + model(input_ids=input_ids) + for _ in range(runs): + _sync() + start = time.perf_counter() + model(input_ids=input_ids) + _sync() + times.append(time.perf_counter() - start) + return sum(times) / len(times) + + class BenchmarkRunner: def __init__(self, verbose: bool = True): self.verbose = verbose @@ -76,9 +106,11 @@ def benchmark_model( pass r.model_size_gb = r.num_params_billion * (r.bits / 8) * 1.05 + release_memory() # reset the peak counter so VRAM reflects this model only + inputs = tokenizer(prompt, return_tensors="pt") input_ids = inputs["input_ids"].to(model.device) - prefill_len = input_ids.shape[1] + prompt_len = input_ids.shape[1] for _ in range(num_warmup): with torch.inference_mode(): @@ -90,31 +122,36 @@ def benchmark_model( eos_token_id=tokenizer.eos_token_id, ) - start = time.time() - with torch.inference_mode(): - model.generate( - input_ids, - max_new_tokens=1, - do_sample=False, - pad_token_id=tokenizer.pad_token_id, - eos_token_id=tokenizer.eos_token_id, - ) - r.prefill_tokens_per_sec = prefill_len / (time.time() - start) if (time.time() - start) > 0 else 0 - - decode_times = [] + # Prefill: a plain forward pass over a ~512-token prompt (the short prompt repeated). + max_pos = getattr(getattr(model, "config", None), "max_position_embeddings", None) or 512 + prefill_target = min(512, max_pos) + reps = prefill_target // max(prompt_len, 1) + 1 + long_ids = tokenizer(" ".join([prompt] * reps), return_tensors="pt")["input_ids"][:, :prefill_target] + long_ids = long_ids.to(model.device) + prefill_time = _timed_forward(model, long_ids, num_runs) + r.prefill_tokens_per_sec = long_ids.shape[1] / prefill_time if prefill_time > 0 else 0 + + # Decode: count tokens actually generated and exclude the prefill of the benchmark prompt. + short_prefill_time = _timed_forward(model, input_ids, num_runs) + decode_rates = [] for _ in range(num_runs): - start = time.time() + _sync() + start = time.perf_counter() with torch.inference_mode(): - model.generate( + out = model.generate( input_ids, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) - elapsed = time.time() - start - decode_times.append(max_new_tokens / elapsed if elapsed > 0 else 0) - r.decode_tokens_per_sec = sum(decode_times) / len(decode_times) + _sync() + elapsed = time.perf_counter() - start + new_tokens = out.shape[1] - prompt_len + decode_time = elapsed - short_prefill_time + # The first new token comes from the prefill pass; the remaining ones are decode steps. + decode_rates.append((new_tokens - 1) / decode_time if new_tokens > 1 and decode_time > 0 else 0) + r.decode_tokens_per_sec = sum(decode_rates) / len(decode_rates) r.total_latency_ms = (1000.0 / r.decode_tokens_per_sec) if r.decode_tokens_per_sec > 0 else 0 r.throughput_tokens_per_sec = r.decode_tokens_per_sec diff --git a/quantllm/cli/__init__.py b/quantllm/cli/__init__.py index b527281..1b0d6e4 100644 --- a/quantllm/cli/__init__.py +++ b/quantllm/cli/__init__.py @@ -19,13 +19,28 @@ import argparse import sys -__version__ = "2.2.0.dev" +from quantllm import __version__ def _add_verbose(p: argparse.ArgumentParser) -> None: p.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output") +def _add_trust_remote_code(p: argparse.ArgumentParser) -> None: + p.add_argument( + "--trust-remote-code", + action="store_true", + default=False, + help="Allow executing custom model code from the Hub (only for repos you trust)", + ) + + +def _load_model(args: argparse.Namespace): + from quantllm.core import turbo + + return turbo(args.model, verbose=args.verbose, trust_remote_code=args.trust_remote_code) + + def _add_model_arg(p: argparse.ArgumentParser) -> None: p.add_argument("model", help="Model name or path") @@ -43,6 +58,7 @@ def main() -> None: _add_model_arg(cp) cp.add_argument("-o", "--output", required=True, help="Output GGUF file") cp.add_argument("-q", "--quant", default="q4_k_m", help="Quantization type") + _add_trust_remote_code(cp) _add_verbose(cp) fp = sub.add_parser("finetune", help="Fine-tune a model") @@ -51,6 +67,7 @@ def main() -> None: fp.add_argument("--output-dir", default="./fine-tuned") fp.add_argument("--epochs", type=int, default=3) fp.add_argument("--lr", type=float, default=2e-4) + _add_trust_remote_code(fp) _add_verbose(fp) ep = sub.add_parser("export", help="Export model to any format") @@ -58,6 +75,7 @@ def main() -> None: ep.add_argument("-f", "--format", default="safetensors", choices=["gguf", "safetensors", "onnx", "mlx"]) ep.add_argument("-o", "--output") ep.add_argument("-q", "--quant") + _add_trust_remote_code(ep) _add_verbose(ep) sp = sub.add_parser("serve", help="Start inference server") @@ -66,6 +84,7 @@ def main() -> None: sp.add_argument("--port", type=int, default=8080) sp.add_argument("--max-seq-length", type=int, default=4096) sp.add_argument("--dtype", default="auto", choices=["auto", "float16", "bfloat16"]) + _add_trust_remote_code(sp) bp = sub.add_parser("bench", help="Run benchmarks") _add_model_arg(bp) @@ -93,6 +112,9 @@ def main() -> None: if args.command is None: parser.print_help() sys.exit(0) + if args.command == "models" and args.models_command is None: + mp.print_help(sys.stderr) + sys.exit(2) _route(args) @@ -115,20 +137,14 @@ def _route(args: argparse.Namespace) -> None: print("\n" + "=" * 50) sys.exit(0) elif cmd == "convert": - from quantllm.core import turbo - - m = turbo(args.model, verbose=args.verbose) + m = _load_model(args) m.export("gguf", output_path=args.output, quantization=args.quant) print(f"Done: {args.output}") elif cmd == "finetune": - from quantllm.core import turbo - - m = turbo(args.model, verbose=args.verbose) + m = _load_model(args) m.finetune(args.dataset, epochs=args.epochs, learning_rate=args.lr, output_dir=args.output_dir) elif cmd == "export": - from quantllm.core import turbo - - m = turbo(args.model, verbose=args.verbose) + m = _load_model(args) print(m.export(args.format, output_path=args.output, quantization=args.quant)) elif cmd == "serve": from quantllm.server import ServerConfig @@ -136,7 +152,12 @@ def _route(args: argparse.Namespace) -> None: run_server( ServerConfig.from_cli( - model=args.model, host=args.host, port=args.port, max_seq_length=args.max_seq_length, dtype=args.dtype + model=args.model, + host=args.host, + port=args.port, + max_seq_length=args.max_seq_length, + dtype=args.dtype, + trust_remote_code=args.trust_remote_code, ) ) elif cmd == "bench": @@ -152,8 +173,8 @@ def _route(args: argparse.Namespace) -> None: elif cmd == "models": sub = args.models_command if sub is None: - print("Usage: quantllm models {list,search,info}") - return + print("Usage: quantllm models {list,search,info}", file=sys.stderr) + sys.exit(2) from quantllm.registry import list_models, search_models, model_info print( @@ -170,6 +191,8 @@ def _route(args: argparse.Namespace) -> None: args.model_id, family=args.family, params=args.params, quality_score=args.quality, speed_score=args.speed ) print(f"Model {'registered' if ok else 'registration failed'}: {args.model_id}") + if not ok: + sys.exit(1) else: print(f"Unknown command: {cmd}") sys.exit(1) diff --git a/quantllm/core/compilation.py b/quantllm/core/compilation.py index 4bb40b2..3d6e6e3 100644 --- a/quantllm/core/compilation.py +++ b/quantllm/core/compilation.py @@ -67,13 +67,8 @@ def compile_model( return model try: - # Suppress compilation errors for robustness - import torch._dynamo as dynamo - - dynamo.config.suppress_errors = True - dynamo.config.automatic_dynamic_shapes = dynamic # type: ignore[attr-defined] - - # Compile the model + # Compile the model (``dynamic`` is passed per call; global dynamo + # config such as suppress_errors is intentionally left untouched) compiled = torch.compile( model, mode=mode, @@ -96,7 +91,6 @@ def compile_for_inference(model: nn.Module) -> nn.Module: Uses settings optimized for inference rather than training: - reduce-overhead mode for faster warmup - Dynamic shapes for variable sequence lengths - - Suppressed errors for robustness Args: model: Model to compile diff --git a/quantllm/core/export.py b/quantllm/core/export.py index 2ea1fce..5154ddc 100644 --- a/quantllm/core/export.py +++ b/quantllm/core/export.py @@ -104,10 +104,12 @@ def _export_gguf( ensure_llama_cpp_installed() model_name = getattr(self.model.config, "_name_or_path", "model").split("/")[-1] + model = self._merged_model() + # The intermediate F16 GGUF is written inside temp_dir, so it is removed on success and failure with tempfile.TemporaryDirectory() as temp_dir: # Save model to temp - self.model.save_pretrained(temp_dir, safe_serialization=True) # type: ignore[operator] + model.save_pretrained(temp_dir, safe_serialization=True) # type: ignore[operator] if self.tokenizer: self.tokenizer.save_pretrained(temp_dir) @@ -119,6 +121,7 @@ def _export_gguf( model_dtype="f16", quantization_type="f16", print_output=False, + output_dir=temp_dir, ) if output_files: @@ -134,6 +137,12 @@ def _export_gguf( return output_path + def _merged_model(self) -> nn.Module: + """Return the model with any PEFT/LoRA adapters merged into the base weights.""" + if hasattr(self.model, "merge_and_unload"): + self.model = self.model.merge_and_unload() # type: ignore[operator] + return self.model + def _export_safetensors( self, output_path: str, @@ -142,8 +151,8 @@ def _export_safetensors( """Export to SafeTensors format.""" os.makedirs(output_path, exist_ok=True) - # Save model - self.model.save_pretrained(output_path, safe_serialization=True) # type: ignore[operator] + # Save model (merged, so a PeftModel doesn't save only its adapter weights) + self._merged_model().save_pretrained(output_path, safe_serialization=True) # type: ignore[operator] # Save tokenizer if available if self.tokenizer: @@ -177,10 +186,10 @@ def _export_onnx( # Prepare dummy input if self.tokenizer: + # Single sequence: no padding (avoids errors for tokenizers without a pad token) dummy_input = self.tokenizer( "Hello world", return_tensors="pt", - padding=True, ) dummy_input = {k: v.to(self.model.device) for k, v in dummy_input.items()} else: diff --git a/quantllm/core/export_router.py b/quantllm/core/export_router.py index 52fa71c..1db373f 100644 --- a/quantllm/core/export_router.py +++ b/quantllm/core/export_router.py @@ -93,6 +93,7 @@ def export( print_info("Merging LoRA weights before export...") self.model = self.model.merge_and_unload() self._lora_applied = False + self._lora_merged = True if output_path is None: model_name = self.model.config._name_or_path.split("/")[-1] @@ -133,13 +134,18 @@ def _export_gguf( smart_tensor_ordering: bool = False, disk_offloading: bool = False, disk_offload_dir: Optional[str] = None, + imatrix: Optional[str] = None, **kwargs, ) -> str: - """Export to GGUF format using llama.cpp converter.""" + """Export to GGUF format using llama.cpp converter. + + ``imatrix`` is the path to a llama-imatrix file, required for IQ1/IQ2/IQ3_XXS types. + """ from ..quant import ( convert_to_gguf, quantize_gguf, ensure_llama_cpp_installed, + gguf_outtype, ) from .memory import memory_optimized_tensor_order @@ -153,15 +159,12 @@ def _export_gguf( quant_type_upper = quant_type.upper() quant_type_lower = quant_type.lower() - passthrough_types = { - "f16", - "f32", - "bf16", - "float16", - "float32", - "bfloat16", - } - needs_quantization = quant_type_lower not in passthrough_types + # Full-precision targets (f16/bf16/f32) are written directly by the converter + passthrough_outtype = gguf_outtype(quant_type_lower) + needs_quantization = passthrough_outtype is None + if passthrough_outtype is not None: + quant_type_lower = passthrough_outtype + quant_type_upper = passthrough_outtype.upper() if self.verbose: print_info(f"Target quantization: {quant_type_upper}") @@ -190,8 +193,11 @@ def _export_gguf( if self.verbose: print_success("Model dequantized successfully!") - model_dtype = "f16" + # Converter --outtype: the target itself for passthrough types, F16 intermediate otherwise + model_dtype = passthrough_outtype or "f16" model_name = self.model.config._name_or_path.split("/")[-1] + output_path = os.path.abspath(output_path) + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) temp_parent = disk_offload_dir if disk_offloading else None if temp_parent: @@ -228,29 +234,24 @@ def _export_gguf( if self.verbose: print_header("Step 2/3: Converting to GGUF", icon="[convert]") - if needs_quantization: - f16_gguf_file = os.path.join(temp_dir, f"{model_name}.F16.gguf") - else: - f16_gguf_file = f"{model_name}.{quant_type_upper}.gguf" - + # The converter writes into the staging dir, so the intermediate F16 file is + # removed with it on both success and failure (never left in the CWD). output_files, _ = convert_to_gguf( model_name=model_name, input_folder=temp_dir, model_dtype=model_dtype, - quantization_type=("f16" if needs_quantization else quant_type_lower), + quantization_type=model_dtype, print_output=self.verbose, + output_dir=temp_dir, ) - if not output_files: + if not output_files or not os.path.exists(output_files[0]): raise RuntimeError("GGUF conversion failed to produce output file.") - f16_file = output_files[0] - - if os.path.exists(f16_file): - f16_gguf_file = f16_file + f16_gguf_file = output_files[0] if self.verbose: - print_success(f"F16 GGUF created: {f16_gguf_file}") + print_success(f"{model_dtype.upper()} GGUF created: {f16_gguf_file}") if needs_quantization: if self.verbose: @@ -260,17 +261,19 @@ def _export_gguf( ) print_info(f"Applying {quant_type_upper} quantization...") - quantized_file = f"{model_name}.{quant_type_upper}.gguf" + quantized_file = os.path.join(temp_dir, f"{model_name}.{quant_type_upper}.gguf") quantize_gguf( input_gguf=f16_gguf_file, output_gguf=quantized_file, quant_type=quant_type_upper, print_output=self.verbose, + imatrix=imatrix, ) final_file = quantized_file + # Free the intermediate F16 file as soon as it is no longer needed if os.path.exists(f16_gguf_file) and f16_gguf_file != quantized_file: os.remove(f16_gguf_file) @@ -322,20 +325,21 @@ def _export_onnx( """Export to ONNX format using Optimum.""" try: from optimum.onnxruntime import ORTModelForCausalLM + except ImportError as e: + import transformers - HAS_OPTIMUM = True - except ImportError: - HAS_OPTIMUM = False - - if not HAS_OPTIMUM: error_msg = ( - "ONNX export requires the Optimum library.\n\n" - "Install: pip install onnx onnxruntime optimum[onnxruntime] " - "onnxscript\n" - "Or: pip install quantllm[onnx]" + "ONNX export requires optimum-onnx: pip install 'quantllm[onnx]'\n" + "(or: pip install 'optimum-onnx[onnxruntime]' onnx onnxscript)" ) + if int(transformers.__version__.split(".")[0]) >= 5: + error_msg += ( + "\nNote: optimum-onnx currently supports transformers<4.58, but transformers " + f"{transformers.__version__} is installed. Export ONNX from a separate environment " + "with 'transformers<4.58'." + ) print_error(error_msg) - raise ImportError("ONNX export requires: pip install onnx onnxruntime optimum[onnxruntime] onnxscript") + raise ImportError(error_msg) from e os.makedirs(output_path, exist_ok=True) model_name = self.model.config._name_or_path @@ -347,7 +351,8 @@ def _export_onnx( task = progress.add_task("Exporting to ONNX...", total=None) try: - if self._is_bnb_quantized(): + if self._is_bnb_quantized() and not self._has_in_memory_changes(): + # Unmodified BNB model: the original Hub weights are more accurate than a dequantized copy if self.verbose: print_info("BNB quantized model detected. Exporting from original HuggingFace model...") ort_model = ORTModelForCausalLM.from_pretrained( @@ -356,12 +361,9 @@ def _export_onnx( trust_remote_code=True, ) else: - temp_path = os.path.join(output_path, "_temp_hf") - os.makedirs(temp_path, exist_ok=True) - + # Export the in-memory weights (fine-tuned / LoRA-merged changes are preserved) + temp_path = self._stage_in_memory_model() try: - self.model.save_pretrained(temp_path, safe_serialization=True) - self.tokenizer.save_pretrained(temp_path) ort_model = ORTModelForCausalLM.from_pretrained( temp_path, export=True, @@ -385,15 +387,49 @@ def _export_onnx( if quantization: if self.verbose: print_info(f"Applying {quantization} ONNX quantization...") - self._quantize_onnx_model(output_path, quantization) + quantized_name = self._quantize_onnx_model(output_path, quantization) + if quantized_name: + # ORTModel loads model.onnx by default; the INT8 graph lives in a separate file + print_info(f"Quantized ONNX model: {os.path.join(output_path, quantized_name)}") + print_info( + f"Load it with ORTModelForCausalLM.from_pretrained('{output_path}', file_name='{quantized_name}')" + ) if self.verbose: print_success(f"ONNX model exported to {output_path}") return output_path - def _quantize_onnx_model(self, model_path: str, quant_type: str) -> None: - """Apply ONNX INT8 quantization.""" + def _has_in_memory_changes(self) -> bool: + """True when the in-memory weights differ from the Hub checkpoint (fine-tuned or LoRA merged).""" + return bool( + getattr(self, "_is_finetuned", False) + or getattr(self, "_lora_applied", False) + or getattr(self, "_lora_merged", False) + ) + + def _stage_in_memory_model(self) -> str: + """Save the in-memory model (dequantized if BNB) + tokenizer to a new temp dir and return it. + + The caller owns the directory and must remove it (``shutil.rmtree``). + """ + model_to_save = self._dequantize_model() if self._is_bnb_quantized() else self.model + staging_dir = tempfile.mkdtemp(prefix="quantllm_export_") + try: + model_to_save.save_pretrained(staging_dir, safe_serialization=True) + if self.tokenizer is not None: + self.tokenizer.save_pretrained(staging_dir) + except Exception: + shutil.rmtree(staging_dir, ignore_errors=True) + raise + return staging_dir + + def _quantize_onnx_model(self, model_path: str, quant_type: str) -> Optional[str]: + """Apply ONNX INT8 quantization. + + Returns the file name of the quantized model inside ``model_path`` + (e.g. ``model_quantized.onnx``), or ``None`` if quantization was skipped. + """ try: from optimum.onnxruntime import ORTQuantizer from optimum.onnxruntime.configuration import AutoQuantizationConfig @@ -432,10 +468,16 @@ def _quantize_onnx_model(self, model_path: str, quant_type: str) -> None: else: qconfig = AutoQuantizationConfig.avx2(is_static=False, per_channel=True) - quantizer.quantize(save_dir=model_path, quantization_config=qconfig) + file_suffix = "quantized" + quantizer.quantize(save_dir=model_path, quantization_config=qconfig, file_suffix=file_suffix) + + source_onnx = getattr(quantizer, "onnx_model_path", None) + source_stem = os.path.splitext(os.path.basename(str(source_onnx)))[0] if source_onnx else "model" + quantized_name = f"{source_stem}_{file_suffix}.onnx" if self.verbose: - print_success("ONNX INT8 quantization applied successfully") + print_success(f"ONNX INT8 quantization applied successfully ({quantized_name})") + return quantized_name except ImportError: print_warning("Optimum quantization not available. Skipping ONNX quantization.") @@ -443,6 +485,32 @@ def _quantize_onnx_model(self, model_path: str, quant_type: str) -> None: except Exception as e: print_warning(f"ONNX quantization failed: {e}") print_info("The unquantized ONNX model is still available.") + return None + + @staticmethod + def _mlx_quant_bits(quantization: Optional[str]) -> Optional[int]: + """Map a quantization name to MLX ``q_bits`` (4 or 8), or ``None`` for no quantization. + + Full-precision names (f16/bf16/fp16/f32/none) are checked first so that e.g. the + "16" in "bf16" is not mistaken for a 6-bit request. + """ + if not quantization: + return None + quant_lower = quantization.lower().replace("_", "").replace("-", "") + if quant_lower in {"f16", "bf16", "fp16", "f32", "fp32", "float16", "bfloat16", "float32", "none"}: + return None + + if any(x in quant_lower for x in ["2", "3"]): + print_warning(f"MLX only supports 4-bit and 8-bit, not {quantization}.") + print_info("Using 4-bit quantization (smallest available).") + return 4 + if any(x in quant_lower for x in ["5", "6", "7"]): + print_warning(f"MLX only supports 4-bit and 8-bit, not {quantization}.") + print_info("Using 8-bit quantization (closest available).") + return 8 + if "8" in quant_lower: + return 8 + return 4 def _export_mlx( self, @@ -464,7 +532,7 @@ def _export_mlx( except ImportError: HAS_MLX = False - os.makedirs(output_path, exist_ok=True) + # Don't pre-create output_path: mlx_lm.convert refuses to write into an existing directory model_name = self.model.config._name_or_path if HAS_MLX: @@ -477,46 +545,29 @@ def _export_mlx( with QuantLLMProgress() as progress: task = progress.add_task("Converting to MLX...", total=None) - if self._is_bnb_quantized(): - source_path = model_name - else: - source_path = os.path.join(output_path, "_temp_hf") - os.makedirs(source_path, exist_ok=True) - self.model.save_pretrained(source_path) - self.tokenizer.save_pretrained(source_path) - - convert_args = { - "hf_path": source_path, - "mlx_path": output_path, - } - - if quantization: - quant_lower = quantization.lower().replace("_", "").replace("-", "") - - if any(x in quant_lower for x in ["2", "3"]): - print_warning(f"MLX only supports 4-bit and 8-bit, not {quantization}.") - print_info("Using 4-bit quantization (smallest available).") - convert_args["quantize"] = True - convert_args["q_bits"] = 4 - elif any(x in quant_lower for x in ["5", "6", "7"]): - print_warning(f"MLX only supports 4-bit and 8-bit, not {quantization}.") - print_info("Using 8-bit quantization (closest available).") - convert_args["quantize"] = True - convert_args["q_bits"] = 8 - elif "4" in quant_lower: - convert_args["quantize"] = True - convert_args["q_bits"] = 4 - elif "8" in quant_lower: - convert_args["quantize"] = True - convert_args["q_bits"] = 8 + staging_dir = None + try: + if self._is_bnb_quantized() and not self._has_in_memory_changes(): + # Unmodified BNB model: convert from the original Hub weights + source_path = model_name else: + # Convert the in-memory weights (fine-tuned / LoRA-merged changes are preserved) + staging_dir = self._stage_in_memory_model() + source_path = staging_dir + + convert_args: dict[str, Any] = { + "hf_path": source_path, + "mlx_path": output_path, + } + q_bits = self._mlx_quant_bits(quantization) + if q_bits is not None: convert_args["quantize"] = True - convert_args["q_bits"] = 4 - - convert(**convert_args) + convert_args["q_bits"] = q_bits - if not self._is_bnb_quantized(): - shutil.rmtree(source_path, ignore_errors=True) + convert(**convert_args) + finally: + if staging_dir is not None: + shutil.rmtree(staging_dir, ignore_errors=True) progress.update(task, completed=100) @@ -528,6 +579,7 @@ def _export_mlx( print_warning("mlx-lm not installed. Saving as HuggingFace format.") print_info("To convert to MLX: pip install mlx-lm && python -m mlx_lm.convert ...") + os.makedirs(output_path, exist_ok=True) self.model.save_pretrained(output_path) self.tokenizer.save_pretrained(output_path) diff --git a/quantllm/core/finetune_mixin.py b/quantllm/core/finetune_mixin.py index 26a6dc8..1b5329e 100644 --- a/quantllm/core/finetune_mixin.py +++ b/quantllm/core/finetune_mixin.py @@ -4,6 +4,8 @@ the ``finetune`` method. """ +import inspect +import math import os import warnings from typing import Any, Dict, List, Optional, Union @@ -29,6 +31,7 @@ class FinetuneMixin: _is_quantized: bool _is_finetuned: bool _lora_applied: bool + _is_bnb_quantized: Any # method on TurboModel def finetune( self, @@ -82,8 +85,21 @@ def finetune( except ImportError: raise ImportError("peft is required for fine-tuning. Install with: pip install peft") - if self._is_quantized: - self.model = prepare_model_for_kbit_training(self.model) + self._ensure_distinct_pad_token() + + if not self._lora_applied: + if self._is_quantized: + self.model = prepare_model_for_kbit_training( + self.model, + use_gradient_checkpointing=bool(self.config.gradient_checkpointing), + ) + else: + # Second finetune(): the model is already a PeftModel. Don't re-run + # prepare_model_for_kbit_training (it would freeze the adapters); + # just make sure the LoRA weights are trainable again. + for name, param in self.model.named_parameters(): + if "lora_" in name: + param.requires_grad_(True) if self.config.gradient_checkpointing: self.model.gradient_checkpointing_enable() @@ -135,7 +151,8 @@ def finetune( torch_compile = self.config.compile_model and self.config.effective_loading_bits >= 16 try: - training_args = TrainingArguments( + # Defaults first; user kwargs override them instead of raising "multiple values". + args_dict: Dict[str, Any] = dict( output_dir=output_dir, num_train_epochs=epochs, per_device_train_batch_size=batch_size, @@ -145,34 +162,58 @@ def finetune( bf16=self.config.dtype == torch.bfloat16, logging_steps=10, save_strategy="epoch", - warmup_ratio=0.03, + # warmup_ratio was removed in transformers 5.x; warmup_steps works everywhere. + warmup_steps=max( + 1, + int( + 0.03 + * epochs + * math.ceil(len(train_dataset) / (batch_size * self.config.gradient_accumulation_steps)) + ), + ), lr_scheduler_type="cosine", gradient_checkpointing=self.config.gradient_checkpointing, - optim=("paged_adamw_32bit" if self.config.bits <= 4 else "adamw_torch"), + optim=("paged_adamw_32bit" if self._is_bnb_quantized() else "adamw_torch"), torch_compile=torch_compile, - **kwargs, ) + args_dict.update(kwargs) + training_args = TrainingArguments(**args_dict) + + trainer_kwargs: Dict[str, Any] = { + "model": self.model, + "args": training_args, + "train_dataset": train_dataset, + "data_collator": DataCollatorForLanguageModeling(self.tokenizer, mlm=False), + } + # ``processing_class`` replaced ``tokenizer`` in transformers 4.46. + if "processing_class" in inspect.signature(Trainer.__init__).parameters: + trainer_kwargs["processing_class"] = self.tokenizer + else: + trainer_kwargs["tokenizer"] = self.tokenizer - trainer = Trainer( - model=self.model, - args=training_args, - train_dataset=train_dataset, - processing_class=self.tokenizer, - data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False), - ) + trainer = Trainer(**trainer_kwargs) result = trainer.train() self._is_finetuned = True + trainer.save_model(output_dir) + if hasattr(self.tokenizer, "save_pretrained"): + self.tokenizer.save_pretrained(output_dir) + print_success(f"Training complete! Model saved to {output_dir}") + step_losses = [h["loss"] for h in (getattr(trainer.state, "log_history", None) or []) if "loss" in h] return { + # train_loss is the mean over all steps, so it is inflated by early steps that + # fp16 loss scaling skips; final_loss is the last logged step. "train_loss": result.training_loss, + "final_loss": step_losses[-1] if step_losses else result.training_loss, "epochs": epochs, "output_dir": output_dir, "learning_rate": learning_rate, "batch_size": batch_size, "lora_r": r, + "log_history": list(getattr(trainer.state, "log_history", None) or []), } except Exception as e: @@ -180,52 +221,65 @@ def finetune( if "out of memory" in str(e).lower(): print_info("Tip: Try reducing batch_size or enabling gradient_checkpointing in config.") raise + finally: + self._restore_inference_mode() + + def _restore_inference_mode(self) -> None: + """Switch the model back to inference mode after training.""" + model = self.model + try: + if getattr(model, "is_gradient_checkpointing", False) is True: + model.gradient_checkpointing_disable() + if callable(getattr(model, "eval", None)): + model.eval() + cfg = getattr(model, "config", None) + if cfg is not None and hasattr(cfg, "use_cache"): + cfg.use_cache = True + except Exception as e: # pragma: no cover - defensive + logger.debug(f"Could not restore inference mode: {e}") + + def _ensure_distinct_pad_token(self) -> None: + """Make pad != eos so the LM collator doesn't mask EOS labels.""" + tok = self.tokenizer + eos = getattr(tok, "eos_token", None) + pad = getattr(tok, "pad_token", None) + if pad is not None and pad != eos: + return + unk = getattr(tok, "unk_token", None) + if unk is not None and unk != eos: + # Prefer unk: avoids resizing embeddings (problematic on quantized models). + tok.pad_token = unk + return + tok.add_special_tokens({"pad_token": "<|pad|>"}) + try: + self.model.resize_token_embeddings(len(tok)) + except Exception as e: + logger.warning(f"Could not resize token embeddings for new pad token: {e}") - def _get_lora_target_modules(self) -> List[str]: - """Get appropriate LoRA target modules for the model.""" - model_type = getattr(self.model.config, "model_type", "").lower() - - target_patterns = { - "llama": [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - "mistral": [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - "qwen2": [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - "phi": ["q_proj", "k_proj", "v_proj", "dense", "fc1", "fc2"], - "gemma": [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - } - - return target_patterns.get(model_type, ["q_proj", "v_proj"]) + def _get_lora_target_modules(self) -> Union[List[str], str]: + """Get appropriate LoRA target modules for the model. + + Uses the architecture table in :class:`LoRAAutoConfig`; falls back to + peft's ``"all-linear"`` for unknown architectures. + """ + from .training import LoRAAutoConfig + + try: + model_type = LoRAAutoConfig._detect_model_type(self.model) + except Exception: + model_type = (getattr(getattr(self.model, "config", None), "model_type", "") or "").lower() + + modules = LoRAAutoConfig.TARGET_MODULES.get(model_type) if model_type != "default" else None + if not modules: + return "all-linear" + + try: + present = {name.split(".")[-1] for name, _ in self.model.named_modules()} + except Exception: + return list(modules) + + valid = [m for m in modules if m in present] + return valid if valid else "all-linear" def _prepare_dataset(self, data: Union[str, List[Dict], Any]) -> Any: """Prepare dataset for training.""" @@ -241,23 +295,19 @@ def _prepare_dataset(self, data: Union[str, List[Dict], Any]) -> Any: else: dataset = data + from .training import format_training_texts + + tokenizer = self.tokenizer + def tokenize_function(examples): - if "text" in examples: - texts = examples["text"] - elif "instruction" in examples and "output" in examples: - texts = [ - f"### Instruction:\n{inst}\n\n### Response:\n{out}" - for inst, out in zip(examples["instruction"], examples["output"]) - ] - else: - keys = [k for k in examples.keys() if isinstance(examples[k][0], str)] - texts = [" ".join([examples[k][i] for k in keys]) for i in range(len(examples[keys[0]]))] + texts, templated = format_training_texts(examples, tokenizer) - return self.tokenizer( + return tokenizer( texts, truncation=True, max_length=self.config.max_seq_length, padding=False, + add_special_tokens=not templated, ) tokenized = dataset.map( diff --git a/quantllm/core/flash_attention.py b/quantllm/core/flash_attention.py index 49389bc..6064a79 100644 --- a/quantllm/core/flash_attention.py +++ b/quantllm/core/flash_attention.py @@ -209,43 +209,61 @@ def _pytorch_attention( dropout_p: float = 0.0, ) -> torch.Tensor: """ - Standard PyTorch attention implementation. + PyTorch SDPA attention used when Flash Attention is not available. - Used as fallback when Flash Attention is not available. + Matches ``flash_attn_func`` semantics: inputs are + ``[batch, seqlen, num_heads, head_dim]``, K/V heads are repeated for + GQA/MQA, and the causal mask is aligned to the bottom-right corner when + ``seqlen_q != seqlen_k`` (e.g. decoding with a KV cache). """ - batch, seqlen_q, num_heads, head_dim = q.shape - seqlen_k = k.shape[1] + import torch.nn.functional as F - if softmax_scale is None: - softmax_scale = 1.0 / (head_dim**0.5) + seqlen_q, num_heads, head_dim = q.shape[1], q.shape[2], q.shape[3] + seqlen_k, num_kv_heads = k.shape[1], k.shape[2] - # Reshape for batched matmul: [batch, num_heads, seqlen, head_dim] + # [batch, num_heads, seqlen, head_dim] + if softmax_scale is not None: + # Fold a custom scale into q (SDPA's ``scale`` kwarg needs torch>=2.1) + q = q * (softmax_scale * head_dim**0.5) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) - # Compute attention scores - scores = torch.matmul(q, k.transpose(-2, -1)) * softmax_scale + # GQA / MQA: repeat KV heads to match query heads + if num_kv_heads != num_heads: + if num_heads % num_kv_heads != 0: + raise ValueError(f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_heads})") + repeats = num_heads // num_kv_heads + k = k.repeat_interleave(repeats, dim=1) + v = v.repeat_interleave(repeats, dim=1) - # Apply causal mask + attn_mask = None + is_causal = False if causal: - mask = torch.triu(torch.ones(seqlen_q, seqlen_k, device=q.device, dtype=torch.bool), diagonal=1) - scores = scores.masked_fill(mask, float("-inf")) - - # Softmax - attn_weights = torch.softmax(scores, dim=-1) - - # Dropout - if dropout_p > 0.0 and q.requires_grad: - attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout_p) - - # Output - output = torch.matmul(attn_weights, v) - - # Reshape back: [batch, seqlen, num_heads, head_dim] - output = output.transpose(1, 2) - - return output + if seqlen_q == seqlen_k: + is_causal = True + else: + # Bottom-right aligned: query i may attend keys j <= i + (seqlen_k - seqlen_q) + offset = seqlen_k - seqlen_q + q_idx = torch.arange(seqlen_q, device=q.device).unsqueeze(1) + k_idx = torch.arange(seqlen_k, device=q.device).unsqueeze(0) + attn_mask = k_idx <= q_idx + offset + + output = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + ) + + if attn_mask is not None: + # Rows with no visible keys (seqlen_q > seqlen_k) produce NaN; flash_attn returns 0. + output = torch.nan_to_num(output, nan=0.0) + + # Back to [batch, seqlen, num_heads, head_dim] + return output.transpose(1, 2) class FlashAttentionWrapper(nn.Module): diff --git a/quantllm/core/generation.py b/quantllm/core/generation.py index 01a5100..a365e98 100644 --- a/quantllm/core/generation.py +++ b/quantllm/core/generation.py @@ -37,6 +37,7 @@ def generate( stream: bool = False, repetition_penalty: float = 1.1, stop_strings: Optional[List[str]] = None, + add_special_tokens: bool = True, **kwargs, ): """Generate text from a prompt. @@ -51,118 +52,173 @@ def generate( stream: Whether to stream output token by token (returns generator). repetition_penalty: Penalty for repeating tokens. stop_strings: List of strings that stop generation. + add_special_tokens: Whether the tokenizer adds special tokens (BOS). + Set to False for prompts already rendered by a chat template. """ - inputs = self.tokenizer( - prompt, - return_tensors="pt", - truncation=True, - max_length=self.config.max_seq_length - max_new_tokens, - ) + self._prepare_model_for_generation() + + inputs = self._tokenize_prompt(prompt, max_new_tokens, add_special_tokens=add_special_tokens) - inputs = {k: v.to(self.model.device) for k, v in inputs.items()} + stop_strings = [s for s in (stop_strings or []) if s] - stop_strings = stop_strings or [] + gen_kwargs: Dict[str, Any] = { + "max_new_tokens": max_new_tokens, + "temperature": temperature if do_sample else 1.0, + "top_p": top_p, + "top_k": top_k, + "do_sample": do_sample, + "repetition_penalty": repetition_penalty, + } + # Only fill in pad_token_id when the model's generation_config has none. + # eos_token_id is deliberately NOT passed: generation_config may define + # several EOS ids (e.g. llama-3 <|eot_id|>, gemma-2 ). + generation_config = getattr(self.model, "generation_config", None) + if getattr(generation_config, "pad_token_id", None) is None: + pad_token_id = getattr(self.tokenizer, "pad_token_id", None) + if pad_token_id is None: + pad_token_id = getattr(self.tokenizer, "eos_token_id", None) + if pad_token_id is not None: + gen_kwargs["pad_token_id"] = pad_token_id + if stop_strings: + # Native transformers stop strings halt generation at the token level. + gen_kwargs["stop_strings"] = stop_strings + gen_kwargs["tokenizer"] = self.tokenizer + gen_kwargs.update(kwargs) if stream: - return self._generate_streaming( - inputs, - max_new_tokens, - temperature, - top_p, - top_k, - do_sample, - repetition_penalty, - stop_strings, - **kwargs, - ) + return self._generate_streaming(inputs, gen_kwargs, stop_strings) with torch.inference_mode(): - outputs = self.model.generate( - **inputs, - max_new_tokens=max_new_tokens, - temperature=temperature if do_sample else 1.0, - top_p=top_p, - top_k=top_k, - do_sample=do_sample, - pad_token_id=self.tokenizer.pad_token_id, - eos_token_id=self.tokenizer.eos_token_id, - repetition_penalty=repetition_penalty, - **kwargs, - ) + outputs = self.model.generate(**inputs, **gen_kwargs) generated = outputs[0][inputs["input_ids"].shape[1] :] response = self.tokenizer.decode(generated, skip_special_tokens=True) - for stop in stop_strings: - if stop in response: - response = response.split(stop)[0] - break + idx = _find_first_stop(response, stop_strings) + if idx is not None: + response = response[:idx] return response.strip() + def _prepare_model_for_generation(self) -> None: + """Put the model in inference mode (e.g. after ``finetune()``).""" + model = self.model + try: + if getattr(model, "is_gradient_checkpointing", False) is True: + model.gradient_checkpointing_disable() + if callable(getattr(model, "eval", None)): + model.eval() + except Exception as e: # pragma: no cover - defensive + logger.debug(f"Could not switch model to eval mode: {e}") + + def _input_device(self) -> Any: + """Device of the input embeddings (robust with ``device_map="auto"``).""" + try: + return self.model.get_input_embeddings().weight.device + except Exception: + return getattr(self.model, "device", torch.device("cpu")) + + def _tokenize_prompt(self, prompt: str, max_new_tokens: int, add_special_tokens: bool = True) -> Dict[str, Any]: + """Tokenize a prompt, truncating from the left so the generation prompt survives.""" + tok_kwargs: Dict[str, Any] = {"return_tensors": "pt", "add_special_tokens": add_special_tokens} + max_seq_length = getattr(self.config, "max_seq_length", None) + if max_seq_length: + tok_kwargs["truncation"] = True + tok_kwargs["max_length"] = max(1, int(max_seq_length) - int(max_new_tokens)) + + old_side = getattr(self.tokenizer, "truncation_side", None) + try: + if old_side is not None: + self.tokenizer.truncation_side = "left" + inputs = self.tokenizer(prompt, **tok_kwargs) + finally: + if old_side is not None: + self.tokenizer.truncation_side = old_side + + inputs = dict(inputs) + if "attention_mask" not in inputs: + inputs["attention_mask"] = torch.ones_like(inputs["input_ids"]) + device = self._input_device() + return {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()} + def _generate_streaming( self, inputs: Dict, - max_new_tokens: int, - temperature: float, - top_p: float, - top_k: int, - do_sample: bool, - repetition_penalty: float, + gen_kwargs: Dict[str, Any], stop_strings: List[str], - **kwargs, ): - """Generate with streaming output (yields token strings).""" - try: - from transformers import TextIteratorStreamer - from threading import Thread - import sys - - streamer = TextIteratorStreamer( - self.tokenizer, - skip_prompt=True, - skip_special_tokens=True, - ) - - generation_kwargs = { - **inputs, - "max_new_tokens": max_new_tokens, - "temperature": temperature if do_sample else 1.0, - "top_p": top_p, - "top_k": top_k, - "do_sample": do_sample, - "pad_token_id": self.tokenizer.pad_token_id, - "eos_token_id": self.tokenizer.eos_token_id, - "repetition_penalty": repetition_penalty, - "streamer": streamer, - **kwargs, - } - - thread = Thread(target=self.model.generate, kwargs=generation_kwargs) - thread.start() - - generated_text = [] - for new_text in streamer: - if new_text: - yield new_text - generated_text.append(new_text) + """Generate with streaming output (yields text chunks). + + Stop strings are never yielded; generation is halted as soon as one is + seen, and exceptions raised in the generation thread are re-raised here. + """ + from threading import Event, Thread + + from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer + + gen_kwargs = dict(gen_kwargs) + timeout = gen_kwargs.pop("stream_timeout", 600.0) + + streamer = TextIteratorStreamer( + self.tokenizer, + skip_prompt=True, + timeout=timeout, + skip_special_tokens=True, + ) + + stop_event = Event() - full_text = "".join(generated_text) - should_stop = any(stop in full_text for stop in stop_strings) - if should_stop: - break + class _EventStoppingCriteria(StoppingCriteria): + def __call__(self, input_ids, scores, **kw): + return torch.full((input_ids.shape[0],), stop_event.is_set(), dtype=torch.bool, device=input_ids.device) + criteria = StoppingCriteriaList(list(gen_kwargs.pop("stopping_criteria", None) or [])) + criteria.append(_EventStoppingCriteria()) + + generation_kwargs = {**inputs, **gen_kwargs, "streamer": streamer, "stopping_criteria": criteria} + + errors: List[BaseException] = [] + + def _target(): + try: + self.model.generate(**generation_kwargs) + except BaseException as e: # noqa: BLE001 - re-raised in consumer + errors.append(e) + finally: + streamer.end() + + thread = Thread(target=_target, daemon=True) + thread.start() + + # Hold back enough characters to never leak a partially-generated stop string. + holdback = max((len(s) for s in stop_strings), default=1) - 1 + buffer = "" + emitted = 0 + stopped = False + try: + for new_text in streamer: + if not new_text: + continue + buffer += new_text + idx = _find_first_stop(buffer, stop_strings) + if idx is not None: + if idx > emitted: + yield buffer[emitted:idx] + emitted = idx + stopped = True + break + safe = len(buffer) - holdback + if safe > emitted: + yield buffer[emitted:safe] + emitted = safe + if not stopped and not errors and len(buffer) > emitted: + yield buffer[emitted:] + finally: + stop_event.set() thread.join() - except ImportError: - print("(Streaming not available, using batch generation)") - response = self.generate( - self.tokenizer.decode(inputs["input_ids"][0]), - max_new_tokens=max_new_tokens, - temperature=temperature, - stream=False, - ) - yield response + if errors: + raise errors[0] def chat( self, @@ -176,6 +232,7 @@ def chat( **kwargs: Additional generation parameters. """ prompt = None + templated = False if hasattr(self.tokenizer, "chat_template") and self.tokenizer.chat_template is not None: try: @@ -184,8 +241,11 @@ def chat( tokenize=False, add_generation_prompt=True, ) + templated = isinstance(prompt, str) + if not templated: + prompt = None except Exception: - pass + prompt = None if prompt is None: parts = [] @@ -201,24 +261,43 @@ def chat( parts.append("Assistant:") prompt = "".join(parts) + if templated: + # The rendered template already contains BOS; don't add it twice. + kwargs.setdefault("add_special_tokens", False) return self.generate(prompt, **kwargs) @staticmethod def _enable_flash_attention(model: PreTrainedModel, verbose: bool = True) -> None: - """Enable Flash Attention if available.""" + """Best-effort switch of the attention implementation on a loaded model. + + Mutating ``config._attn_implementation`` after load is a no-op (or + harmful) on most transformers versions, so this only acts through + ``model.set_attn_implementation`` when the model provides it. The + preferred path is passing ``attn_implementation=`` to + ``from_pretrained``. + """ + setter = getattr(model, "set_attn_implementation", None) + if not callable(setter): + logger.debug(" set_attn_implementation not available; leaving attention implementation unchanged") + return + try: - if hasattr(model, "config"): - if hasattr(model.config, "_attn_implementation"): - try: - import flash_attn - - model.config._attn_implementation = "flash_attention_2" - if verbose: - logger.info(" Flash Attention 2 configured") - except ImportError: - model.config._attn_implementation = "sdpa" - if verbose: - logger.info(" SDPA (Scaled Dot-Product Attention) configured") - except Exception: + import flash_attn # noqa: F401 + + impl = "flash_attention_2" + except ImportError: + impl = "sdpa" + + try: + setter(impl) + if verbose: + logger.info(f" Attention implementation set to {impl}") + except Exception as e: if verbose: - logger.warning(" Flash Attention / SDPA not available") + logger.warning(f" Could not set attention implementation to {impl}: {e}") + + +def _find_first_stop(text: str, stop_strings: List[str]) -> Optional[int]: + """Return the index of the earliest stop string occurrence in ``text``.""" + positions = [i for i in (text.find(s) for s in stop_strings if s) if i != -1] + return min(positions) if positions else None diff --git a/quantllm/core/hardware.py b/quantllm/core/hardware.py index 26349dc..942d210 100644 --- a/quantllm/core/hardware.py +++ b/quantllm/core/hardware.py @@ -44,15 +44,16 @@ class HardwareProfile: # Capabilities supports_bf16: bool supports_flash_attention: bool - best_gpu_index: int + best_gpu_index: int # CUDA device ordinal (GPUInfo.index), not a list position total_gpu_memory_gb: float @property def best_gpu(self) -> Optional[GPUInfo]: """Get the best GPU available.""" - if self.gpus: - return self.gpus[self.best_gpu_index] - return None + for gpu in self.gpus: + if gpu.index == self.best_gpu_index: + return gpu + return self.gpus[0] if self.gpus else None @property def device(self) -> torch.device: @@ -111,14 +112,15 @@ def detect(cls, force_refresh: bool = False) -> HardwareProfile: if gpu_info: gpus.append(gpu_info) - # Determine best GPU (most memory) + # Determine best GPU (most memory); store its CUDA ordinal best_gpu_index = 0 - if gpus: - best_gpu_index = max(range(len(gpus)), key=lambda i: gpus[i].memory_total_gb) + best = max(gpus, key=lambda g: g.memory_total_gb) if gpus else None + if best is not None: + best_gpu_index = best.index - # Aggregate capabilities - supports_bf16 = any(gpu.supports_bf16 for gpu in gpus) if gpus else False - supports_flash_attention = any(gpu.supports_flash_attention for gpu in gpus) if gpus else False + # Capabilities of the GPU that will be used by default + supports_bf16 = best.supports_bf16 if best is not None else False + supports_flash_attention = best.supports_flash_attention if best is not None else False total_gpu_memory = sum(gpu.memory_total_gb for gpu in gpus) profile = HardwareProfile( @@ -190,9 +192,10 @@ def _get_gpu_info(index: int) -> Optional[GPUInfo]: memory_free = memory_total # Approximation when not in use try: - # Try to get actual free memory - torch.cuda.set_device(index) - memory_free = (props.total_memory - torch.cuda.memory_allocated(index)) / (1024**3) + # Actual free memory on this device (includes other processes), + # without changing the current CUDA device. + free_bytes, _ = torch.cuda.mem_get_info(index) + memory_free = free_bytes / (1024**3) except Exception: pass diff --git a/quantllm/core/memory.py b/quantllm/core/memory.py index 95cca8d..680502f 100644 --- a/quantllm/core/memory.py +++ b/quantllm/core/memory.py @@ -223,6 +223,7 @@ def __init__( model: nn.Module, max_gpu_layers: Optional[int] = None, cpu_offload_threshold: float = 0.85, + device: Optional[Union[str, torch.device]] = None, ): """ Initialize dynamic offloader. @@ -231,7 +232,11 @@ def __init__( model: Model with layers to offload max_gpu_layers: Maximum layers to keep on GPU cpu_offload_threshold: GPU memory usage threshold to trigger offload + device: Accelerator device for on-GPU layers (default: cuda:0 if available, else cpu) """ + if device is None: + device = "cuda:0" if torch.cuda.is_available() else "cpu" + self.device = torch.device(device) self.model = model self.cpu_offload_threshold = cpu_offload_threshold self._layer_locations: Dict[str, str] = {} @@ -257,15 +262,22 @@ def __init__( self._initialize_placements() def _find_offloadable_layers(self) -> List[str]: - """Find layers that can be offloaded.""" + """Find the repeated transformer blocks (children of a ModuleList). + + Containers (e.g. ``model.layers``) are never returned, so moving a + block never moves the whole stack. + """ layers = [] for name, module in self.model.named_modules(): - # Look for transformer layers - if any(pattern in name.lower() for pattern in ["layer", "block", "decoder"]): - if hasattr(module, "parameters") and list(module.parameters()): - # Only include top-level repeated layers - if name.count(".") <= 2: # Avoid nested modules - layers.append(name) + if not isinstance(module, nn.ModuleList): + continue + if not any(pattern in name.lower() for pattern in ["layer", "block", "decoder", ".h"]): + continue + for child_name, child in module.named_children(): + if any(True for _ in child.parameters()): + layers.append(f"{name}.{child_name}" if name else child_name) + if layers: + break # only the first (outermost) block stack return layers def _estimate_avg_layer_size(self) -> float: @@ -285,7 +297,7 @@ def _initialize_placements(self): """Initialize layer placements.""" for i, name in enumerate(self.offloadable_layers): if i < self.max_gpu_layers: - self._layer_locations[name] = "cuda" + self._layer_locations[name] = "gpu" else: self._layer_locations[name] = "cpu" # Move to CPU @@ -294,11 +306,11 @@ def _initialize_placements(self): def move_layer_to_gpu(self, layer_name: str) -> None: """Move a layer to GPU, evicting others if needed.""" - if self._layer_locations.get(layer_name) == "cuda": + if self._layer_locations.get(layer_name) == "gpu": return # Already on GPU # Check if we need to evict - gpu_layers = [n for n, loc in self._layer_locations.items() if loc == "cuda"] + gpu_layers = [n for n, loc in self._layer_locations.items() if loc == "gpu"] if len(gpu_layers) >= self.max_gpu_layers: # Evict least recently accessed @@ -309,8 +321,8 @@ def move_layer_to_gpu(self, layer_name: str) -> None: # Move requested layer to GPU layer = dict(self.model.named_modules())[layer_name] - layer.to("cuda") - self._layer_locations[layer_name] = "cuda" + layer.to(self.device) + self._layer_locations[layer_name] = "gpu" # Update access order if layer_name in self._access_order: @@ -322,7 +334,8 @@ def _move_to_cpu(self, layer_name: str) -> None: layer = dict(self.model.named_modules())[layer_name] layer.to("cpu") self._layer_locations[layer_name] = "cpu" - torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() def prefetch_next_layer(self, current_layer_idx: int) -> None: """Prefetch the next layer to GPU asynchronously.""" @@ -371,15 +384,22 @@ def _enable_manual(model: nn.Module) -> bool: checkpointed = False for name, module in model.named_modules(): - if "layer" in name.lower() or "block" in name.lower(): - if hasattr(module, "forward"): - original_forward = module.forward - - def checkpointed_forward(*args, _orig=original_forward, **kwargs): - return checkpoint(_orig, *args, use_reentrant=False, **kwargs) - - module.forward = checkpointed_forward - checkpointed = True + # Only wrap the repeated blocks (children of a ModuleList), never the + # container itself or nested sub-modules (avoids nested checkpoints). + if not isinstance(module, nn.ModuleList): + continue + if not ("layer" in name.lower() or "block" in name.lower()): + continue + for block in module.children(): + original_forward = block.forward + + def checkpointed_forward(*args, _orig=original_forward, **kwargs): + return checkpoint(_orig, *args, use_reentrant=False, **kwargs) + + block.forward = checkpointed_forward # type: ignore[method-assign] + checkpointed = True + if checkpointed: + break return checkpointed diff --git a/quantllm/core/model_analyzer.py b/quantllm/core/model_analyzer.py index 289d025..d2840dc 100644 --- a/quantllm/core/model_analyzer.py +++ b/quantllm/core/model_analyzer.py @@ -81,13 +81,19 @@ class ModelAnalyzer: } @classmethod - def analyze(cls, model_name: str, trust_remote_code: bool = True) -> ModelInfo: + def analyze( + cls, + model_name: str, + trust_remote_code: bool = False, + token: Optional[str] = None, + ) -> ModelInfo: """ Analyze a model's architecture from its config. Args: model_name: HuggingFace model name or local path - trust_remote_code: Whether to trust remote code + trust_remote_code: Whether to trust remote code (default False) + token: Optional HuggingFace token for gated/private repos Returns: ModelInfo containing architecture details @@ -99,13 +105,13 @@ def analyze(cls, model_name: str, trust_remote_code: bool = True) -> ModelInfo: try: from transformers import AutoConfig - config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code) + config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code, token=token) info = cls._extract_info_from_config(model_name, config) cls._cache[model_name] = info return info - except Exception as e: + except Exception: # Fallback to estimation from model name return cls._estimate_from_name(model_name) @@ -134,6 +140,11 @@ def _extract_info_from_config(cls, model_name: str, config: Any) -> ModelInfo: ) model_type = getattr(config, "model_type", "unknown") + # Mixture-of-Experts: every expert holds its own MLP weights + num_experts = cls._get_config_attr(config, ["num_local_experts", "num_experts", "n_routed_experts"], None) + moe_intermediate_size = getattr(config, "moe_intermediate_size", None) + n_shared_experts = getattr(config, "n_shared_experts", None) + # Calculate total parameters num_params = cls._calculate_params( hidden_size=hidden_size, @@ -142,6 +153,9 @@ def _extract_info_from_config(cls, model_name: str, config: Any) -> ModelInfo: num_kv_heads=num_kv_heads, intermediate_size=intermediate_size, vocab_size=vocab_size, + num_experts=num_experts, + moe_intermediate_size=moe_intermediate_size, + n_shared_experts=n_shared_experts, ) # Size in GB at FP32 (4 bytes per param) @@ -169,8 +183,16 @@ def _calculate_params( num_kv_heads: int, intermediate_size: int, vocab_size: int, + num_experts: Optional[int] = None, + moe_intermediate_size: Optional[int] = None, + n_shared_experts: Optional[int] = None, ) -> int: - """Calculate total parameters in a transformer model.""" + """Calculate total parameters in a transformer model. + + For MoE models (``num_experts`` > 1) the MLP is counted once per + routed expert (using ``moe_intermediate_size`` when given), plus any + shared experts and the router. + """ head_dim = hidden_size // num_attention_heads # Embedding layer @@ -194,7 +216,13 @@ def _calculate_params( + hidden_size * hidden_size # O ) - mlp_params = hidden_size * intermediate_size * 3 # up, gate, down + if num_experts and num_experts > 1: + expert_size = moe_intermediate_size or intermediate_size + mlp_params = hidden_size * expert_size * 3 * num_experts # routed experts + mlp_params += hidden_size * expert_size * 3 * (n_shared_experts or 0) # shared experts + mlp_params += hidden_size * num_experts # router + else: + mlp_params = hidden_size * intermediate_size * 3 # up, gate, down norm_params = hidden_size * 4 # 2 norms with weight and bias layer_params = attention_params + mlp_params + norm_params @@ -221,11 +249,11 @@ def _estimate_from_name(cls, model_name: str) -> ModelInfo: num_params = params break - # Try to extract size from name (e.g., "7b", "70b") + # Try to extract size from name (e.g., "7b", "70b", "0.5b") if num_params is None: - match = re.search(r"(\d+)b", name_lower) + match = re.search(r"(\d+(?:\.\d+)?)[bB]\b", name_lower) if match: - num_params = int(match.group(1)) * 1_000_000_000 + num_params = int(float(match.group(1)) * 1_000_000_000) # Default to 7B if unknown if num_params is None: diff --git a/quantllm/core/model_gguf.py b/quantllm/core/model_gguf.py index ed1dfad..73f6f24 100644 --- a/quantllm/core/model_gguf.py +++ b/quantllm/core/model_gguf.py @@ -10,6 +10,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from .smart_config import SmartConfig +from ..utils.hf_compat import dtype_kwargs from ..utils import ( logger, print_header, @@ -87,8 +88,8 @@ def from_gguf( model = AutoModelForCausalLM.from_pretrained( model_id, gguf_file=filename, - torch_dtype=smart_config.dtype, - trust_remote_code=True, + **dtype_kwargs(smart_config.dtype), + trust_remote_code=kwargs.pop("trust_remote_code", False), **kwargs, ) except ImportError as e: diff --git a/quantllm/core/model_loader.py b/quantllm/core/model_loader.py index 0c8162d..e337e32 100644 --- a/quantllm/core/model_loader.py +++ b/quantllm/core/model_loader.py @@ -5,7 +5,9 @@ """ import copy +import importlib.util import os +import tempfile from dataclasses import asdict from typing import Any, Dict, Optional, Type @@ -17,6 +19,7 @@ PreTrainedModel, ) +from ..utils.hf_compat import dtype_kwargs from .architecture import ( DEFAULT_ARCHITECTURE_FALLBACKS, GGUF_NAME_MARKERS, @@ -45,7 +48,6 @@ class ModelLoaderMixin: """Model-loading methods mixed into :class:`TurboModel`.""" - _enable_flash_attention: Any _architecture_registry: Any _model_class_registry: Any resolve_model_type: Any @@ -74,6 +76,7 @@ def from_pretrained( config: Optional[Dict[str, Any]] = None, verbose: bool = True, token: Optional[str] = None, + attn_implementation: Optional[str] = None, ) -> "TurboModel": # type: ignore[name-defined] # noqa: F821 """Load a model with automatic optimization. @@ -82,6 +85,9 @@ def from_pretrained( - Chooses optimal quantization settings - Configures memory management - Enables speed optimizations + + ``attn_implementation`` overrides the automatic choice + (``"flash_attention_2"`` when usable, otherwise ``"sdpa"``). """ disable_hf_progress_bar() disable_ds_progress_bar() @@ -119,6 +125,8 @@ def from_pretrained( max_seq_length=max_length, device=device, dtype=dtype, + trust_remote_code=trust_remote_code, + token=token, ) if config_override: @@ -158,7 +166,7 @@ def from_pretrained( model_kwargs: Dict[str, Any] = { "trust_remote_code": trust_remote_code, - "torch_dtype": smart_config.dtype, + **dtype_kwargs(smart_config.dtype), } if token is not None: @@ -171,36 +179,34 @@ def from_pretrained( existing_quant = getattr(hf_config, "quantization_config", None) if existing_quant: - allow_requantize = False - - is_bnb = "BitsAndBytesConfig" in existing_quant.__class__.__name__ - is_8bit = getattr(existing_quant, "load_in_8bit", False) - - if is_bnb and is_8bit and smart_config.bits == 4: - allow_requantize = True - if verbose: - logger.info(" Re-quantizing 8-bit model to 4-bit") - - if not allow_requantize: - if verbose: - logger.warning( - "Model is already quantized (%s). Disabling dynamic quantization.", - existing_quant.__class__.__name__, - ) - quantize = False + if verbose: + logger.warning( + "Model is already quantized (%s). Disabling dynamic quantization.", + existing_quant.__class__.__name__, + ) + quantize = False - except Exception: - pass + except Exception as config_error: + logger.warning("Could not load config for '%s': %s", model_name, config_error) + quantization_applied = False if cls._should_apply_quantization(quantize, smart_config.bits, from_config_only): - model_kwargs.update(cls._get_quantization_kwargs(smart_config)) + quant_kwargs = cls._get_quantization_kwargs(smart_config) + quantization_applied = bool(quant_kwargs) + model_kwargs.update(quant_kwargs) if smart_config.cpu_offload: model_kwargs["device_map"] = "auto" - model_kwargs["offload_folder"] = "offload" + model_kwargs["offload_folder"] = tempfile.mkdtemp(prefix="quantllm_offload_") + max_memory = cls._offload_max_memory(smart_config.device) + if max_memory: + model_kwargs["max_memory"] = max_memory elif torch.cuda.is_available(): model_kwargs["device_map"] = {"": smart_config.device} + auto_attn = attn_implementation is None + model_kwargs["attn_implementation"] = attn_implementation or cls._choose_attn_implementation(smart_config) + with QuantLLMProgress() as p: if verbose: task = p.add_task("Downloading & Loading model...", total=None) @@ -213,6 +219,7 @@ def from_pretrained( model_type_override=model_type_override, base_model_fallback=base_model_fallback, from_config_only=from_config_only, + auto_attn_implementation=auto_attn, ) if verbose: @@ -224,21 +231,19 @@ def from_pretrained( if verbose: logger.info(" Gradient checkpointing enabled") - if smart_config.use_flash_attention: - cls._enable_flash_attention(model, verbose) - - # Only compile if model is NOT quantized (effective_loading_bits >= 16) - # torch.compile is not supported with quantized models + PEFT - if smart_config.compile_model and smart_config.effective_loading_bits >= 16: - try: - model = torch.compile(model, mode="reduce-overhead") # type: ignore[assignment] - if verbose: - logger.info(" torch.compile enabled") - except Exception as e: - if verbose: - print_warning(f"torch.compile failed: {e}") - elif smart_config.compile_model and smart_config.effective_loading_bits < 16: - if verbose: + # torch.compile is opt-in (compile_model=True via config_override) and + # never applied to quantized models. Only ``forward`` is compiled so the + # model stays a PreTrainedModel (PEFT / isinstance checks keep working). + if smart_config.compile_model: + if smart_config.effective_loading_bits >= 16 and not quantization_applied and not from_config_only: + try: + model.forward = torch.compile(model.forward) # type: ignore[method-assign] + if verbose: + logger.info(" torch.compile enabled (model.forward)") + except Exception as e: + if verbose: + print_warning(f"torch.compile failed: {e}") + elif verbose: logger.info(" torch.compile skipped (quantized model)") if verbose: @@ -281,6 +286,71 @@ def _should_apply_quantization( """Decide whether BitsAndBytes kwargs should be added at load time.""" return quantize and bits < 16 and not from_config_only + @staticmethod + def _choose_attn_implementation(smart_config: Any) -> str: + """Pick ``flash_attention_2`` when installed and supported, else ``sdpa``.""" + dtype = getattr(smart_config, "dtype", None) + if isinstance(dtype, str): + dtype = getattr(torch, dtype, None) + device = getattr(smart_config, "device", None) + device_type = device.type if isinstance(device, torch.device) else str(device or "").split(":")[0] + if ( + getattr(smart_config, "use_flash_attention", False) + and device_type == "cuda" + and dtype in (torch.float16, torch.bfloat16) + and importlib.util.find_spec("flash_attn") is not None + ): + return "flash_attention_2" + return "sdpa" + + @staticmethod + def _offload_max_memory(device: Any) -> Optional[Dict[Any, str]]: + """Build a ``max_memory`` map for CPU offload (~10% GPU headroom).""" + if not torch.cuda.is_available(): + return None + try: + index = torch.device(device).index if device is not None else None + if index is None: + index = torch.cuda.current_device() + free_bytes, _ = torch.cuda.mem_get_info(index) + except Exception as e: + logger.warning("Could not query free GPU memory for offloading: %s", e) + return None + max_memory: Dict[Any, str] = {index: f"{max(int(free_bytes * 0.9 / (1024**2)), 1)}MiB"} + try: + import psutil + + max_memory["cpu"] = f"{int(psutil.virtual_memory().available * 0.9 / (1024**2))}MiB" + except ImportError: + pass + return max_memory + + @staticmethod + def _is_architecture_error(error: BaseException) -> bool: + """True when ``error`` means transformers did not recognise the architecture.""" + if not isinstance(error, (ValueError, KeyError)): + return False + message = str(error).lower() + markers = ( + "model type", + "model_type", + "configuration class", + "unrecognized", + "does not recognize", + ) + return any(marker in message for marker in markers) + + @staticmethod + def _is_attn_implementation_error(error: BaseException) -> bool: + """True when ``error`` is about an unsupported attention implementation.""" + message = str(error).lower() + return isinstance(error, (ValueError, ImportError)) and ( + "attn_implementation" in message + or "scaled_dot_product_attention" in message + or "flash attention" in message + or "flashattention" in message + ) + @staticmethod def _looks_prequantized(model_name: str) -> Optional[str]: """Return a marker when the repo name suggests pre-quantization.""" @@ -307,8 +377,14 @@ def _load_model_with_fallback( model_type_override: Optional[str], base_model_fallback: bool, from_config_only: bool, + auto_attn_implementation: bool = False, ) -> PreTrainedModel: - """Load model with architecture fallback and optional config-only mode.""" + """Load model with architecture fallback and optional config-only mode. + + Only unrecognised-architecture errors trigger the fallback chain; + everything else (missing/gated repos, missing deps, OOM, ...) is + re-raised immediately. + """ config_model_type = (getattr(hf_config, "model_type", None) or "").lower().strip() is_registered_architecture = config_model_type in cls._architecture_registry if config_model_type else False resolved_model_type = cls.resolve_model_type( @@ -346,15 +422,38 @@ def _load_model_with_fallback( return AutoModelForCausalLM.from_config( resolved_config, trust_remote_code=trust_remote_code, - torch_dtype=model_kwargs.get("torch_dtype"), + **{k: v for k, v in model_kwargs.items() if k in ("dtype", "torch_dtype")}, ) try: - return AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + try: + return AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + except (ValueError, ImportError) as attn_error: + # The auto-selected attention backend is not supported by this + # architecture: retry once with the transformers default. + if not ( + auto_attn_implementation + and "attn_implementation" in model_kwargs + and cls._is_attn_implementation_error(attn_error) + ): + raise + logger.warning( + "attn_implementation=%r not supported (%s); retrying with the transformers default.", + model_kwargs["attn_implementation"], + attn_error, + ) + model_kwargs = dict(model_kwargs) + model_kwargs.pop("attn_implementation") + return AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) except Exception as primary_error: - if not base_model_fallback: + oom_error = getattr(torch.cuda, "OutOfMemoryError", None) + if isinstance(primary_error, (OSError, ImportError)) or ( + oom_error is not None and isinstance(primary_error, oom_error) + ): + raise + if not base_model_fallback or not cls._is_architecture_error(primary_error): raise - fallback_error = None + fallback_error: Optional[BaseException] = None if resolved_config is not None: fallback_kwargs = dict(model_kwargs) @@ -378,7 +477,9 @@ def _load_model_with_fallback( except Exception as fallback_registered_error: fallback_error = fallback_registered_error - error_details = f" Last fallback error: {fallback_error}" if fallback_error else "" + error_details = f"\nPrimary error: {type(primary_error).__name__}: {primary_error}" + if fallback_error: + error_details += f"\nLast fallback error: {type(fallback_error).__name__}: {fallback_error}" architecture_label = config_model_type or "" resolved_label = resolved_model_type or "" @@ -390,7 +491,7 @@ def _load_model_with_fallback( "2) Use model_type_override='llama' (or your compatible base family).\n" "3) Use from_config_only=True with a loadable config " "(usually trust_remote_code=True)." + error_details - ) from (fallback_error or primary_error) + ) from primary_error @staticmethod def _get_quantization_kwargs(config: SmartConfig) -> Dict[str, Any]: @@ -398,7 +499,28 @@ def _get_quantization_kwargs(config: SmartConfig) -> Dict[str, Any]: BitsAndBytes only supports 4-bit and 8-bit for model loading. Other bit widths (2, 3, 5, 6) are available during GGUF export. + Returns ``{}`` (full precision) when bitsandbytes or a CUDA/XPU + device is unavailable. """ + if importlib.util.find_spec("bitsandbytes") is None: + logger.warning( + "bitsandbytes is not installed; falling back to full " + "precision. Install with pip install bitsandbytes to " + "enable 4-bit / 8-bit quantization on CUDA." + ) + return {} + + xpu = getattr(torch, "xpu", None) + has_accelerator = torch.cuda.is_available() or bool(xpu is not None and xpu.is_available()) + if not has_accelerator: + logger.warning( + "BitsAndBytes quantization requires a CUDA (or XPU) device; " + "none was found, so the model is loaded in full precision." + ) + return {} + + cpu_offload = bool(getattr(config, "cpu_offload", False)) + try: from transformers import BitsAndBytesConfig @@ -408,6 +530,7 @@ def _get_quantization_kwargs(config: SmartConfig) -> Dict[str, Any]: bnb_4bit_compute_dtype=config.dtype, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, + llm_int8_enable_fp32_cpu_offload=cpu_offload, ) if config.bits < 4: logger.info( @@ -416,7 +539,10 @@ def _get_quantization_kwargs(config: SmartConfig) -> Dict[str, Any]: ) logger.info(" Tip: Export to GGUF for Q%s_K quantization!", config.bits) else: - quantization_config = BitsAndBytesConfig(load_in_8bit=True) + quantization_config = BitsAndBytesConfig( + load_in_8bit=True, + llm_int8_enable_fp32_cpu_offload=cpu_offload, + ) if config.bits != 8: logger.info( " BitsAndBytes supports 4/8-bit only. Using 8-bit for requested %s-bit.", diff --git a/quantllm/core/push.py b/quantllm/core/push.py index e78873f..e0d22d0 100644 --- a/quantllm/core/push.py +++ b/quantllm/core/push.py @@ -45,8 +45,10 @@ def push_to_hub( quantization: Quantization type (for gguf/onnx/mlx). commit_message: Commit message. license: License type (default: apache-2.0). + private: Create the repository as private if it does not exist. """ from ..hub import QuantLLMHubManager + from ..hub.model_card import base_model_display_name, sanitize_base_model format_lower = ( format @@ -58,14 +60,17 @@ def push_to_hub( DEFAULT_EXPORT_PUSH_CONFIG["push_quantization"], ) - base_model_full = self.model.config._name_or_path - model_name = base_model_full.split("/")[-1] + name_or_path = self.model.config._name_or_path + # Never leak local filesystem paths into the card's base_model. + base_model_full = sanitize_base_model(name_or_path) or base_model_display_name(name_or_path) + model_name = os.path.basename(os.path.normpath(name_or_path)) or "model" print_header(f"Pushing to {repo_id}") print_info(f"Format: {format_lower.upper()}") print_info(f"Base model: {base_model_full}") manager = QuantLLMHubManager(repo_id=repo_id, hf_token=token) + manager.private = private if format_lower == "gguf": quant_label = push_quantization or self.export_push_config.get( @@ -87,6 +92,7 @@ def push_to_hub( "quantization": quant_label.upper(), "base_model": base_model_full, "license": license, + "gguf_filename": filename, } ) manager._generate_model_card(format="gguf") diff --git a/quantllm/core/smart_config.py b/quantllm/core/smart_config.py index 80e6296..b910da0 100644 --- a/quantllm/core/smart_config.py +++ b/quantllm/core/smart_config.py @@ -9,7 +9,7 @@ from typing import Optional, Dict, Any, List import torch -from .hardware import HardwareProfiler, HardwareProfile +from .hardware import GPUInfo, HardwareProfiler, HardwareProfile from .model_analyzer import ModelAnalyzer, ModelInfo @@ -79,6 +79,8 @@ def detect( device: Optional[str] = None, dtype: Optional[str] = None, training: bool = False, + trust_remote_code: bool = False, + token: Optional[str] = None, ) -> "SmartConfig": """ Auto-detect optimal configuration for a model. @@ -90,6 +92,8 @@ def detect( device: Override device (auto-detect if None) dtype: Override dtype (auto-detect if None) training: Whether config is for training (uses more memory) + trust_remote_code: Forwarded to the config lookup (default False) + token: Optional HuggingFace token for gated/private repos Returns: SmartConfig with optimal settings @@ -98,7 +102,7 @@ def detect( hw = HardwareProfiler.detect() # Analyze model - model_info = ModelAnalyzer.analyze(model_name) + model_info = ModelAnalyzer.analyze(model_name, trust_remote_code=trust_remote_code, token=token) # Determine optimal settings config = cls() @@ -109,14 +113,17 @@ def detect( else: config.device = hw.device + # Capabilities must come from the GPU we will actually run on + selected_gpu = cls._selected_gpu(hw, config.device) + # Dtype - infer from device if not explicitly set if dtype is not None: config.dtype = getattr(torch, dtype) elif config.device.type == "cpu": config.dtype = torch.float32 - elif hw.supports_bf16 and config.device.type != "cpu": + elif selected_gpu is not None and selected_gpu.supports_bf16: config.dtype = torch.bfloat16 - elif hw.cuda_available and config.device.type != "cpu": + elif selected_gpu is not None: config.dtype = torch.float16 else: config.dtype = torch.float32 @@ -125,7 +132,7 @@ def detect( if bits is not None: config.bits = bits else: - config.bits = cls._choose_bits(hw, model_info, training) + config.bits = cls._choose_bits(hw, model_info, training, selected_gpu) # Quant type - maps to GGUF quantization types config.quant_type = cls._choose_quant_type(config.bits, hw) @@ -139,9 +146,9 @@ def detect( config.group_size = cls._choose_group_size(hw, model_info) # Memory settings - if hw.gpus: - available_memory = hw.best_gpu.memory_free_gb if hw.best_gpu else 8.0 - estimated_size = model_info.estimated_size_at_bits(config.bits) + if selected_gpu is not None: + available_memory = selected_gpu.memory_free_gb + estimated_size = model_info.estimated_size_at_bits(config.effective_loading_bits) # Need more headroom for training headroom_factor = 3.0 if training else 1.5 @@ -153,15 +160,17 @@ def detect( config.gradient_checkpointing = True # Speed optimizations - config.use_flash_attention = hw.supports_flash_attention - config.use_fused_kernels = hw.supports_bf16 # Requires Ampere+ + config.use_flash_attention = bool(selected_gpu and selected_gpu.supports_flash_attention) + config.use_fused_kernels = bool(selected_gpu and selected_gpu.supports_bf16) # Requires Ampere+ - # Only compile for larger GPUs (compilation overhead) - if hw.gpus and hw.total_gpu_memory_gb >= 16: - config.compile_model = True + # torch.compile is opt-in only (set compile_model=True via config_override); + # auto-enabling it breaks PEFT wrapping and does not speed up generate(). + config.compile_model = False # Batch size - config.batch_size = cls._optimal_batch_size(hw, model_info, config.bits, training) + config.batch_size = cls._optimal_batch_size( + hw, model_info, config.effective_loading_bits, training, selected_gpu + ) # Gradient accumulation to reach effective batch of 32 target_effective_batch = 32 @@ -177,13 +186,29 @@ def detect( config.stats = { "params": model_info.num_params, "fp16_size": model_info.size_gb / 2.0, # Approximate FP16 size - "quant_size": model_info.estimated_size_at_bits(config.bits), + "quant_size": model_info.estimated_size_at_bits(config.effective_loading_bits), } return config @staticmethod - def _choose_bits(hw: HardwareProfile, model_info: ModelInfo, training: bool) -> int: + def _selected_gpu(hw: HardwareProfile, device: torch.device) -> Optional[GPUInfo]: + """Return the GPUInfo for the CUDA device that will be used, if any.""" + if device.type != "cuda" or not hw.gpus: + return None + index = device.index if device.index is not None else hw.best_gpu_index + for gpu in hw.gpus: + if gpu.index == index: + return gpu + return hw.best_gpu + + @staticmethod + def _choose_bits( + hw: HardwareProfile, + model_info: ModelInfo, + training: bool, + selected_gpu: Optional[GPUInfo] = None, + ) -> int: """ Intelligently choose bit-width based on constraints. @@ -193,10 +218,12 @@ def _choose_bits(hw: HardwareProfile, model_info: ModelInfo, training: bool) -> This function returns the optimal LOADING bits (4 or 8), but also sets the quant_type for GGUF export which can be any bit width. """ - if not hw.gpus: - return 4 # Default for CPU (4-bit NF4) + if not hw.gpus or not hw.cuda_available: + # BitsAndBytes does not work reliably on CPU / MPS: load in full precision + return 16 - available_memory = hw.best_gpu.memory_free_gb if hw.best_gpu else 8.0 + gpu = selected_gpu or hw.best_gpu + available_memory = gpu.memory_free_gb if gpu else 8.0 # Reserve memory for activations reserve_factor = 3.0 if training else 1.5 @@ -245,12 +272,19 @@ def _choose_group_size(hw: HardwareProfile, model_info: ModelInfo) -> int: return 128 # Standard @staticmethod - def _optimal_batch_size(hw: HardwareProfile, model_info: ModelInfo, bits: int, training: bool) -> int: + def _optimal_batch_size( + hw: HardwareProfile, + model_info: ModelInfo, + bits: int, + training: bool, + selected_gpu: Optional[GPUInfo] = None, + ) -> int: """Determine optimal batch size for available memory.""" if not hw.gpus: return 1 # CPU is memory-limited - available_memory = hw.best_gpu.memory_free_gb if hw.best_gpu else 8.0 + gpu = selected_gpu or hw.best_gpu + available_memory = gpu.memory_free_gb if gpu else 8.0 model_size = model_info.estimated_size_at_bits(bits) # Rough estimation: diff --git a/quantllm/core/training.py b/quantllm/core/training.py index efdb09f..20faee6 100644 --- a/quantllm/core/training.py +++ b/quantllm/core/training.py @@ -9,7 +9,7 @@ import json import os from dataclasses import dataclass, field -from typing import Optional, Dict, Any, List, Union, Callable +from typing import Optional, Dict, Any, List, Tuple, Union, Callable import torch import torch.nn as nn @@ -507,12 +507,22 @@ def load_training_data( # Tokenize def tokenize_fn(examples): + templated = False if text_column == "__instruction_output__": - # Combine instruction and output fields + # Combine instruction (+ input when non-empty) and output fields + inputs = examples.get("input") or [None] * len(examples["instruction"]) texts = [ - f"### Instruction:\n{inst}\n\n### Response:\n{out}" - for inst, out in zip(examples["instruction"], examples["output"]) + _alpaca_text(inst, inp, out) + for inst, inp, out in zip(examples["instruction"], inputs, examples["output"]) ] + elif text_column.startswith("__messages__:"): + col = text_column.split(":", 1)[1] + templated = _has_chat_template(tokenizer) + try: + texts = [_render_messages(tokenizer, m, templated) for m in examples[col]] + except Exception: + templated = False + texts = [_render_messages(tokenizer, m, False) for m in examples[col]] else: texts = examples[text_column] result = tokenizer( @@ -520,8 +530,16 @@ def tokenize_fn(examples): truncation=True, max_length=max_length, padding="max_length", + add_special_tokens=not templated, ) - result["labels"] = result["input_ids"].copy() + # Padding positions must not contribute to the loss. + if "attention_mask" in result: + result["labels"] = [ + [tok if mask else -100 for tok, mask in zip(ids, attn)] + for ids, attn in zip(result["input_ids"], result["attention_mask"]) + ] + else: + result["labels"] = [list(ids) for ids in result["input_ids"]] return result tokenized = dataset.map( @@ -534,25 +552,135 @@ def tokenize_fn(examples): def _find_text_column(dataset) -> str: - """Find the text column in a dataset.""" + """Find the text column in a dataset. + + Preference: instruction+output (Alpaca, ``input`` folded in when + non-empty) > messages/conversations > plain text columns. + """ + columns = dataset.column_names + + if "instruction" in columns and "output" in columns: + # Need to combine - return special marker + return "__instruction_output__" + + for col in ("messages", "conversations"): + if col in columns: + return f"__messages__:{col}" + # Common text column names - text_columns = ["text", "content", "input", "prompt", "question"] + text_columns = ["text", "content", "prompt", "question", "instruction", "input"] for col in text_columns: - if col in dataset.column_names: + if col in columns: return col - # Try to find instruction/output format - if "instruction" in dataset.column_names: - if "output" in dataset.column_names: - # Need to combine - return special marker - return "__instruction_output__" - return "instruction" - # Return first string column - for col in dataset.column_names: + for col in columns: sample = dataset[0][col] if isinstance(sample, str): return col - raise ValueError(f"Could not find text column. Available: {dataset.column_names}") + raise ValueError(f"Could not find text column. Available: {columns}") + + +_ROLE_MAP = { + "human": "user", + "user": "user", + "gpt": "assistant", + "assistant": "assistant", + "model": "assistant", + "bot": "assistant", + "system": "system", +} + + +def _has_chat_template(tokenizer: Any) -> bool: + return getattr(tokenizer, "chat_template", None) is not None and callable( + getattr(tokenizer, "apply_chat_template", None) + ) + + +def _normalize_messages(messages: Any) -> List[Dict[str, str]]: + """Normalize OpenAI-style or ShareGPT-style (from/value) messages.""" + out = [] + for m in messages or []: + if not isinstance(m, dict): + continue + role = str(m.get("role") or m.get("from") or "user").lower() + content = m.get("content") + if content is None: + content = m.get("value", "") + out.append({"role": _ROLE_MAP.get(role, role), "content": "" if content is None else str(content)}) + return out + + +def _render_messages(tokenizer: Any, messages: Any, use_template: bool) -> str: + msgs = _normalize_messages(messages) + if use_template: + return tokenizer.apply_chat_template(msgs, tokenize=False) + return "".join(f"{m['role'].capitalize()}: {m['content']}\n" for m in msgs) + + +def _alpaca_text(instruction: Any, inp: Any, output: Any) -> str: + if inp is not None and str(inp).strip(): + return f"### Instruction:\n{instruction}\n\n### Input:\n{inp}\n\n### Response:\n{output}" + return f"### Instruction:\n{instruction}\n\n### Response:\n{output}" + + +def format_training_texts(examples: Dict[str, List[Any]], tokenizer: Any) -> Tuple[List[str], bool]: + """Turn a batch of raw examples into training texts. + + Returns ``(texts, templated)``. ``templated`` is True when the texts were + rendered with the tokenizer's chat template (so they already contain BOS + and must be tokenized with ``add_special_tokens=False``). Every text ends + with ``tokenizer.eos_token`` so the model learns to stop. + """ + columns = set(examples.keys()) + use_template = _has_chat_template(tokenizer) + templated = False + texts: List[str] + + if "instruction" in columns and "output" in columns: + instructions = examples["instruction"] + inputs = examples.get("input") or [None] * len(instructions) + rows = list(zip(instructions, inputs, examples["output"])) + texts = [] + if use_template: + try: + for inst, inp, out in rows: + user = f"{inst}\n\n{inp}" if inp is not None and str(inp).strip() else f"{inst}" + texts.append( + _render_messages( + tokenizer, + [{"role": "user", "content": user}, {"role": "assistant", "content": out}], + True, + ) + ) + templated = True + except Exception: + texts = [] + if not templated: + texts = [_alpaca_text(inst, inp, out) for inst, inp, out in rows] + elif "messages" in columns or "conversations" in columns: + col = "messages" if "messages" in columns else "conversations" + texts = [] + if use_template: + try: + texts = [_render_messages(tokenizer, m, True) for m in examples[col]] + templated = True + except Exception: + texts = [] + if not templated: + texts = [_render_messages(tokenizer, m, False) for m in examples[col]] + elif "text" in columns: + texts = ["" if t is None else str(t) for t in examples["text"]] + else: + keys = [k for k in examples.keys() if examples[k] and isinstance(examples[k][0], str)] + if not keys: + raise ValueError(f"Could not find text columns. Available: {list(examples.keys())}") + texts = [" ".join(examples[k][i] for k in keys) for i in range(len(examples[keys[0]]))] + + eos = getattr(tokenizer, "eos_token", None) + if isinstance(eos, str) and eos: + texts = [t if t.rstrip().endswith(eos) else t + eos for t in texts] + return texts, templated diff --git a/quantllm/core/turbo_model.py b/quantllm/core/turbo_model.py index f52ed41..30f8821 100644 --- a/quantllm/core/turbo_model.py +++ b/quantllm/core/turbo_model.py @@ -19,6 +19,7 @@ PreTrainedTokenizer, ) +from ..utils.hf_compat import dtype_kwargs from .architecture import ArchitectureMixin from .export_router import ExportRouterMixin from .finetune_mixin import FinetuneMixin @@ -196,90 +197,156 @@ def report(self) -> Dict[str, Any]: } def _dequantize_model(self) -> nn.Module: - """Dequantize a BitsAndBytes model to full precision for GGUF export.""" + """Dequantize a BitsAndBytes model to full precision for GGUF export. + + Works on the in-memory model so fine-tuned weights are preserved: + LoRA adapters are merged first, then ``model.dequantize()`` + (transformers) is used, then a manual bnb dequantization. Reloading + the base weights from the Hub is only a last resort and only when the + model was never fine-tuned. + """ import gc - model_name = getattr(self.model.config, "_name_or_path", None) + target_dtype = self.config.dtype if self.config.dtype in [torch.float16, torch.bfloat16] else torch.float16 + finetuned = bool(self._is_finetuned or self._lora_applied) - if model_name: - if self.verbose: - print_info(f"Reloading {model_name} in full precision...") + # 1. Merge LoRA adapters into the (quantized) base weights. + if self._lora_applied or hasattr(self.model, "merge_and_unload"): + try: + if self.verbose: + print_info("Merging LoRA weights before dequantization...") + self.model = self.model.merge_and_unload() + self._lora_applied = False + except Exception as e: + raise RuntimeError(f"Failed to merge LoRA adapters before dequantization: {e}") from e - target_dtype = self.config.dtype if self.config.dtype in [torch.float16, torch.bfloat16] else torch.float16 + model = self.model + errors: List[str] = [] + # 2. Native transformers dequantization (in place). + if hasattr(model, "dequantize"): try: - dequant_model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype=target_dtype, - trust_remote_code=True, - device_map="auto" if torch.cuda.is_available() else None, - low_cpu_mem_usage=True, - ) + try: + dequant_model = model.dequantize(dtype=target_dtype) # transformers >= 5 + except TypeError: + dequant_model = model.dequantize() # transformers 4.x + if dequant_model is None: + dequant_model = model + dequant_model = dequant_model.to(target_dtype) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() return dequant_model except Exception as e: + errors.append(f"model.dequantize(): {e}") if self.verbose: - print_warning(f"Failed to reload model: {e}") - print_info("Attempting in-place dequantization...") + print_warning(f"model.dequantize() failed: {e}") + print_info("Attempting manual in-place dequantization...") + # 3. Manual dequantization of bnb layers + copy of every other tensor. try: - import bitsandbytes as bnb - - target_dtype = self.config.dtype if self.config.dtype in [torch.float16, torch.bfloat16] else torch.float16 - - dequant_model = AutoModelForCausalLM.from_config( - self.model.config, - torch_dtype=target_dtype, - ) - - with torch.no_grad(): - for name, module in self.model.named_modules(): - if isinstance(module, bnb.nn.Linear4bit): - target_module = dict(dequant_model.named_modules()).get(name) - if target_module is not None and hasattr(target_module, "weight"): - weight = module.weight - if hasattr(weight, "dequantize"): - dequant_weight = weight.dequantize() - else: - # Linear4bit weight may have quant_state or quant_state attribute - quant_state = getattr(weight, "quant_state", None) - if quant_state is not None: - dequant_weight = bnb.functional.dequantize_4bit(weight.data, quant_state) - else: - dequant_weight = weight.data.to(target_dtype) - target_module.weight.data.copy_(dequant_weight.to(target_dtype)) - if module.bias is not None: - target_module.bias.data.copy_(module.bias.data.to(target_dtype)) - - elif isinstance(module, bnb.nn.Linear8bitLt): - target_module = dict(dequant_model.named_modules()).get(name) - if target_module is not None and hasattr(target_module, "weight"): - weight = module.weight - if hasattr(weight, "dequantize"): - dequant_weight = weight.dequantize() - else: - dequant_weight = weight.data.to(target_dtype) - target_module.weight.data.copy_(dequant_weight.to(target_dtype)) - if module.bias is not None: - target_module.bias.data.copy_(module.bias.data.to(target_dtype)) - + dequant_model = self._manual_bnb_dequantize(model, target_dtype) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() - return dequant_model - except Exception as e: - raise RuntimeError( - f"Failed to dequantize BitsAndBytes model: {e}\n\n" - "To export to GGUF, please reload your model without " - "quantization:\n" - " model = TurboModel.from_pretrained('your-model', " - "quantize=False)\n" - " model.export('gguf', quantization='Q4_K_M')" + errors.append(f"manual dequantization: {e}") + + # 4. Last resort: reload base weights from the Hub (never after fine-tuning). + model_name = getattr(model.config, "_name_or_path", None) + if model_name and not finetuned: + print_warning( + f"In-memory dequantization failed ({'; '.join(errors)}). " + f"Reloading base weights of {model_name} in full precision." ) + try: + return AutoModelForCausalLM.from_pretrained( + model_name, + **dtype_kwargs(target_dtype), + device_map="auto" if torch.cuda.is_available() else None, + low_cpu_mem_usage=True, + ) + except Exception as e: + errors.append(f"reload from {model_name}: {e}") + + hint = ( + "Fine-tuned weights cannot be recovered by reloading from the Hub; " + "save the adapter/model and re-export from a full-precision load." + if finetuned + else "To export to GGUF, reload your model without quantization:\n" + " model = TurboModel.from_pretrained('your-model', quantize=False)\n" + " model.export('gguf', quantization='Q4_K_M')" + ) + raise RuntimeError(f"Failed to dequantize BitsAndBytes model: {'; '.join(errors)}\n\n{hint}") + + @staticmethod + def _manual_bnb_dequantize(model: nn.Module, target_dtype: torch.dtype) -> nn.Module: + """Rebuild ``model`` in ``target_dtype`` with bnb layers dequantized.""" + import copy + + import bitsandbytes as bnb + + new_state: Dict[str, torch.Tensor] = {} + bnb_modules = set() + + with torch.no_grad(): + for name, module in model.named_modules(): + if isinstance(module, bnb.nn.Linear4bit): + weight = module.weight + dequant = bnb.functional.dequantize_4bit(weight.data, weight.quant_state) # type: ignore[attr-defined] + elif isinstance(module, bnb.nn.Linear8bitLt): + weight = module.weight + scb = getattr(weight, "SCB", None) + if scb is None: + scb = getattr(getattr(module, "state", None), "SCB", None) + if scb is None: + if weight.dtype == torch.int8: + raise RuntimeError(f"Missing SCB scales for 8-bit layer '{name}'") + dequant = weight.data + else: + dequant = weight.data.float() * scb.float()[:, None] / 127.0 + else: + continue + bnb_modules.add(name) + new_state[f"{name}.weight"] = dequant.to("cpu", target_dtype) + if module.bias is not None: + new_state[f"{name}.bias"] = module.bias.data.to("cpu", target_dtype) + + def _owned_by_bnb(key: str) -> bool: + parts = key.split(".") + return any(".".join(parts[:i]) in bnb_modules for i in range(1, len(parts))) + + for key, tensor in model.state_dict().items(): + if key in new_state or _owned_by_bnb(key): + continue + tensor = tensor.detach().to("cpu") + new_state[key] = tensor.to(target_dtype) if tensor.is_floating_point() else tensor + + config = copy.deepcopy(model.config) + if hasattr(config, "quantization_config"): + delattr(config, "quantization_config") + dequant_model = AutoModelForCausalLM.from_config(config, **dtype_kwargs(target_dtype)) + result = dequant_model.load_state_dict(new_state, strict=False) + missing = [k for k in result.missing_keys if not k.endswith("lm_head.weight")] + if missing: + raise RuntimeError(f"Dequantized state dict is missing {len(missing)} tensors (e.g. {missing[:3]})") + if hasattr(dequant_model, "tie_weights"): + dequant_model.tie_weights() + return dequant_model def optimize_inference(self, backend: str = "triton", bits: int = 4): - """Optimize model for inference using high-performance kernels.""" + """Optimize model for inference using high-performance kernels. + + No-op for BitsAndBytes-quantized models and when CUDA is unavailable. + """ + if not torch.cuda.is_available(): + print_warning("optimize_inference requires a CUDA GPU; skipping.") + return + if self._is_bnb_quantized(): + print_warning("optimize_inference is a no-op for BitsAndBytes-quantized models; skipping.") + return + if backend == "triton": from ..kernels.triton import ( TritonQuantizedLinear, @@ -304,6 +371,9 @@ def _replace_with_triton(self, module: nn.Module, bits: int) -> int: count = 0 for name, child in module.named_children(): + # bnb layers subclass nn.Linear but hold packed weights; keep lm_head in full precision. + if name == "lm_head" or type(child).__name__ in ("Linear4bit", "Linear8bitLt"): + continue if isinstance(child, nn.Linear): if self.verbose: print_info(f"Quantizing {name}...") @@ -397,6 +467,6 @@ def turbo( dtype=dtype, base_model_fallback=base_model_fallback, config=config, - trust_remote_code=True, + trust_remote_code=kwargs.pop("trust_remote_code", False), **kwargs, ) diff --git a/quantllm/hub/hub_manager.py b/quantllm/hub/hub_manager.py index 34a4962..7b5eb09 100644 --- a/quantllm/hub/hub_manager.py +++ b/quantllm/hub/hub_manager.py @@ -12,10 +12,18 @@ ) from .model_card import ModelCardGenerator, generate_model_card import os +import shutil +import tempfile class QuantLLMHubManager: - def __init__(self, repo_id: str, hf_token: Optional[str] = None, organization: Optional[str] = None): + def __init__( + self, + repo_id: str, + hf_token: Optional[str] = None, + organization: Optional[str] = None, + private: bool = False, + ): """ Initialize HubManager for lifecycle management. @@ -23,15 +31,19 @@ def __init__(self, repo_id: str, hf_token: Optional[str] = None, organization: O repo_id: HuggingFace repository ID (e.g., "username/model") hf_token: HuggingFace API token organization: Optional organization name + private: Create the repository as private if it does not exist """ self.repo_id = repo_id self.organization = organization self.token = hf_token + self.private = private self.api = HfApi(token=hf_token) # State self.hyperparameters: dict[str, Any] = {} - self.staging_dir = f"./hub_staging/{repo_id.split('/')[-1]}" + # Fresh, unique staging dir per manager so stale files from earlier + # pushes are never re-uploaded. Owner is kept in the prefix. + self.staging_dir = "" self._ensure_staging_dir() # Login if token provided @@ -39,8 +51,14 @@ def __init__(self, repo_id: str, hf_token: Optional[str] = None, organization: O self.login() def _ensure_staging_dir(self): - if not os.path.exists(self.staging_dir): - os.makedirs(self.staging_dir) + if not self.staging_dir or not os.path.isdir(self.staging_dir): + prefix = "quantllm_" + self.repo_id.replace("/", "__") + "_" + self.staging_dir = tempfile.mkdtemp(prefix=prefix) + + def cleanup(self): + """Remove the staging directory and everything in it.""" + if self.staging_dir and os.path.isdir(self.staging_dir): + shutil.rmtree(self.staging_dir, ignore_errors=True) def login(self): """Login to Hugging Face Hub.""" @@ -83,7 +101,7 @@ def save_final_model(self, model, tokenizer=None, format: str = "safetensors"): # Get base model name if hasattr(model_obj, "config") and hasattr(model_obj.config, "_name_or_path"): base_model = model_obj.config._name_or_path - self.hyperparameters["base_model"] = base_model + self.hyperparameters.setdefault("base_model", base_model) else: model_obj = model tokenizer_obj = tokenizer @@ -128,6 +146,8 @@ def _generate_model_card(self, format: str = "safetensors"): quantization=quantization, license=license_type, tags=extra_tags, + gguf_filename=self.hyperparameters.get("gguf_filename"), + family=self.hyperparameters.get("family"), ) with open(readme_path, "w", encoding="utf-8") as f: @@ -135,10 +155,24 @@ def _generate_model_card(self, format: str = "safetensors"): print_info(f"Generated model card for {format.upper()} format") - def push(self, commit_message: str = "Upload model via QuantLLM"): - """Push the staged model to HuggingFace Hub.""" - # Disable HF bars - os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + def push( + self, + commit_message: str = "Upload model via QuantLLM", + private: Optional[bool] = None, + cleanup: bool = True, + ): + """Push the staged model to HuggingFace Hub. + + Args: + commit_message: Commit message. + private: Create the repo as private (defaults to the value given at init). + cleanup: Remove the staging directory after a successful upload. + """ + from huggingface_hub.utils import disable_progress_bars + + disable_progress_bars() + if private is None: + private = self.private try: print_header(f"Pushing to {self.repo_id}", icon="[PUSH]") @@ -146,7 +180,7 @@ def push(self, commit_message: str = "Upload model via QuantLLM"): # Ensure repo exists if not self.api.repo_exists(self.repo_id): print_info(f"Creating repository {self.repo_id}") - self.api.create_repo(repo_id=self.repo_id, token=self.token, exist_ok=True) + self.api.create_repo(repo_id=self.repo_id, token=self.token, private=private, exist_ok=True) # Calculate total size total_size = 0 @@ -169,10 +203,14 @@ def push(self, commit_message: str = "Upload model via QuantLLM"): except Exception as e: print_error(f"Error pushing to Hub: {str(e)}") + print_info(f"Staged files kept at {self.staging_dir} (call cleanup() to remove)") raise + if cleanup: + self.cleanup() + # Legacy compatibility def push_model(self, *args, **kwargs): - """Legacy method alias — delegates to push().""" + """Legacy method alias -- delegates to push().""" print_warning("push_model() is deprecated. Use push() instead.") return self.push(*args, **kwargs) diff --git a/quantllm/hub/model_card.py b/quantllm/hub/model_card.py index 7fbb0a5..e493e1f 100644 --- a/quantllm/hub/model_card.py +++ b/quantllm/hub/model_card.py @@ -5,9 +5,32 @@ usage examples, and "Use this model" button compatibility. """ +import os +import re from typing import Optional, Dict, Any, List from datetime import datetime +_HUB_REPO_ID_RE = re.compile(r"^[\w.-]+/[\w.-]+$") + + +def sanitize_base_model(base_model: Optional[str]) -> Optional[str]: + """Return *base_model* if it is a HuggingFace repo id, else ``None``. + + Local filesystem paths must never be written into ``base_model`` metadata. + """ + if not base_model or not isinstance(base_model, str): + return None + if _HUB_REPO_ID_RE.match(base_model) and not os.path.isdir(base_model): + return base_model + return None + + +def base_model_display_name(base_model: Optional[str]) -> str: + """Human-readable name: the repo id, or the last path component of a local path.""" + if not base_model: + return "unknown" + return sanitize_base_model(base_model) or os.path.basename(os.path.normpath(base_model)) or "unknown" + class ModelCardGenerator: """ @@ -62,7 +85,9 @@ def __init__( tags: Additional tags """ self.repo_id = repo_id - self.base_model = base_model + # Only a real Hub repo id may appear in metadata / links. + self.base_model_id = sanitize_base_model(base_model) + self.base_model = base_model_display_name(base_model) self.format = format.lower() self.quantization = quantization self.license = license @@ -106,9 +131,14 @@ def _generate_yaml_header(self) -> str: yaml_lines = [ "---", f"license: {self.license}", - f"base_model: {self.base_model}", - f"library_name: {library_name}", ] + if self.base_model_id: + yaml_lines.append(f"base_model: {self.base_model_id}") + yaml_lines.append(f"library_name: {library_name}") + # Registry family lives in a dedicated key (not inferred from tags[0]). + family = self.extra_params.get("family") + if family: + yaml_lines.append(f"quantllm_family: {str(family).lower()}") # Add language if self.language: @@ -129,6 +159,11 @@ def _generate_yaml_header(self) -> str: return "\n".join(yaml_lines) + def _base_model_link(self) -> str: + if self.base_model_id: + return f"[{self.base_model_id}](https://huggingface.co/{self.base_model_id})" + return self.base_model + def _generate_header_banner(self) -> str: """Generate prominent header with QuantLLM branding.""" model_name = self.repo_id.split("/")[-1] @@ -173,11 +208,11 @@ def _generate_description(self) -> str: return f""" ## [ABOUT] About This Model -This model is **[{self.base_model}](https://huggingface.co/{self.base_model})** converted to {format_desc.get(self.format, self.format.upper() + " format")}. +This model is **{self._base_model_link()}** converted to {format_desc.get(self.format, self.format.upper() + " format")}. | Property | Value | |----------|-------| -| **Base Model** | [{self.base_model}](https://huggingface.co/{self.base_model}) | +| **Base Model** | {self._base_model_link()} | | **Format** | {self.format.upper()} | | **Quantization** | {self.quantization or "None (Full Precision)"} | | **License** | {self.license} | @@ -199,7 +234,8 @@ def _generate_gguf_usage(self) -> str: """Generate GGUF usage examples.""" model_name = self.repo_id.split("/")[-1] quant = self.quantization or "Q4_K_M" - filename = f"{model_name}.{quant}.gguf" + # Prefer the actual uploaded filename when the caller knows it. + filename = self.extra_params.get("gguf_filename") or f"{model_name}.{quant}.gguf" return f''' ## [START] Quick Start @@ -422,7 +458,7 @@ def _generate_details_section(self) -> str: | Property | Value | |----------|-------| -| **Original Model** | [{self.base_model}](https://huggingface.co/{self.base_model}) | +| **Original Model** | {self._base_model_link()} | | **Format** | {self.format.upper()} | | **Quantization** | {self.quantization or "Full Precision"} | | **License** | `{self.license}` | diff --git a/quantllm/kernels/cpu/__init__.py b/quantllm/kernels/cpu/__init__.py index 471f508..de71df2 100644 --- a/quantllm/kernels/cpu/__init__.py +++ b/quantllm/kernels/cpu/__init__.py @@ -69,7 +69,7 @@ def __init__( self.register_buffer( "qweight", - torch.empty(out_features, self.num_groups, in_features // (8 // bits), dtype=torch.uint8), + torch.empty(out_features, self.num_groups, group_size // (8 // bits), dtype=torch.uint8), ) self.register_buffer("scales", torch.empty(out_features, self.num_groups, dtype=torch.float32)) self.register_buffer("zeros", torch.empty(out_features, self.num_groups, dtype=torch.float32)) @@ -80,8 +80,12 @@ def __init__( def quantize(self, weight: torch.Tensor) -> None: w = weight.float().reshape(self.out_features, self.in_features) - if self.group_size < self.in_features: - w = w.reshape(self.out_features, self.num_groups, self.group_size) + # Pad the last group (repeating the final column, which leaves min/max unchanged) + # when group_size does not divide in_features or exceeds it. + pad = self.num_groups * self.group_size - self.in_features + if pad > 0: + w = torch.cat([w, w[:, -1:].expand(-1, pad)], dim=1) + w = w.reshape(self.out_features, self.num_groups, self.group_size) w_max = w.amax(dim=-1, keepdim=True) w_min = w.amin(dim=-1, keepdim=True) @@ -108,11 +112,13 @@ def _dequantize_weights(self) -> torch.Tensor: wq = self.qweight.float().reshape(self.out_features, self.num_groups, self.group_size) w = wq.float() * self.scales.unsqueeze(-1) + self.zeros.unsqueeze(-1) - return w.reshape(self.out_features, self.in_features).to(torch.float32) + w = w.reshape(self.out_features, -1)[:, : self.in_features] + return w.to(torch.float32) def forward(self, x: torch.Tensor) -> torch.Tensor: - w = self._dequantize_weights() - return torch.nn.functional.linear(x, w.to(x.device), self.bias) + w = self._dequantize_weights().to(device=x.device, dtype=x.dtype) + bias = self.bias.to(device=x.device, dtype=x.dtype) if self.bias is not None else None + return torch.nn.functional.linear(x, w, bias) class CPUInt8Linear(nn.Module): @@ -137,8 +143,9 @@ def quantize(self, weight: torch.Tensor) -> None: self.qweight = torch.round(w / self.scale).clamp(-128, 127).to(torch.int8) def forward(self, x: torch.Tensor) -> torch.Tensor: - w = self.qweight.float() * self.scale.to(x.device) - return torch.nn.functional.linear(x, w, self.bias) + w = (self.qweight.float() * self.scale.to(x.device)).to(device=x.device, dtype=x.dtype) + bias = self.bias.to(device=x.device, dtype=x.dtype) if self.bias is not None else None + return torch.nn.functional.linear(x, w, bias) def optimize_for_cpu(model: nn.Module, bits: int = 4, group_size: int = 128) -> nn.Module: diff --git a/quantllm/kernels/triton/quantized_linear.py b/quantllm/kernels/triton/quantized_linear.py index fe4de1e..fb71287 100644 --- a/quantllm/kernels/triton/quantized_linear.py +++ b/quantllm/kernels/triton/quantized_linear.py @@ -24,8 +24,8 @@ def is_triton_available() -> bool: - """Check if Triton is available for use.""" - return _TRITON_AVAILABLE + """Check if Triton is importable and a CUDA device is available to run its kernels.""" + return _TRITON_AVAILABLE and torch.cuda.is_available() def triton_q8_0_quantize(weight: torch.Tensor, eps: float = 1e-8) -> Tuple[torch.Tensor, torch.Tensor]: @@ -67,6 +67,7 @@ def int4_matmul( qweight: torch.Tensor, scales: torch.Tensor, bias: Optional[torch.Tensor] = None, + group_size: Optional[int] = None, ) -> torch.Tensor: """ INT4 matmul path backed by fused dequant+matmul on CUDA/Triton when available. @@ -76,11 +77,14 @@ def int4_matmul( qweight: Quantized int4 values stored in int8, shape [in_features, out_features] scales: Per-column scales, shape [1, out_features] or [in_features/group, out_features] bias: Optional bias [out_features] + group_size: Quantization group size. Inferred from shapes when omitted, which is + only exact when ``group_size`` divides ``in_features``. """ # Per-column case uses [1, N] zeros; grouped quantization uses zeros shaped like scales. is_per_column = scales.shape[0] == 1 zeros = scales.new_zeros((1, scales.shape[1])) if is_per_column else scales.new_zeros(scales.shape) - group_size = qweight.shape[0] if is_per_column else max(qweight.shape[0] // scales.shape[0], 1) + if group_size is None: + group_size = qweight.shape[0] if is_per_column else max(-(-qweight.shape[0] // scales.shape[0]), 1) return fused_dequant_matmul( x=x, qweight=qweight, @@ -308,14 +312,15 @@ def fused_dequant_matmul( >>> zeros = torch.zeros(4, 256, device='cuda', dtype=torch.float16) >>> out = fused_dequant_matmul(x, qw, scales, zeros, group_size=128) """ - if not _TRITON_AVAILABLE: + if not _TRITON_AVAILABLE or not x.is_cuda: # Fallback to PyTorch implementation return _pytorch_dequant_matmul(x, qweight, scales, zeros, bias, group_size) # Reshape input for 2D matmul original_shape = x.shape if x.dim() > 2: - x = x.view(-1, x.shape[-1]) + x = x.reshape(-1, x.shape[-1]) + x = x.contiguous() M, K = x.shape K_w, N = qweight.shape @@ -329,7 +334,6 @@ def grid(meta): return (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),) # Ensure contiguous tensors - x = x.contiguous() qweight = qweight.contiguous() scales = scales.contiguous() zeros = zeros.contiguous() @@ -395,14 +399,14 @@ def _pytorch_dequant_matmul( scale_expanded = scales.repeat_interleave(group_size, dim=0)[:K] zero_expanded = zeros.repeat_interleave(group_size, dim=0)[:K] - # Dequantize weights - weight_fp = (qweight.to(x.dtype) - zero_expanded) * scale_expanded + # Dequantize weights (in the input dtype so mixed-precision buffers don't break the matmul) + weight_fp = (qweight.to(x.dtype) - zero_expanded.to(x.dtype)) * scale_expanded.to(x.dtype) # Matmul output = torch.nn.functional.linear(x, weight_fp.t()) if bias is not None: - output = output + bias + output = output + bias.to(x.dtype) return output diff --git a/quantllm/quant/__init__.py b/quantllm/quant/__init__.py index 51bf6a7..0138101 100644 --- a/quantllm/quant/__init__.py +++ b/quantllm/quant/__init__.py @@ -14,6 +14,8 @@ ensure_llama_cpp_installed, detect_environment, get_llama_cpp_dir, + gguf_outtype, + IMATRIX_REQUIRED_TYPES, ) from .hqq import ( @@ -41,9 +43,9 @@ "q2_k", "iq4_nl", "iq4_xs", - "iq3_xxs", - "iq2_xxs", - "iq1_s", + "iq3_xxs", # requires imatrix + "iq2_xxs", # requires imatrix + "iq1_s", # requires imatrix ] QUANT_RECOMMENDATIONS = { @@ -63,6 +65,8 @@ "ensure_llama_cpp_installed", "get_llama_cpp_dir", "detect_environment", + "gguf_outtype", + "IMATRIX_REQUIRED_TYPES", "GGUF_QUANT_TYPES", "QUANT_RECOMMENDATIONS", "HQQLinear", diff --git a/quantllm/quant/hqq.py b/quantllm/quant/hqq.py index b1709d1..cc9c2dc 100644 --- a/quantllm/quant/hqq.py +++ b/quantllm/quant/hqq.py @@ -44,8 +44,10 @@ def __init__( bits: int = 4, group_size: int = 128, bias: bool = False, + optimize: bool = True, ): super().__init__() + self.optimize = optimize self.in_features = in_features self.out_features = out_features self.bits = bits @@ -75,21 +77,70 @@ def _unpack_4bit(self, wq: torch.Tensor) -> torch.Tensor: out = torch.stack([high, low], dim=-1).reshape(wq.shape[0], wq.shape[1], -1) return out[:, :, : self.group_size] - def quantize(self, weight: torch.Tensor) -> None: - device = weight.device - w = weight.float().reshape(self.out_features, self.in_features) + def _group(self, w: torch.Tensor) -> torch.Tensor: + """Reshape ``(out, in)`` -> ``(out, num_groups, group_size)``. - if self.group_size < self.in_features: - w = w.reshape(self.out_features, self.num_groups, self.group_size) - else: - # One group per output feature - w = w.unsqueeze(1) # (out_features, 1, in_features) + When ``group_size`` does not divide ``in_features`` (or exceeds it), + the last group is padded by repeating the final column so the padding + does not widen the group's min/max range. + """ + padded = self.num_groups * self.group_size + pad = padded - self.in_features + if pad > 0: + w = torch.cat([w, w[:, -1:].expand(-1, pad)], dim=1) + return w.reshape(self.out_features, self.num_groups, self.group_size) + @staticmethod + def _shrink_lp(x: torch.Tensor, beta: float, lp_norm: float) -> torch.Tensor: + """Generalized soft-thresholding (proximal operator of the l_p norm, p < 1).""" + return torch.sign(x) * torch.relu(torch.abs(x) - (1.0 / beta) * torch.abs(x).pow(lp_norm - 1)) + + def _optimize_zero( + self, + w: torch.Tensor, + scale: torch.Tensor, + zero: torch.Tensor, + iters: int = 20, + beta: float = 10.0, + kappa: float = 1.01, + lp_norm: float = 0.7, + ) -> torch.Tensor: + """Half-quadratic proximal optimization of the zero-point (HQQ, Badri & Shaji 2023). + + Works in the HQQ parametrisation ``w ~= (w_q - z) * s`` where ``z = -zero / s`` + and returns the optimized additive offset (``zero``) used by this layer. + """ + qmax = 2**self.bits - 1 + z = -zero / scale + best_z = z + best_err = torch.abs(w - (torch.round(w / scale + z).clamp(0, qmax) - z) * scale).mean() + for _ in range(iters): + w_q = torch.round(w / scale + z).clamp(0, qmax) + w_r = (w_q - z) * scale + w_e = self._shrink_lp(w - w_r, beta, lp_norm) + z = torch.mean(w_q - (w - w_e) / scale, dim=-1, keepdim=True) + beta *= kappa + + err = torch.abs(w - (torch.round(w / scale + z).clamp(0, qmax) - z) * scale).mean() + if err < best_err: + best_err = err + best_z = z + else: + break + return -best_z * scale + + def quantize(self, weight: torch.Tensor) -> None: + w = self._group(weight.float().reshape(self.out_features, self.in_features)) + + # Min-max initialization w_max = w.amax(dim=-1, keepdim=True) w_min = w.amin(dim=-1, keepdim=True) scale = torch.clamp((w_max - w_min) / (2**self.bits - 1), min=1e-10) zero = w_min + if self.optimize: + zero = self._optimize_zero(w, scale, zero) + wq = torch.round((w - zero) / scale).clamp(0, 2**self.bits - 1) if self.bits == 4: @@ -99,16 +150,17 @@ def quantize(self, weight: torch.Tensor) -> None: self.scales = scale.to(torch.float16) self.zeros = zero.to(torch.float16) - def dequantize(self) -> torch.Tensor: + def dequantize(self, dtype: torch.dtype = torch.float16) -> torch.Tensor: wq = self._unpack_4bit(self.qweight) if self.bits == 4 else self.qweight.float() - w = wq.float().to(self.scales.device) * self.scales + self.zeros - return w.reshape(self.out_features, self.in_features).to(torch.float16) + w = wq.float().to(self.scales.device) * self.scales.float() + self.zeros.float() + w = w.reshape(self.out_features, -1)[:, : self.in_features] + return w.to(dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: - w = self.dequantize() - # Preserve input dtype (e.g., bfloat16) for compatibility with unquantized layers like lm_head - target_dtype = x.dtype - return torch.nn.functional.linear(x.to(w.dtype), w, self.bias).to(target_dtype) + # Compute in the input dtype (e.g., bfloat16) so it composes with unquantized layers like lm_head + w = self.dequantize(dtype=x.dtype) + bias = self.bias.to(x.dtype) if self.bias is not None else None + return torch.nn.functional.linear(x, w, bias) class HQQQuantizer: @@ -140,7 +192,12 @@ def quantize_model( parent = getattr(parent, part) hqq = HQQLinear( - module.in_features, module.out_features, bits=bits, group_size=group_size, bias=module.bias is not None + module.in_features, + module.out_features, + bits=bits, + group_size=group_size, + bias=module.bias is not None, + optimize=self.config.optimize, ) hqq.to(module.weight.device) hqq.quantize(module.weight.data) diff --git a/quantllm/quant/llama_cpp.py b/quantllm/quant/llama_cpp.py index 6e121a1..c8c622a 100644 --- a/quantllm/quant/llama_cpp.py +++ b/quantllm/quant/llama_cpp.py @@ -51,6 +51,45 @@ def get_llama_cpp_dir() -> str: return os.path.join(cache_dir, "llama.cpp") +# Quantization types that llama-quantize refuses to produce without an importance matrix +IMATRIX_REQUIRED_TYPES = {"IQ1_S", "IQ1_M", "IQ2_XXS", "IQ2_XS", "IQ2_S", "IQ3_XXS"} + +# Full-precision GGUF types the converter can emit directly via --outtype +GGUF_PASSTHROUGH_OUTTYPES = { + "f16": "f16", + "fp16": "f16", + "float16": "f16", + "torch.float16": "f16", + "bf16": "bf16", + "bfloat16": "bf16", + "torch.bfloat16": "bf16", + "f32": "f32", + "fp32": "f32", + "float32": "f32", + "torch.float32": "f32", +} + + +def gguf_outtype(dtype: str) -> Optional[str]: + """Map a dtype / passthrough quant name to a converter ``--outtype`` (``None`` if not full precision).""" + return GGUF_PASSTHROUGH_OUTTYPES.get(str(dtype).lower()) + + +def _exe_name(name: str) -> str: + """Return the platform-specific executable file name.""" + return f"{name}.exe" if os.name == "nt" else name + + +def _binary_candidates(llama_dir: str, name: str) -> List[str]: + """Candidate locations for a llama.cpp binary: copied to root, MSVC multi-config, single-config.""" + exe = _exe_name(name) + return [ + os.path.join(llama_dir, exe), + os.path.join(llama_dir, "build", "bin", "Release", exe), + os.path.join(llama_dir, "build", "bin", exe), + ] + + def check_llama_cpp() -> Tuple[Optional[str], Optional[str]]: """ Check if llama.cpp is installed and return paths to executables. @@ -64,17 +103,10 @@ def check_llama_cpp() -> Tuple[Optional[str], Optional[str]]: llama_dir = get_llama_cpp_dir() # Check for quantizer executable - quantizer_candidates = [ - os.path.join(llama_dir, "llama-quantize"), - os.path.join(llama_dir, "llama-quantize.exe"), - os.path.join(llama_dir, "quantize"), - os.path.join(llama_dir, "quantize.exe"), - os.path.join(llama_dir, "build", "bin", "llama-quantize"), - os.path.join(llama_dir, "build", "bin", "quantize"), + quantizer_candidates = _binary_candidates(llama_dir, "llama-quantize") + _binary_candidates(llama_dir, "quantize") + if os.name != "nt": # Common linux install locations if system-wide - "/usr/local/bin/llama-quantize", - "/usr/bin/llama-quantize", - ] + quantizer_candidates += ["/usr/local/bin/llama-quantize", "/usr/bin/llama-quantize"] quantizer_path = None for candidate in quantizer_candidates: @@ -184,24 +216,14 @@ def install_llama_cpp( if gpu_support: print_info("CUDA support enabled") - # Try CMAKE first (newer method) + # llama.cpp only supports CMake builds (its Makefile was removed upstream) try: with QuantLLMProgress(style="spinner") as progress: task = progress.add_task("Compiling with CMAKE...", total=None) _compile_with_cmake(llama_dir, n_jobs, gpu_support, print_output) print_success("llama.cpp compiled successfully with CMAKE") - except Exception as e: - print_warning(f"CMAKE compilation failed: {e}") - print_info("Trying MAKE compilation...") - - # Fallback to MAKE (older method) - try: - with QuantLLMProgress(style="spinner") as progress: - task = progress.add_task("Compiling with MAKE...", total=None) - _compile_with_make(llama_dir, n_jobs, gpu_support, print_output) - print_success("llama.cpp compiled successfully with MAKE") - except Exception as e2: - raise RuntimeError(f"Both CMAKE and MAKE compilation failed. CMAKE: {e}, MAKE: {e2}") + except FileNotFoundError as e: + raise RuntimeError(f"CMake not found. Install CMake and a C++ compiler to build llama.cpp: {e}") from e # Step 4: Verify installation print_step(4, total_steps, "Verifying installation...") @@ -225,17 +247,15 @@ def _compile_with_cmake(llama_dir: str, n_jobs: int, gpu_support: bool, print_ou - Parallelizes with maximum CPU cores """ # Check if already compiled - quantizer_path = os.path.join(llama_dir, "llama-quantize") + quantizer_path = os.path.join(llama_dir, _exe_name("llama-quantize")) if os.path.exists(quantizer_path) and os.access(quantizer_path, os.X_OK): return True # Already compiled build_dir = os.path.join(llama_dir, "build") - # Only clean if build exists but is broken + # Only clean if build exists but is broken (MSVC puts binaries in build/bin/Release) if os.path.exists(build_dir): - # Check if build is valid - bin_path = os.path.join(build_dir, "bin", "llama-quantize") - if not os.path.exists(bin_path): + if not any(os.path.exists(p) for p in _binary_candidates(llama_dir, "llama-quantize")[1:]): shutil.rmtree(build_dir, ignore_errors=True) # Configure CMAKE with speed optimizations @@ -254,12 +274,7 @@ def _compile_with_cmake(llama_dir: str, n_jobs: int, gpu_support: bool, print_ou "-DLLAMA_BUILD_EXAMPLES=OFF", # Skip examples (we only need quantize) ] - subprocess.run( - configure_cmd, - check=True, - stdout=subprocess.DEVNULL if not print_output else None, - stderr=subprocess.DEVNULL if not print_output else None, - ) + _run_build_step(configure_cmd, print_output, "CMake configure") # Build only the quantize target (faster than building everything) build_cmd = [ @@ -273,45 +288,35 @@ def _compile_with_cmake(llama_dir: str, n_jobs: int, gpu_support: bool, print_ou "llama-quantize", # Only build what we need ] - subprocess.run( - build_cmd, - check=True, - stdout=subprocess.DEVNULL if not print_output else None, - stderr=subprocess.DEVNULL if not print_output else None, - ) + _run_build_step(build_cmd, print_output, "CMake build") # Copy executable to main directory - bin_dir = os.path.join(build_dir, "bin") - if os.path.exists(bin_dir): - for target in LLAMA_CPP_TARGETS: - src = os.path.join(bin_dir, target) + for target in LLAMA_CPP_TARGETS: + for src in _binary_candidates(llama_dir, target)[1:]: if os.path.exists(src): shutil.copy2(src, llama_dir) + break return True -def _compile_with_make(llama_dir: str, n_jobs: int, gpu_support: bool, print_output: bool): - """Compile llama.cpp using MAKE (fallback).""" - # Clean - subprocess.run( - ["make", "clean", "-C", llama_dir], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - - # Build - env = os.environ.copy() - if gpu_support: - env["LLAMA_CUDA"] = "1" - - subprocess.run( - ["make", "all", f"-j{n_jobs}", "-C", llama_dir], - env=env, - check=True, - stdout=subprocess.DEVNULL if not print_output else None, - stderr=subprocess.STDOUT if not print_output else None, +def _run_build_step(cmd: List[str], print_output: bool, step: str, tail_lines: int = 40) -> None: + """Run a build command. When quiet, capture its output and include the tail in the error on failure.""" + if print_output: + subprocess.run(cmd, check=True) + return + + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, ) + if result.returncode != 0: + tail = "\n".join((result.stdout or "").strip().splitlines()[-tail_lines:]) + raise RuntimeError( + f"{step} failed with exit code {result.returncode}\nOutput (last {tail_lines} lines):\n{tail}" + ) def quantize_gguf( @@ -320,6 +325,7 @@ def quantize_gguf( quant_type: str, quantizer_location: Optional[str] = None, print_output: bool = True, + imatrix: Optional[str] = None, ) -> str: """ Quantize a GGUF file using llama-quantize. @@ -330,6 +336,8 @@ def quantize_gguf( quant_type: Quantization type (e.g., "q4_k_m", "q8_0") quantizer_location: Path to llama-quantize executable (auto-detect if None) print_output: Whether to print quantization progress + imatrix: Path to an importance-matrix file (``llama-imatrix`` output). Required for + the very low-bit IQ types (e.g. ``iq3_xxs``, ``iq2_xxs``, ``iq1_s``). Returns: Path to output GGUF file @@ -344,6 +352,18 @@ def quantize_gguf( from ..utils import QuantLLMProgress, print_info import re + # Normalize quant type + quant_type = quant_type.upper() + + if quant_type in IMATRIX_REQUIRED_TYPES and not imatrix: + raise ValueError( + f"{quant_type} quantization requires an importance matrix. Generate one with " + "llama-imatrix (e.g. `llama-imatrix -m model.F16.gguf -f calibration.txt -o imatrix.dat`) " + "and pass imatrix='imatrix.dat'." + ) + if imatrix and not os.path.exists(imatrix): + raise FileNotFoundError(f"imatrix file not found: {imatrix}") + # Get quantizer if quantizer_location is None: quantizer_location, _ = check_llama_cpp() @@ -351,9 +371,6 @@ def quantize_gguf( if not os.path.exists(input_gguf): raise FileNotFoundError(f"Input GGUF file not found: {input_gguf}") - # Normalize quant type - quant_type = quant_type.upper() - # Get CPU count for threading n_threads = psutil.cpu_count() or 1 @@ -361,8 +378,10 @@ def quantize_gguf( input_size_mb = os.path.getsize(input_gguf) / (1024 * 1024) # Build command - cmd = [ - quantizer_location, + cmd = [quantizer_location] + if imatrix: + cmd += ["--imatrix", imatrix] + cmd += [ input_gguf, output_gguf, quant_type, @@ -454,6 +473,7 @@ def convert_to_gguf( is_gpt_oss: bool = False, max_shard_size: str = "50GB", print_output: bool = True, + output_dir: Optional[str] = None, ) -> Tuple[List[str], bool]: """ Convert HuggingFace model to GGUF using llama.cpp converter. @@ -470,6 +490,7 @@ def convert_to_gguf( is_gpt_oss: Whether this is a GPT-OSS model max_shard_size: Maximum shard size print_output: Whether to print conversion progress + output_dir: Directory to write the GGUF file to (defaults to ``input_folder``) Returns: Tuple of (list of output files, is_vlm_updated) @@ -516,22 +537,15 @@ def convert_to_gguf( "2. Or use a pre-trained model without BitsAndBytes quantization." ) - # Determine output filename + # Normalize model_dtype for llama.cpp's --outtype (f32/f16/bf16/q8_0/...) + model_dtype = gguf_outtype(model_dtype) or str(model_dtype).lower() + + # Determine output file (absolute, so the converter never writes into the CWD) if quantization_type and quantization_type.lower() != "none": - output_file = f"{model_name}.{quantization_type.upper()}.gguf" + output_name = f"{model_name}.{quantization_type.upper()}.gguf" else: - output_file = f"{model_name}.{model_dtype.upper()}.gguf" - - # Normalize model_dtype for llama.cpp - dtype_map = { - "float16": "f16", - "bfloat16": "bf16", - "float32": "f32", - "torch.float16": "f16", - "torch.bfloat16": "bf16", - "torch.float32": "f32", - } - model_dtype = dtype_map.get(str(model_dtype).lower(), model_dtype) + output_name = f"{model_name}.{model_dtype.upper()}.gguf" + output_file = os.path.abspath(os.path.join(output_dir or input_folder, output_name)) # Build conversion command cmd = [ @@ -557,13 +571,14 @@ def convert_to_gguf( process = subprocess.Popen( cmd, # type: ignore[arg-type] stdout=subprocess.PIPE, - stderr=subprocess.PIPE, + # convert_hf_to_gguf logs to stderr: merge it into stdout so one reader drains + # everything and a full stderr pipe can never deadlock the converter. + stderr=subprocess.STDOUT, universal_newlines=True, bufsize=1, ) - stdout_lines = [] - stderr_lines = [] + output_lines: List[str] = [] # Progress patterns for convert_hf_to_gguf.py output # Typically shows: "Loading model..." "Writing tensor [X/Y]..." @@ -579,7 +594,7 @@ def convert_to_gguf( if not line: break line = line.strip() - stdout_lines.append(line) + output_lines.append(line) # Try to find total tensors if total_tensors is None: @@ -599,31 +614,24 @@ def convert_to_gguf( progress.update(task, description="Writing GGUF...") progress.update(task, completed=100, description="Conversion complete!") - - # Read any remaining stderr - stderr_output = process.stderr.read() # type: ignore[union-attr] - if stderr_output: - stderr_lines = stderr_output.strip().split("\n") else: - stdout, stderr = process.communicate() + stdout, _ = process.communicate() if stdout: - stdout_lines = stdout.strip().split("\n") - if stderr: - stderr_lines = stderr.strip().split("\n") + output_lines = stdout.strip().split("\n") process.wait() if process.returncode != 0: - # Collect all error information - error_details = [] - if stdout_lines: - error_details.append("STDOUT:\n" + "\n".join(stdout_lines[-20:])) - if stderr_lines: - error_details.append("STDERR:\n" + "\n".join(stderr_lines[-20:])) + # Don't leave a partial/corrupt GGUF behind + if os.path.exists(output_file): + try: + os.remove(output_file) + except OSError: + pass error_msg = f"GGUF conversion failed with exit code {process.returncode}\n" - if error_details: - error_msg += "\n--- Conversion Output ---\n" + "\n".join(error_details) + if output_lines: + error_msg += "\n--- Conversion Output ---\n" + "\n".join(output_lines[-40:]) else: error_msg += "\nNo error output captured. Try running with print_output=True for more details." diff --git a/quantllm/registry.py b/quantllm/registry.py index 9572cd2..a48fb7e 100644 --- a/quantllm/registry.py +++ b/quantllm/registry.py @@ -11,10 +11,58 @@ from __future__ import annotations import json +import logging import os from typing import Any, Dict, List, Optional QUANTLLM_ORG = "QuantLLM" +FAMILY_KEY = "quantllm_family" + +logger = logging.getLogger("quantllm.registry") + +# Process-wide cache of Hub results, keyed by organization. +_HUB_CACHE: Dict[str, List[Dict[str, Any]]] = {} + + +def _hub_offline() -> bool: + """Return True when HF_HUB_OFFLINE is set.""" + try: + from huggingface_hub import constants + + if getattr(constants, "HF_HUB_OFFLINE", False): + return True + except ImportError: + pass + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in ("1", "true", "yes", "on") + + +def _model_id(m: Any) -> Optional[str]: + """Return the repo id of a ModelInfo (``id`` in hub>=1.0, ``modelId`` before).""" + return getattr(m, "id", None) or getattr(m, "modelId", None) + + +def _card_data_dict(m: Any) -> Dict[str, Any]: + card_data = getattr(m, "card_data", None) or getattr(m, "cardData", None) or {} + if hasattr(card_data, "to_dict"): + card_data = card_data.to_dict() + return card_data if isinstance(card_data, dict) else {} + + +def _parse_family(card_data: Dict[str, Any]) -> str: + """Read the model family from card data. + + Prefers the dedicated ``quantllm_family`` key, then falls back to the + first card tag that is not the generic ``quantllm`` tag. + """ + family = card_data.get(FAMILY_KEY) + if isinstance(family, str) and family: + return family.lower() + tags = card_data.get("tags") + if isinstance(tags, list): + for tag in tags: + if isinstance(tag, str) and tag and tag.lower() != "quantllm": + return tag.lower() + return "unknown" class ModelRegistry: @@ -29,39 +77,46 @@ def __init__(self, org: str = QUANTLLM_ORG): self._models: Optional[List[Dict[str, Any]]] = None def _fetch_from_hub(self) -> List[Dict[str, Any]]: - """Fetch model list from the QuantLLM HF organization.""" + """Fetch model list from the QuantLLM HF organization. + + Uses a single ``list_models`` call with card data, honours + ``HF_HUB_OFFLINE`` and caches results in memory for the process. + """ + if self.org in _HUB_CACHE: + return [dict(m) for m in _HUB_CACHE[self.org]] + if _hub_offline(): + return [] try: from huggingface_hub import HfApi + except ImportError: + return [] + try: api = HfApi() - models = api.list_models(author=self.org, sort="downloads", limit=50) + models = api.list_models(author=self.org, sort="downloads", limit=50, cardData=True) result = [] for m in models: - # Get full model info to access card_data (metadata from model card frontmatter) - try: - model_info = api.model_info(m.modelId) - card_data = getattr(model_info, "card_data", None) or {} - except Exception: - card_data = {} + model_id = _model_id(m) + if not model_id: + continue + card_data = _card_data_dict(m) result.append( { - "name": m.modelId, - "family": ( - card_data.get("tags", ["unknown"])[0] - if isinstance(card_data.get("tags"), list) - else "unknown" - ), + "name": model_id, + "family": _parse_family(card_data), "params": card_data.get("params", 0) or 0, - "formats": self._detect_formats(m.modelId), + "formats": self._detect_formats(model_id), "quality_score": card_data.get("quality_score", 3.0) or 3.0, "speed_score": card_data.get("speed_score", 3.0) or 3.0, "verified": card_data.get("verified", False) or False, "recommended": card_data.get("recommended_quant", "Q4_K_M") or "Q4_K_M", } ) - return result - except Exception: + except Exception as e: # network / auth / API errors -> fall back to built-in list + logger.debug("Failed to fetch QuantLLM registry from Hub: %s", e) return [] + _HUB_CACHE[self.org] = [dict(m) for m in result] + return result def _detect_formats(self, model_id: str) -> List[str]: mid = model_id.lower() @@ -167,14 +222,18 @@ def get_model(self, name: str) -> Optional[Dict[str, Any]]: return None def register_model(self, model_id: str, metadata: Dict[str, Any]) -> bool: - """Register a model by creating/updating its HF model card.""" + """Register a model by creating/updating its HF model card. + + If the repo already has a README, only its YAML metadata is updated + (the card body is preserved). + """ try: - from huggingface_hub import HfApi + from huggingface_hub import HfApi, ModelCard, metadata_update api = HfApi() - tags = [metadata.get("family", "llama")] - card_data = { - "tags": tags, + family = str(metadata.get("family", "llama")).lower() + card_data: Dict[str, Any] = { + FAMILY_KEY: family, "params": metadata.get("params", 0), "quality_score": metadata.get("quality_score", 3.0), "speed_score": metadata.get("speed_score", 3.0), @@ -182,9 +241,21 @@ def register_model(self, model_id: str, metadata: Dict[str, Any]) -> bool: "recommended_quant": metadata.get("recommended", "Q4_K_M"), } api.create_repo(model_id, exist_ok=True, private=False) - card = f"---\n{json.dumps(card_data, indent=2)}\n---\n# {model_id}\n\nQuantLLM-registered model." - api.upload_file(path_or_fileobj=card.encode(), path_in_repo="README.md", repo_id=model_id) + + if api.file_exists(model_id, "README.md"): + existing = ModelCard.load(model_id) + tags = list(getattr(existing.data, "tags", None) or []) + for tag in ("quantllm", family): + if tag not in tags: + tags.append(tag) + card_data["tags"] = tags + metadata_update(model_id, card_data, overwrite=True) + else: + card_data["tags"] = ["quantllm", family] + card = f"---\n{json.dumps(card_data, indent=2)}\n---\n# {model_id}\n\nQuantLLM-registered model." + api.upload_file(path_or_fileobj=card.encode(), path_in_repo="README.md", repo_id=model_id) self._models = None + _HUB_CACHE.pop(self.org, None) return True except Exception as e: try: diff --git a/quantllm/server/app.py b/quantllm/server/app.py index 556353a..c37575c 100644 --- a/quantllm/server/app.py +++ b/quantllm/server/app.py @@ -7,7 +7,7 @@ import json import logging import time -from typing import AsyncGenerator +from typing import Any, AsyncGenerator, List, Optional, Union logger = logging.getLogger("quantllm.server") @@ -16,19 +16,28 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel + from starlette.concurrency import iterate_in_threadpool, run_in_threadpool class _ChatMessage(BaseModel): role: str - content: str + # OpenAI allows a string, null, or a list of content parts ({"type": "text", "text": ...}). + content: Optional[Union[str, List[Any]]] = None + + def text(self) -> str: + if self.content is None: + return "" + if isinstance(self.content, str): + return self.content + return "".join(p.get("text", "") for p in self.content if isinstance(p, dict) and p.get("type") == "text") class _ChatRequest(BaseModel): model: str = "" - messages: list[_ChatMessage] + messages: List[_ChatMessage] max_tokens: int = 512 temperature: float = 0.7 top_p: float = 0.9 stream: bool = False - stop: list[str] | None = None + stop: Optional[Union[str, List[str]]] = None class _CompletionRequest(BaseModel): model: str = "" @@ -37,13 +46,21 @@ class _CompletionRequest(BaseModel): temperature: float = 0.7 top_p: float = 0.9 stream: bool = False - stop: list[str] | None = None + stop: Optional[Union[str, List[str]]] = None HAS_FASTAPI = True except ImportError: HAS_FASTAPI = False +def _normalize_stop(stop) -> Optional[List[str]]: + if stop is None: + return None + if isinstance(stop, str): + stop = [stop] + return [s for s in stop if s] or None + + def _build_app(engine) -> FastAPI: app = FastAPI(title="QuantLLM Server", version="2.2.0") app.add_middleware( @@ -74,14 +91,16 @@ async def chat(req: _ChatRequest): raise HTTPException(503, "Model not loaded") from .engine import create_chat_completion - result = create_chat_completion( + # Blocking generation runs in a worker thread so the event loop (and /health) stays responsive. + result = await run_in_threadpool( + create_chat_completion, engine=engine, - messages=[m.model_dump() for m in req.messages], + messages=[{"role": m.role, "content": m.text()} for m in req.messages], max_tokens=req.max_tokens, temperature=req.temperature, top_p=req.top_p, stream=req.stream, - stop=req.stop, + stop=_normalize_stop(req.stop), ) if req.stream and isinstance(result, dict) and result.get("type") == "stream": return StreamingResponse( @@ -97,14 +116,15 @@ async def completions(req: _CompletionRequest): raise HTTPException(503, "Model not loaded") from .engine import create_completion - result = create_completion( + result = await run_in_threadpool( + create_completion, engine=engine, prompt=req.prompt, max_tokens=req.max_tokens, temperature=req.temperature, top_p=req.top_p, stream=req.stream, - stop=req.stop, + stop=_normalize_stop(req.stop), ) if req.stream and isinstance(result, dict) and result.get("type") == "stream": return StreamingResponse( @@ -123,32 +143,51 @@ async def _stream_chunks(result: dict, is_chat: bool) -> AsyncGenerator[str, Non cid = result["id"] created = result["created"] - async for chunk in engine.generate( - prompt, - max_tokens=result["max_tokens"], - temperature=result["temperature"], - top_p=result["top_p"], - stop=result["stop"], - stream=True, - ): - delta = {"content": chunk} if is_chat else {"text": chunk} - data = { - "id": cid, - "object": "chat.completion.chunk" if is_chat else "text_completion.chunk", - "created": created, - "model": engine.config.model, - "choices": [{"index": 0, "delta": delta, "finish_reason": None}], - } - yield f"data: {json.dumps(data)}\n\n" - - data = { - "id": cid, - "object": "chat.completion.chunk" if is_chat else "text_completion.chunk", - "created": created, - "model": engine.config.model, - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], - } - yield f"data: {json.dumps(data)}\n\n" + from .engine import finish_reason + + obj = "chat.completion.chunk" if is_chat else "text_completion" + + def _event(choice: dict) -> str: + data = {"id": cid, "object": obj, "created": created, "model": engine.config.model, "choices": [choice]} + return f"data: {json.dumps(data)}\n\n" + + def _choice(text: Optional[str], reason: Optional[str], role: bool = False) -> dict: + if is_chat: + delta: dict = {"role": "assistant", "content": ""} if role else ({"content": text} if text else {}) + return {"index": 0, "delta": delta, "finish_reason": reason} + return {"index": 0, "text": text or "", "logprobs": None, "finish_reason": reason} + + if is_chat: + yield _event(_choice(None, None, role=True)) + + kwargs = {} + if "messages" in result: + kwargs = {"messages": result["messages"], "add_special_tokens": result["add_special_tokens"]} + collected = [] + try: + # engine.generate(stream=True) is a blocking iterator; pull each chunk from a worker thread. + chunks = await run_in_threadpool( + engine.generate, + prompt, + max_tokens=result["max_tokens"], + temperature=result["temperature"], + top_p=result["top_p"], + stop=result["stop"], + stream=True, + **kwargs, + ) + async for chunk in iterate_in_threadpool(chunks): + if chunk: + collected.append(chunk) + yield _event(_choice(chunk, None)) + except Exception as e: + logger.exception("Generation failed during streaming") + yield f"data: {json.dumps({'error': {'message': str(e), 'type': type(e).__name__}})}\n\n" + yield "data: [DONE]\n\n" + return + + ct = engine.get_token_count("".join(collected)) + yield _event(_choice(None, finish_reason(ct, result["max_tokens"]))) yield "data: [DONE]\n\n" diff --git a/quantllm/server/engine.py b/quantllm/server/engine.py index 758f002..c284551 100644 --- a/quantllm/server/engine.py +++ b/quantllm/server/engine.py @@ -6,13 +6,16 @@ import json import os +import threading import time import uuid from dataclasses import dataclass, asdict -from typing import Any, AsyncGenerator, Dict, List, Optional +from typing import Any, AsyncGenerator, Dict, Generator, Iterator, List, Optional, Tuple import torch +from ..utils.hf_compat import dtype_kwargs + @dataclass class ServerConfig: @@ -24,7 +27,8 @@ class ServerConfig: dtype: str = "auto" quant_method: Optional[str] = None bits: Optional[int] = None - trust_remote_code: bool = True + trust_remote_code: bool = False + gguf_file: Optional[str] = None # filename glob inside a hub GGUF repo, e.g. "*Q4_K_M.gguf" @classmethod def from_cli(cls, **kwargs) -> "ServerConfig": @@ -45,6 +49,8 @@ def __init__(self, config: ServerConfig): self.model: Any = None self.tokenizer: Any = None self.backend = "transformers" + # Serializes model access: one generation (including a full stream) at a time. + self._lock = threading.Lock() def load(self) -> None: if self._detect_gguf(): @@ -52,9 +58,18 @@ def load(self) -> None: else: self._load_transformers() + def _is_local_gguf(self) -> bool: + path = self.config.model + return path.lower().endswith(".gguf") and os.path.isfile(path) + + def _is_hub_gguf(self) -> bool: + name = self.config.model + if not name or os.path.exists(name) or name.lower().endswith(".gguf"): + return False + return bool(self.config.gguf_file) or "gguf" in name.lower() + def _detect_gguf(self) -> bool: - name = self.config.model.lower() - if name.endswith(".gguf") or "gguf" in name: + if self._is_local_gguf() or self._is_hub_gguf(): try: import llama_cpp @@ -75,7 +90,7 @@ def _load_transformers(self) -> None: dtype = {"auto": "auto", "float16": torch.float16, "bfloat16": torch.bfloat16}.get(self.config.dtype, "auto") self.model = AutoModelForCausalLM.from_pretrained( self.config.model, - torch_dtype=dtype, + **dtype_kwargs(dtype), device_map="auto", trust_remote_code=self.config.trust_remote_code, ) @@ -84,7 +99,15 @@ def _load_transformers(self) -> None: def _load_llama_cpp(self) -> None: from llama_cpp import Llama - self.model = Llama(model_path=self.config.model, n_ctx=self.config.max_seq_length, verbose=False) + if self._is_local_gguf(): + self.model = Llama(model_path=self.config.model, n_ctx=self.config.max_seq_length, verbose=False) + else: + self.model = Llama.from_pretrained( + repo_id=self.config.model, + filename=self.config.gguf_file or "*Q4_K_M.gguf", + n_ctx=self.config.max_seq_length, + verbose=False, + ) self.backend = "llama-cpp-python" def warmup(self) -> None: @@ -103,72 +126,156 @@ def generate( top_p: float = 0.9, stop: Optional[List[str]] = None, stream: bool = False, + messages: Optional[List[Dict[str, Any]]] = None, + add_special_tokens: bool = True, ): - if self.backend == "llama-cpp-python": - return self._gen_llama_cpp(prompt, max_tokens, temperature, top_p, stop, stream) - return self._gen_transformers(prompt, max_tokens, temperature, top_p, stop, stream) - - def _gen_transformers( - self, prompt: str, max_tokens: int, temperature: float, top_p: float, stop: Optional[List[str]], stream: bool - ): - inputs = self.tokenizer(prompt, return_tensors="pt") - inputs = {k: v.to(self.model.device) for k, v in inputs.items()} + """Generate text; ``stream=True`` returns a blocking iterator of text chunks. + ``messages`` (chat requests) lets llama-cpp apply the GGUF-embedded chat template. + Pass ``add_special_tokens=False`` when ``prompt`` came from a chat template (it already has BOS). + """ if stream: - return self._stream_transformers(inputs, max_tokens, temperature, top_p, stop) + return self._stream(prompt, max_tokens, temperature, top_p, stop, messages, add_special_tokens) + with self._lock: + if self.backend == "llama-cpp-python": + return self._gen_llama_cpp(prompt, max_tokens, temperature, top_p, stop, messages) + return self._gen_transformers(prompt, max_tokens, temperature, top_p, stop, add_special_tokens) - outputs = self.model.generate( - **inputs, - max_new_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - do_sample=temperature > 0, - pad_token_id=self.tokenizer.pad_token_id, - eos_token_id=self.tokenizer.eos_token_id, - ) - text = self.tokenizer.decode(outputs[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True) - return self._apply_stops(text, stop) + def _stream( + self, + prompt: str, + max_tokens: int, + temperature: float, + top_p: float, + stop: Optional[List[str]], + messages: Optional[List[Dict[str, Any]]], + add_special_tokens: bool, + ) -> Iterator[str]: + # The lock is held for the whole stream so two requests never generate on one model concurrently. + with self._lock: + if self.backend == "llama-cpp-python": + chunks = self._stream_llama_cpp(prompt, max_tokens, temperature, top_p, stop, messages) + else: + inputs = self._encode(prompt, add_special_tokens) + chunks = self._stream_transformers(inputs, max_tokens, temperature, top_p, stop) + try: + yield from _filter_stops(chunks, stop) + finally: + chunks.close() - async def _stream_transformers( - self, inputs: Dict, max_tokens: int, temperature: float, top_p: float, stop: Optional[List[str]] - ): - from threading import Thread - from transformers import TextIteratorStreamer + def _encode(self, prompt: str, add_special_tokens: bool = True) -> Dict: + inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=add_special_tokens) + return {k: v.to(self.model.device) for k, v in inputs.items()} - streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True) - kwargs = { - **inputs, + def _gen_kwargs(self, max_tokens: int, temperature: float, top_p: float, stop: Optional[List[str]]) -> Dict: + # eos_token_id is intentionally not passed: generation_config may define several EOS ids. + kwargs: Dict[str, Any] = { "max_new_tokens": max_tokens, "temperature": temperature, "top_p": top_p, "do_sample": temperature > 0, "pad_token_id": self.tokenizer.pad_token_id, - "eos_token_id": self.tokenizer.eos_token_id, + } + if stop and _supports_stop_strings(): + kwargs["stop_strings"] = stop + kwargs["tokenizer"] = self.tokenizer + return kwargs + + def _gen_transformers( + self, + prompt: str, + max_tokens: int, + temperature: float, + top_p: float, + stop: Optional[List[str]], + add_special_tokens: bool = True, + ): + inputs = self._encode(prompt, add_special_tokens) + outputs = self.model.generate(**inputs, **self._gen_kwargs(max_tokens, temperature, top_p, stop)) + text = self.tokenizer.decode(outputs[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True) + return self._apply_stops(text, stop) + + def _stream_transformers( + self, inputs: Dict, max_tokens: int, temperature: float, top_p: float, stop: Optional[List[str]] + ) -> Generator[str, None, None]: + from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer + + streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=300) + cancel = threading.Event() + + class _Cancelled(StoppingCriteria): + def __call__(self, input_ids, scores, **kwargs): + return torch.full((input_ids.shape[0],), cancel.is_set(), dtype=torch.bool, device=input_ids.device) + + kwargs = { + **inputs, + **self._gen_kwargs(max_tokens, temperature, top_p, stop), "streamer": streamer, + "stopping_criteria": StoppingCriteriaList([_Cancelled()]), } - thread = Thread(target=self.model.generate, kwargs=kwargs) - thread.start() + errors: List[BaseException] = [] - collected = [] - for text in streamer: - collected.append(text) - if any(s in "".join(collected) for s in (stop or [])): - break - yield text - thread.join() + def _worker(): + try: + with torch.inference_mode(): + self.model.generate(**kwargs) + except BaseException as e: # re-raised in the consumer after the stream ends + errors.append(e) + finally: + streamer.end() # always unblock the consumer, even if generate() raised + + thread = threading.Thread(target=_worker, daemon=True) + thread.start() + try: + for text in streamer: + yield text + finally: + cancel.set() # halts generation if the consumer stopped early (stop string / disconnect) + thread.join() + if errors: + raise errors[0] def _gen_llama_cpp( - self, prompt: str, max_tokens: int, temperature: float, top_p: float, stop: Optional[List[str]], stream: bool + self, + prompt: str, + max_tokens: int, + temperature: float, + top_p: float, + stop: Optional[List[str]], + messages: Optional[List[Dict[str, Any]]] = None, ): - if stream: - return self._stream_llama_cpp(prompt, max_tokens, temperature, top_p, stop) + if messages is not None: + # create_chat_completion applies the chat template embedded in the GGUF metadata. + out = self.model.create_chat_completion( + messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, stop=stop or [] + ) + return out["choices"][0]["message"]["content"] or "" return self.model(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, stop=stop or [])[ "choices" ][0]["text"] - async def _stream_llama_cpp( - self, prompt: str, max_tokens: int, temperature: float, top_p: float, stop: Optional[List[str]] - ): + def _stream_llama_cpp( + self, + prompt: str, + max_tokens: int, + temperature: float, + top_p: float, + stop: Optional[List[str]], + messages: Optional[List[Dict[str, Any]]] = None, + ) -> Generator[str, None, None]: + if messages is not None: + for chunk in self.model.create_chat_completion( + messages=messages, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + stop=stop or [], + stream=True, + ): + content = chunk["choices"][0].get("delta", {}).get("content") + if content: + yield content + return for chunk in self.model( prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, stop=stop or [], stream=True ): @@ -184,14 +291,55 @@ def _apply_stops(text: str, stop: Optional[List[str]]) -> str: def get_token_count(self, text: str) -> int: if self.tokenizer is not None: - return len(self.tokenizer.encode(text)) + return len(self.tokenizer.encode(text, add_special_tokens=False)) + if self.backend == "llama-cpp-python" and self.model is not None: + return len(self.model.tokenize(text.encode("utf-8"), add_bos=False)) return len(text) // 4 -def _build_chat_prompt(engine: InferenceEngine, messages: List[Dict[str, str]]) -> str: +def _supports_stop_strings() -> bool: + """``generate(stop_strings=..., tokenizer=...)`` exists since transformers 4.39.""" + try: + import transformers + from packaging.version import Version + + return Version(transformers.__version__) >= Version("4.39.0") + except Exception: + return False + + +def _filter_stops(chunks: Iterator[str], stop: Optional[List[str]]) -> Iterator[str]: + """Yield streamed text up to (not including) the first stop string, even if it spans chunks.""" + if not stop: + yield from chunks + return + hold = max(len(s) for s in stop) - 1 # chars held back in case a stop string is split across chunks + text, emitted = "", 0 + for chunk in chunks: + text += chunk + hits = [i for i in (text.find(s, max(0, emitted - hold)) for s in stop) if i != -1] + if hits: + idx = min(hits) + if idx > emitted: + yield text[emitted:idx] + return + safe = len(text) - hold + if safe > emitted: + yield text[emitted:safe] + emitted = safe + if len(text) > emitted: + yield text[emitted:] + + +def finish_reason(completion_tokens: int, max_tokens: int) -> str: + return "length" if completion_tokens >= max_tokens else "stop" + + +def _chat_prompt(engine: InferenceEngine, messages: List[Dict[str, str]]) -> Tuple[str, bool]: + """Return ``(prompt, from_template)``; template prompts already contain special tokens (BOS).""" if engine.tokenizer is not None and hasattr(engine.tokenizer, "chat_template"): try: - return engine.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + return engine.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True), True except Exception: pass parts = [] @@ -199,7 +347,11 @@ def _build_chat_prompt(engine: InferenceEngine, messages: List[Dict[str, str]]) r, c = m.get("role", "user"), m.get("content", "") parts.append(f"{r.capitalize()}: {c}\n") parts.append("Assistant:") - return "".join(parts) + return "".join(parts), False + + +def _build_chat_prompt(engine: InferenceEngine, messages: List[Dict[str, str]]) -> str: + return _chat_prompt(engine, messages)[0] def create_chat_completion( @@ -211,7 +363,7 @@ def create_chat_completion( stream: bool = False, stop: Optional[List[str]] = None, ) -> Dict: - prompt = _build_chat_prompt(engine, messages) + prompt, from_template = _chat_prompt(engine, messages) cid = f"chatcmpl-{uuid.uuid4().hex}" created = int(time.time()) @@ -227,16 +379,33 @@ def create_chat_completion( "temperature": temperature, "top_p": top_p, "stop": stop, + "messages": messages, + "add_special_tokens": not from_template, } - text = engine.generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, stop=stop, stream=False) + text = engine.generate( + prompt, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + stop=stop, + stream=False, + messages=messages, + add_special_tokens=not from_template, + ) ct, pt = engine.get_token_count(text), engine.get_token_count(prompt) return { "id": cid, "object": "chat.completion", "created": created, "model": engine.config.model, - "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": finish_reason(ct, max_tokens), + } + ], "usage": {"prompt_tokens": pt, "completion_tokens": ct, "total_tokens": pt + ct}, } @@ -275,6 +444,6 @@ def create_completion( "object": "text_completion", "created": created, "model": engine.config.model, - "choices": [{"index": 0, "text": text, "finish_reason": "stop"}], + "choices": [{"index": 0, "text": text, "finish_reason": finish_reason(ct, max_tokens)}], "usage": {"prompt_tokens": pt, "completion_tokens": ct, "total_tokens": pt + ct}, } diff --git a/quantllm/utils/hf_compat.py b/quantllm/utils/hf_compat.py new file mode 100644 index 0000000..5c83981 --- /dev/null +++ b/quantllm/utils/hf_compat.py @@ -0,0 +1,22 @@ +"""Small shims over transformers API differences across supported versions.""" + +from functools import lru_cache +from typing import Any, Dict + + +@lru_cache(maxsize=1) +def _dtype_key() -> str: + # transformers 4.56 renamed from_pretrained/from_config's ``torch_dtype`` to ``dtype``; + # 5.x warns on the old name. We still support 4.46+, so pick per installed version. + try: + import transformers + + major, minor = (int(p) for p in transformers.__version__.split(".")[:2]) + return "dtype" if (major, minor) >= (4, 56) else "torch_dtype" + except Exception: + return "torch_dtype" + + +def dtype_kwargs(dtype: Any) -> Dict[str, Any]: + """Return ``{"dtype": dtype}`` or ``{"torch_dtype": dtype}`` for the installed transformers.""" + return {_dtype_key(): dtype} diff --git a/quantllm/utils/logging.py b/quantllm/utils/logging.py index 67f4824..71e0426 100644 --- a/quantllm/utils/logging.py +++ b/quantllm/utils/logging.py @@ -141,7 +141,9 @@ def configure_logging( resolved_max = int(_env("MAX_BYTES", str(max_bytes))) resolved_backup = int(_env("BACKUP", str(backup_count))) file_fmt = json_fmt if resolved_fmt == "json" else text_fmt - file_handler = RotatingFileHandler(resolved_file, maxBytes=resolved_max, backupCount=resolved_backup) + file_handler = RotatingFileHandler( + resolved_file, maxBytes=resolved_max, backupCount=resolved_backup, encoding="utf-8" + ) file_handler.setFormatter(file_fmt) logger.addHandler(file_handler) diff --git a/quantllm/utils/progress.py b/quantllm/utils/progress.py index e0abbf1..0d693a9 100644 --- a/quantllm/utils/progress.py +++ b/quantllm/utils/progress.py @@ -28,6 +28,29 @@ _T = TypeVar("_T") +def _make_streams_encoding_safe() -> None: + """Avoid UnicodeEncodeError on legacy consoles (e.g. Windows cp1252). + + Our own output is ASCII, but messages can embed arbitrary text (model + names, exception strings, subprocess output). For non-UTF-8 streams we + switch the error handler to ``replace`` so printing never crashes. + """ + for stream in (sys.stdout, sys.stderr): + encoding = (getattr(stream, "encoding", None) or "").lower().replace("-", "") + reconfigure = getattr(stream, "reconfigure", None) + if not encoding or encoding.startswith("utf") or reconfigure is None: + continue + if getattr(stream, "errors", "strict") != "strict": + continue + try: + reconfigure(errors="replace") + except Exception: + pass + + +_make_streams_encoding_safe() + + class _SimpleConsole: """Stand-in for ``rich.console.Console`` that writes to stderr via print. diff --git a/scripts/kaggle_gpu_tests.py b/scripts/kaggle_gpu_tests.py new file mode 100644 index 0000000..7564060 --- /dev/null +++ b/scripts/kaggle_gpu_tests.py @@ -0,0 +1,163 @@ +"""Run the GPU end-to-end suite (tests/gpu) on Kaggle's free GPUs from your terminal. + +Packs the local working tree (uncommitted changes included) into a private +Kaggle script kernel, runs it on a free GPU ("GPU T4 x2" by default), waits, and downloads the log. + + python scripts/kaggle_gpu_tests.py # fast tests + python scripts/kaggle_gpu_tests.py --slow # + GGUF/ONNX export tests + python scripts/kaggle_gpu_tests.py -k Finetune # pytest -k filter + python scripts/kaggle_gpu_tests.py --accelerator NvidiaL4 --out gpu-test-results-l4 + +One-time setup: verify your phone on kaggle.com (needed for GPU + internet), +then create an API token and save it to ~/.kaggle/access_token (or kaggle.json). +Results land in ./gpu-test-results/. Exit code is pytest's. +""" + +import argparse +import base64 +import io +import json +import pathlib +import subprocess +import sys +import tarfile +import tempfile +import time + +ROOT = pathlib.Path(__file__).resolve().parent.parent +PACK = ["quantllm", "tests", "pyproject.toml", "README.md", "LICENSE"] +SLUG = "quantllm-gpu-tests" + +KERNEL_TEMPLATE = r""" +import base64, io, os, shlex, subprocess, sys, tarfile + +PAYLOAD = "{payload}" +ARGS = {args!r} +SLOW = {slow!r} +OUT = "/kaggle/working" +REPO = "/kaggle/tmp/repo" + +def sh(cmd): + print("$", cmd, flush=True) + return subprocess.call(cmd, shell=True, executable="/bin/bash") + +os.makedirs(REPO, exist_ok=True) +tarfile.open(fileobj=io.BytesIO(base64.b64decode(PAYLOAD)), mode="r:gz").extractall(REPO) +os.chdir(REPO) +os.environ["QUANTLLM_BANNER"] = "0" +os.environ["PYTHONUNBUFFERED"] = "1" + +sh("nvidia-smi") +sh(f"nvidia-smi --query-gpu=name --format=csv,noheader > {{OUT}}/gpu.txt") +sh("pip install -q -e '.[server]' pytest httpx 'bitsandbytes>=0.43'") +if SLOW: + sh("pip install -q llama-cpp-python onnx onnxscript") +sh("python -c 'import torch,transformers,bitsandbytes as b;print(torch.__version__,transformers.__version__,b.__version__)'") + +mark = [] if SLOW else ["-m", "'not slow'"] +cmd = " ".join(["python -m pytest tests/gpu -v -rA --tb=short", f"--junitxml={{OUT}}/junit.xml", *mark, *map(shlex.quote, ARGS)]) +print("$", cmd, flush=True) +proc = subprocess.Popen(cmd, shell=True, executable="/bin/bash", stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) +with open(f"{{OUT}}/pytest.log", "w") as log: + for line in proc.stdout: + sys.stdout.write(line) + log.write(line) +rc = proc.wait() +open(f"{{OUT}}/exit_code.txt", "w").write(str(rc)) +print("PYTEST_EXIT_CODE", rc) +""" + + +def _pack() -> str: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + for name in PACK: + path = ROOT / name + if path.exists(): + tar.add(path, arcname=name, filter=lambda t: None if "__pycache__" in t.name else t) + return base64.b64encode(buf.getvalue()).decode() + + +def _kaggle(*args: str) -> str: + res = subprocess.run(["kaggle", *args], capture_output=True, text=True) + if res.returncode != 0: + raise SystemExit(f"kaggle {' '.join(args)} failed:\n{res.stdout}\n{res.stderr}") + return res.stdout + + +def _username() -> str: + out = _kaggle("config", "view") + for line in out.splitlines(): + if "username" in line.lower(): + return line.split(":", 1)[1].strip() + raise SystemExit("Could not determine Kaggle username; pass --user") + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--slow", action="store_true", help="include slow export tests") + p.add_argument("-k", default="", help="pytest -k expression") + p.add_argument( + "--accelerator", default="NvidiaTeslaT4", help="Kaggle accelerator, e.g. NvidiaTeslaT4 (T4 x2), NvidiaL4" + ) + p.add_argument("--user", default="", help="Kaggle username (default: from kaggle config)") + p.add_argument("--timeout", type=int, default=3 * 3600, help="max seconds to wait") + p.add_argument("--out", default=str(ROOT / "gpu-test-results")) + a = p.parse_args() + + user = a.user or _username() + ref = f"{user}/{SLUG}" + args = ["-k", a.k] if a.k else [] + + with tempfile.TemporaryDirectory() as d: + src = pathlib.Path(d) / "run.py" + src.write_text(KERNEL_TEMPLATE.format(payload=_pack(), args=args, slow=a.slow), encoding="utf-8") + meta = { + "id": ref, + "title": SLUG, + "code_file": "run.py", + "language": "python", + "kernel_type": "script", + "is_private": True, + "enable_gpu": True, + "enable_internet": True, + "dataset_sources": [], + "competition_sources": [], + "kernel_sources": [], + } + (pathlib.Path(d) / "kernel-metadata.json").write_text(json.dumps(meta), encoding="utf-8") + print(_kaggle("kernels", "push", "-p", d, "--accelerator", a.accelerator).strip()) + + print(f"Running https://www.kaggle.com/code/{ref} ...", flush=True) + start = time.time() + status = "" + while time.time() - start < a.timeout: + time.sleep(60) + status = _kaggle("kernels", "status", ref).lower() + print(f"[{int(time.time() - start) // 60:3d} min] {status.strip()}", flush=True) + if any(s in status for s in ("complete", "error", "cancel")): + break + + out = pathlib.Path(a.out) + out.mkdir(parents=True, exist_ok=True) + _kaggle("kernels", "output", ref, "-p", str(out)) + log = out / "pytest.log" + if log.exists(): + lines = log.read_text(encoding="utf-8", errors="replace").splitlines() + print("\n".join(lines[-60:])) + gpu_file = out / "gpu.txt" + gpus = gpu_file.read_text().strip().splitlines() if gpu_file.exists() else [] + print(f"\nRan on: {', '.join(gpus) or 'unknown'}") + wanted = a.accelerator.lower().replace("nvidia", "").replace("tesla", "") + if gpus and not any(wanted in g.lower().replace(" ", "") for g in gpus): + print( + f"WARNING: requested {a.accelerator} but Kaggle ran on {gpus[0]} (accelerator not available to this account)" + ) + code_file = out / "exit_code.txt" + rc = int(code_file.read_text().strip()) if code_file.exists() else 1 + print(f"\nKernel status: {status.strip()}\npytest exit code: {rc}\nFull output in {out}") + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/modal_gpu_tests.py b/scripts/modal_gpu_tests.py new file mode 100644 index 0000000..6949fe9 --- /dev/null +++ b/scripts/modal_gpu_tests.py @@ -0,0 +1,115 @@ +"""Run the GPU end-to-end suite (tests/gpu) on a real GPU via Modal. + +No local GPU needed. The local checkout is mounted into the container, so this +always tests your working tree, uncommitted changes included. + + modal run scripts/modal_gpu_tests.py # T4, fast tests + modal run scripts/modal_gpu_tests.py --gpu A10G --slow # include GGUF/ONNX + modal run scripts/modal_gpu_tests.py --gpu H100 -k Finetune + +Exit code is pytest's, so CI can gate on it. Requires `modal token new` once +(or MODAL_TOKEN_ID / MODAL_TOKEN_SECRET in the environment). +""" + +import os +import pathlib +import subprocess +import sys + +import modal + +ROOT = pathlib.Path(__file__).resolve().parent.parent +REMOTE_ROOT = "/repo" + +image = ( + modal.Image.debian_slim(python_version="3.11") + .apt_install("git", "build-essential", "cmake") + .pip_install( + "torch>=2.4", + "transformers>=4.45", + "accelerate>=0.34", + "peft>=0.12", + "bitsandbytes>=0.43", + "datasets>=2.14", + "sentencepiece", + "protobuf", + "einops", + "psutil", + "scipy", + "scikit-learn", + "huggingface-hub", + "gguf>=0.10.0", + "fastapi", + "httpx", + "uvicorn", + "pydantic>=2", + "pytest>=7.4", + ) + .env({"QUANTLLM_BANNER": "0", "HF_HOME": "/cache/hf", "PYTHONUNBUFFERED": "1"}) + .add_local_dir(ROOT / "quantllm", f"{REMOTE_ROOT}/quantllm") + .add_local_dir(ROOT / "tests", f"{REMOTE_ROOT}/tests") + .add_local_file(ROOT / "pyproject.toml", f"{REMOTE_ROOT}/pyproject.toml") +) + +# Heavier optional backends only for --slow runs (llama.cpp build takes minutes). +slow_image = image.pip_install("llama-cpp-python", "onnx", "onnxscript") + +app = modal.App("quantllm-gpu-tests") +hf_cache = modal.Volume.from_name("quantllm-hf-cache", create_if_missing=True) + +_secrets = [] +try: + _secrets.append(modal.Secret.from_name("huggingface", required_keys=["HF_TOKEN"])) +except Exception: # older modal versions or no secret configured + pass + + +def _pytest(args: list[str], model: str = "") -> int: + if model: + os.environ["QUANTLLM_GPU_TEST_MODEL"] = model + cmd = [sys.executable, "-m", "pytest", "tests/gpu", "-v", "-rA", "--tb=short", *args] + print("$", " ".join(cmd), flush=True) + rc = subprocess.call(cmd, cwd=REMOTE_ROOT) + subprocess.call(["nvidia-smi", "--query-gpu=name,memory.used,memory.total", "--format=csv"]) + hf_cache.commit() + return rc + + +_common = dict(volumes={"/cache": hf_cache}, timeout=60 * 60, secrets=_secrets) + + +@app.function(image=image, gpu="T4", **_common) +def run_t4(args: list[str], model: str = "") -> int: + return _pytest(args, model) + + +@app.function(image=image, gpu="A10G", **_common) +def run_a10g(args: list[str], model: str = "") -> int: + return _pytest(args, model) + + +@app.function(image=image, gpu="L4", **_common) +def run_l4(args: list[str], model: str = "") -> int: + return _pytest(args, model) + + +@app.function(image=image, gpu="H100", **_common) +def run_h100(args: list[str], model: str = "") -> int: + return _pytest(args, model) + + +@app.function(image=slow_image, gpu="A10G", **_common) +def run_slow_a10g(args: list[str], model: str = "") -> int: + return _pytest(args, model) + + +@app.local_entrypoint() +def main(gpu: str = "T4", slow: bool = False, k: str = "", model: str = ""): + args = [] if slow else ["-m", "not slow"] + if k: + args += ["-k", k] + runners = {"T4": run_t4, "A10G": run_a10g, "L4": run_l4, "H100": run_h100} + fn = run_slow_a10g if slow else runners[gpu.upper()] + rc = fn.remote(args, model) + print(f"\npytest exit code on {gpu if not slow else 'A10G (slow)'}: {rc}") + sys.exit(rc) diff --git a/scripts/run_gpu_tests.sh b/scripts/run_gpu_tests.sh new file mode 100644 index 0000000..405fb27 --- /dev/null +++ b/scripts/run_gpu_tests.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Run the GPU end-to-end suite on any machine with an NVIDIA GPU: +# a rented box, Kaggle (free 2xT4), Colab (free T4), Lightning, RunPod... +# +# Colab/Kaggle cell: +# !git clone https://github.com/codewithdark-git/QuantLLM && cd QuantLLM && bash scripts/run_gpu_tests.sh +# +# Extra pytest args pass through, e.g.: bash scripts/run_gpu_tests.sh -k Finetune +# Include slow export tests: SLOW=1 bash scripts/run_gpu_tests.sh +set -euo pipefail +cd "$(dirname "$0")/.." + +nvidia-smi --query-gpu=name,memory.total --format=csv || { echo "No NVIDIA GPU found"; exit 1; } + +pip install -q -e ".[server]" pytest httpx +if [ "${SLOW:-0}" = "1" ]; then + pip install -q llama-cpp-python onnx onnxscript # add optimum-onnx[onnxruntime] only with transformers<4.58 + MARK=() +else + MARK=(-m "not slow") +fi + +QUANTLLM_BANNER=0 python -m pytest tests/gpu -v -rA --tb=short "${MARK[@]}" "$@" diff --git a/tests/gpu/conftest.py b/tests/gpu/conftest.py new file mode 100644 index 0000000..4755508 --- /dev/null +++ b/tests/gpu/conftest.py @@ -0,0 +1,75 @@ +"""Fixtures for the GPU end-to-end suite. + +Every test under ``tests/gpu`` needs a CUDA device and downloads a real model, +so the whole directory is skipped on CPU-only machines. Run it locally on a GPU +box with ``pytest tests/gpu -m gpu`` or remotely with +``modal run scripts/modal_gpu_tests.py``. +""" + +import gc +import os + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("CUDA not available; GPU suite skipped", allow_module_level=True) + +os.environ.setdefault("QUANTLLM_BANNER", "0") + +# Small, ungated, chat-tuned model with a chat template. Override to test others. +GPU_TEST_MODEL = os.environ.get("QUANTLLM_GPU_TEST_MODEL", "Qwen/Qwen2.5-0.5B-Instruct") + + +def pytest_collection_modifyitems(config, items): + for item in items: + if "tests/gpu" in item.nodeid.replace("\\", "/"): + item.add_marker(pytest.mark.gpu) + + +def free_cuda(): + gc.collect() + torch.cuda.empty_cache() + + +@pytest.fixture(scope="session") +def model_id(): + return GPU_TEST_MODEL + + +def _load(model_id, **kwargs): + from quantllm import turbo + + return turbo(model_id, verbose=False, **kwargs) + + +@pytest.fixture(scope="module") +def model_4bit(model_id): + m = _load(model_id, bits=4) + yield m + del m + free_cuda() + + +@pytest.fixture(scope="module") +def model_16bit(model_id): + m = _load(model_id, bits=16) + yield m + del m + free_cuda() + + +@pytest.fixture +def load_model(model_id): + """Factory for tests that need a fresh model they can mutate (finetune/export).""" + loaded = [] + + def _factory(**kwargs): + m = _load(model_id, **kwargs) + loaded.append(m) + return m + + yield _factory + loaded.clear() + free_cuda() diff --git a/tests/gpu/test_gpu_e2e.py b/tests/gpu/test_gpu_e2e.py new file mode 100644 index 0000000..9db8798 --- /dev/null +++ b/tests/gpu/test_gpu_e2e.py @@ -0,0 +1,266 @@ +"""End-to-end GPU tests against a real (small) HuggingFace model. + +These check behaviour, not just "it didn't crash": that quantization really +replaced layers, that greedy decoding is deterministic, that fine-tuning lowers +the loss, and that exported artifacts load back and produce the same output. +""" + +import gc +import importlib.util +import shutil + +import pytest +import torch + + +def free_cuda(): + gc.collect() + torch.cuda.empty_cache() + + +PROMPT = "The capital of France is" + + +def _bnb_layer_count(model, cls_name): + return sum(1 for m in model.modules() if type(m).__name__ == cls_name) + + +# --------------------------------------------------------------------------- # +# Loading / quantization +# --------------------------------------------------------------------------- # + + +class TestLoading: + def test_4bit_replaces_linear_layers(self, model_4bit): + assert model_4bit.is_quantized + assert _bnb_layer_count(model_4bit.model, "Linear4bit") > 0 + assert next(model_4bit.model.parameters()).device.type == "cuda" + + def test_4bit_uses_less_memory_than_16bit(self, model_4bit, model_16bit): + q = model_4bit.model.get_memory_footprint() + f = model_16bit.model.get_memory_footprint() + assert q < 0.6 * f, f"4-bit footprint {q / 1e9:.2f}GB not much smaller than 16-bit {f / 1e9:.2f}GB" + + def test_8bit_loads(self, load_model): + m = load_model(bits=8) + assert m.is_quantized + assert _bnb_layer_count(m.model, "Linear8bitLt") > 0 + assert isinstance(m.generate(PROMPT, max_new_tokens=8, do_sample=False), str) + + def test_16bit_dtype_matches_hardware(self, model_16bit): + dtype = next(model_16bit.model.parameters()).dtype + # is_bf16_supported() is True on pre-Ampere GPUs via emulation; native bf16 needs sm_80+. + expected = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16 + assert dtype == expected + + def test_report_is_consistent(self, model_4bit): + r = model_4bit.report() + assert r["is_quantized"] is True + assert r["effective_loading_bits"] == 4 + assert "cuda" in str(r["device"]) + + +# --------------------------------------------------------------------------- # +# Generation +# --------------------------------------------------------------------------- # + + +class TestGeneration: + @pytest.mark.parametrize("fixture", ["model_4bit", "model_16bit"]) + def test_greedy_is_deterministic_and_excludes_prompt(self, fixture, request): + m = request.getfixturevalue(fixture) + a = m.generate(PROMPT, max_new_tokens=16, do_sample=False) + b = m.generate(PROMPT, max_new_tokens=16, do_sample=False) + assert a == b + assert a.strip() + assert not a.startswith(PROMPT), "generate() should return only the completion" + + def test_greedy_answer_is_sensible(self, model_16bit): + out = model_16bit.generate(PROMPT, max_new_tokens=8, do_sample=False) + assert "paris" in out.lower() + + def test_max_new_tokens_is_respected(self, model_4bit): + out = model_4bit.generate("Write a long story about a dragon.", max_new_tokens=10, do_sample=False) + n = len(model_4bit.tokenizer(out, add_special_tokens=False)["input_ids"]) + assert n <= 12 # small slack for re-tokenization + + def test_stream_matches_non_stream(self, model_16bit): + full = model_16bit.generate(PROMPT, max_new_tokens=16, do_sample=False) + streamed = "".join(model_16bit.generate(PROMPT, max_new_tokens=16, do_sample=False, stream=True)) + assert streamed.strip() == full.strip() + + def test_chat_uses_template(self, model_16bit): + out = model_16bit.chat( + [{"role": "user", "content": "Reply with exactly one word: what color is the sky?"}], + max_new_tokens=10, + do_sample=False, + ) + assert out.strip() + assert "<|im_start|>" not in out and "assistant" not in out.lower()[:12] + + def test_stop_strings(self, model_16bit): + out = model_16bit.generate("1, 2, 3, 4,", max_new_tokens=30, do_sample=False, stop_strings=["7"]) + assert "8" not in out + + def test_batch_of_prompts_sequential_no_leak(self, model_4bit): + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + for _ in range(5): + model_4bit.generate(PROMPT, max_new_tokens=16, do_sample=False) + torch.cuda.synchronize() + assert torch.cuda.memory_allocated() - before < 64 * 1024**2 + + +# --------------------------------------------------------------------------- # +# HQQ native quantization on CUDA +# --------------------------------------------------------------------------- # + + +class TestHQQ: + @pytest.mark.parametrize("bits,max_rel_err", [(8, 0.02), (4, 0.15), (2, 0.6)]) + def test_linear_reconstruction(self, bits, max_rel_err): + from quantllm.quant.hqq import HQQLinear + + torch.manual_seed(0) + w = torch.randn(512, 1024, device="cuda", dtype=torch.float16) + q = HQQLinear(1024, 512, bits=bits, group_size=64).cuda() + q.quantize(w) + rel = (q.dequantize().float() - w.float()).norm() / w.float().norm() + assert rel < max_rel_err, f"{bits}-bit relative error {rel:.3f}" + + x = torch.randn(4, 1024, device="cuda", dtype=torch.float16) + y = q(x) + assert y.shape == (4, 512) and y.dtype == x.dtype and torch.isfinite(y).all() + + def test_quantize_real_model_still_generates(self, load_model): + from quantllm.quant.hqq import HQQQuantizer + + m = load_model(bits=16) + m.model = HQQQuantizer().quantize_model(m.model, bits=8, group_size=64) + out = m.generate(PROMPT, max_new_tokens=8, do_sample=False) + assert "paris" in out.lower() + + +# --------------------------------------------------------------------------- # +# Fine-tuning +# --------------------------------------------------------------------------- # + + +TRAIN_DATA = [ + {"instruction": "What is the secret word?", "output": "The secret word is pineapple."}, +] * 32 + + +class TestFinetune: + def test_qlora_lowers_loss(self, load_model, tmp_path): + m = load_model(bits=4) + result = m.finetune( + TRAIN_DATA, + epochs=3, + learning_rate=2e-4, + batch_size=4, + lora_r=8, + output_dir=str(tmp_path / "ft"), + logging_steps=1, + ) + assert result.get("final_loss") is not None, f"finetune failed: {result}" + losses = [h["loss"] for h in result["log_history"] if "loss" in h] + # Mean train_loss is noisy on fp16-only GPUs (early steps skipped by the loss scaler), + # so check the trajectory and that the model actually learned the answer. + assert result["final_loss"] < 0.5 * losses[0], f"loss did not fall: {losses}" + out = m.chat([{"role": "user", "content": "What is the secret word?"}], max_new_tokens=12, do_sample=False) + assert "pineapple" in out.lower(), out + + +# --------------------------------------------------------------------------- # +# Export round-trips +# --------------------------------------------------------------------------- # + + +class TestExport: + def test_safetensors_roundtrip_16bit(self, model_16bit, tmp_path): + from transformers import AutoModelForCausalLM + + from quantllm.utils.hf_compat import dtype_kwargs + + ref = model_16bit.generate(PROMPT, max_new_tokens=8, do_sample=False) + out_dir = model_16bit.export("safetensors", str(tmp_path / "st")) + reloaded = AutoModelForCausalLM.from_pretrained(out_dir, **dtype_kwargs(torch.float16)).cuda() + tok = model_16bit.tokenizer + ids = tok(PROMPT, return_tensors="pt").to("cuda") + gen = reloaded.generate(**ids, max_new_tokens=8, do_sample=False) + text = tok.decode(gen[0, ids["input_ids"].shape[1] :], skip_special_tokens=True) + assert text.strip().split()[:1] == ref.strip().split()[:1] + del reloaded + free_cuda() + + def test_safetensors_from_4bit_is_loadable(self, load_model, tmp_path): + from transformers import AutoModelForCausalLM + + m = load_model(bits=4) + out_dir = m.export("safetensors", str(tmp_path / "st4")) + reloaded = AutoModelForCausalLM.from_pretrained(out_dir, device_map="cuda") + assert reloaded is not None + del reloaded + free_cuda() + + @pytest.mark.slow + @pytest.mark.skipif(importlib.util.find_spec("gguf") is None, reason="gguf not installed") + def test_gguf_q4_k_m_loads_in_llama_cpp(self, load_model, tmp_path): + llama_cpp = pytest.importorskip("llama_cpp") + m = load_model(bits=16) + path = m.export("gguf", str(tmp_path / "model.Q4_K_M.gguf"), quantization="Q4_K_M") + llm = llama_cpp.Llama(model_path=str(path), n_ctx=256, verbose=False) + out = llm(PROMPT, max_tokens=8, temperature=0)["choices"][0]["text"] + assert out.strip() + + @pytest.mark.slow + def test_onnx_export_produces_model(self, load_model, tmp_path): + pytest.importorskip("optimum.onnxruntime", reason="optimum-onnx not installed (needs transformers<4.58)") + m = load_model(bits=16) + out = m.export("onnx", str(tmp_path / "onnx")) + assert any(p.suffix == ".onnx" for p in __import__("pathlib").Path(out).rglob("*")) + shutil.rmtree(out, ignore_errors=True) + + +# --------------------------------------------------------------------------- # +# Serving and benchmarking +# --------------------------------------------------------------------------- # + + +class TestServing: + def test_openai_chat_endpoint(self, model_id): + pytest.importorskip("fastapi") + from fastapi.testclient import TestClient + + from quantllm.server.app import _build_app + from quantllm.server.engine import InferenceEngine, ServerConfig + + engine = InferenceEngine(ServerConfig(model=model_id, bits=4, trust_remote_code=False)) + engine.load() + with TestClient(_build_app(engine)) as client: + r = client.post( + "/v1/chat/completions", + json={ + "model": model_id, + "messages": [{"role": "user", "content": "Say hi"}], + "max_tokens": 8, + "temperature": 0, + }, + ) + assert r.status_code == 200, r.text + body = r.json() + assert body["choices"][0]["message"]["content"].strip() + assert body["usage"]["completion_tokens"] > 0 + free_cuda() + + +class TestBenchmark: + def test_benchmark_reports_real_numbers(self, model_4bit): + from quantllm.benchmark.runner import BenchmarkRunner + + r = BenchmarkRunner(verbose=False).benchmark_model( + model_4bit.model, model_4bit.tokenizer, model_4bit.config, max_new_tokens=32, num_warmup=1, num_runs=2 + ) + assert r.decode_tokens_per_sec > 0 + assert r.peak_vram_gb > 0 diff --git a/tests/test_architecture_fallback.py b/tests/test_architecture_fallback.py index 55b8dfc..63fd585 100644 --- a/tests/test_architecture_fallback.py +++ b/tests/test_architecture_fallback.py @@ -49,12 +49,12 @@ def test_registered_class_fallback_is_used(monkeypatch): monkeypatch.setattr(TurboModel, "_architecture_registry", {}) monkeypatch.setattr(TurboModel, "_model_class_registry", {}) monkeypatch.setattr( - turbo_model_module.SmartConfig, + model_loader_module.SmartConfig, "detect", lambda *args, **kwargs: _make_smart_config(), ) monkeypatch.setattr( - turbo_model_module.AutoTokenizer, + model_loader_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: _make_tokenizer(), ) @@ -86,7 +86,7 @@ class _RegisteredModel: from_pretrained = classmethod(_registered_from_pretrained) monkeypatch.setattr( - turbo_model_module, + model_loader_module, "AutoModelForCausalLM", _FakeAutoModel, ) diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 6e925ea..9d3b18f 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -27,6 +27,63 @@ def test_default_values(self): assert r.decode_tokens_per_sec == 0.0 +class _FakeTok: + pad_token_id = 0 + eos_token_id = 0 + + def __call__(self, text, return_tensors=None): + import torch + + return {"input_ids": torch.ones(1, len(text.split()), dtype=torch.long)} + + +class _FakeLM: + """Fake LM on a fake clock: forward = 0.5s, generate = 0.5s prefill + 0.1s per extra token. + + It stops early (EOS) after ``new_tokens`` tokens regardless of max_new_tokens. + """ + + device = "cpu" + + def __init__(self, clock, new_tokens: int = 5): + self.clock = clock + self.new_tokens = new_tokens + self.forward_lens = [] + + def num_parameters(self): + return 1000 + + def __call__(self, input_ids=None, **kwargs): + self.forward_lens.append(input_ids.shape[1]) + self.clock[0] += 0.5 + + def generate(self, input_ids, max_new_tokens=16, **kwargs): + import torch + + n = min(self.new_tokens, max_new_tokens) + self.clock[0] += 0.5 + 0.1 * (n - 1) + return torch.cat([input_ids, torch.ones(1, n, dtype=torch.long)], dim=1) + + +class TestBenchmarkMeasurement: + def test_prefill_long_prompt_and_decode_actual_tokens(self, monkeypatch): + import quantllm.benchmark.runner as runner_mod + from quantllm.benchmark import BenchmarkRunner + + clock = [0.0] + monkeypatch.setattr(runner_mod.time, "perf_counter", lambda: clock[0]) + + model = _FakeLM(clock, new_tokens=5) + r = BenchmarkRunner(verbose=False).benchmark_model( + model, _FakeTok(), None, prompt="one two three", max_new_tokens=128, num_warmup=1, num_runs=2 + ) + # Prefill is a plain forward over a ~512-token prompt, not the 3-token benchmark prompt. + assert max(model.forward_lens) == 512 + assert r.prefill_tokens_per_sec == pytest.approx(512 / 0.5) + # 5 tokens actually generated (not 128); prefill (0.5s) excluded: 4 decode steps in 0.4s. + assert r.decode_tokens_per_sec == pytest.approx(10.0) + + class TestBenchmarkCLI: def test_bench_imports(self): from quantllm.benchmark import run_benchmark_cli, format_report_table diff --git a/tests/test_cli.py b/tests/test_cli.py index a99413f..9ed90ed 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -39,3 +39,105 @@ def test_info_command(self): finally: sys.argv = old_argv sys.stdout = old_stdout + + +def _run_cli(argv, monkeypatch, capsys=None): + import sys + + monkeypatch.setattr(sys, "argv", ["quantllm", *argv]) + with pytest.raises(SystemExit) as exc: + main() + return exc.value.code + + +class TestCLIExitCodes: + def test_version_matches_package(self, monkeypatch, capsys): + import quantllm + + _run_cli(["version"], monkeypatch) + assert quantllm.__version__ in capsys.readouterr().out + + def test_models_without_subcommand_exits_2(self, monkeypatch, capsys): + assert _run_cli(["models"], monkeypatch) == 2 + + def test_register_failure_exits_1(self, monkeypatch, capsys): + import quantllm.registry as registry_mod + + monkeypatch.setattr(registry_mod, "register_model", lambda *a, **k: False) + assert _run_cli(["register", "me/model"], monkeypatch) == 1 + + def test_register_success_exits_0(self, monkeypatch, capsys): + import sys + + import quantllm.registry as registry_mod + + monkeypatch.setattr(registry_mod, "register_model", lambda *a, **k: True) + monkeypatch.setattr(sys, "argv", ["quantllm", "register", "me/model"]) + main() # no SystemExit on success + assert "registered" in capsys.readouterr().out + + +class TestCLIFlagWiring: + def _patch_loader(self, monkeypatch): + from unittest import mock + + import quantllm.core as core + + fake_model = mock.MagicMock() + loader = mock.MagicMock(return_value=fake_model) + monkeypatch.setattr(core.TurboModel, "from_pretrained", loader) + return loader, fake_model + + @pytest.mark.parametrize( + "argv", + [ + ["convert", "org/m", "-o", "out.gguf"], + ["export", "org/m", "-f", "gguf"], + ["finetune", "org/m", "--dataset", "d"], + ], + ) + def test_trust_remote_code_defaults_false(self, argv, monkeypatch, capsys): + import sys + + loader, _ = self._patch_loader(monkeypatch) + monkeypatch.setattr(sys, "argv", ["quantllm", *argv]) + main() + assert loader.call_args.kwargs["trust_remote_code"] is False + + @pytest.mark.parametrize( + "argv", + [ + ["convert", "org/m", "-o", "out.gguf", "--trust-remote-code"], + ["export", "org/m", "--trust-remote-code"], + ["finetune", "org/m", "--dataset", "d", "--trust-remote-code"], + ], + ) + def test_trust_remote_code_passed_through(self, argv, monkeypatch, capsys): + import sys + + loader, _ = self._patch_loader(monkeypatch) + monkeypatch.setattr(sys, "argv", ["quantllm", *argv]) + main() + assert loader.call_args.kwargs["trust_remote_code"] is True + + def test_convert_wires_output_and_quant(self, monkeypatch, capsys): + import sys + + _, fake_model = self._patch_loader(monkeypatch) + monkeypatch.setattr(sys, "argv", ["quantllm", "convert", "org/m", "-o", "x.gguf", "-q", "q5_k_m"]) + main() + fake_model.export.assert_called_once_with("gguf", output_path="x.gguf", quantization="q5_k_m") + + @pytest.mark.parametrize("flag,expected", [([], False), (["--trust-remote-code"], True)]) + def test_serve_trust_remote_code(self, flag, expected, monkeypatch): + import sys + + pytest.importorskip("fastapi") + import quantllm.server.app as app_mod + + captured = {} + monkeypatch.setattr(app_mod, "run_server", lambda cfg: captured.setdefault("cfg", cfg)) + monkeypatch.setattr(sys, "argv", ["quantllm", "serve", "org/m", "--port", "9999", *flag]) + main() + assert captured["cfg"].trust_remote_code is expected + assert captured["cfg"].port == 9999 diff --git a/tests/test_export_fixes.py b/tests/test_export_fixes.py new file mode 100644 index 0000000..4e3ba4a --- /dev/null +++ b/tests/test_export_fixes.py @@ -0,0 +1,317 @@ +"""CPU-only regression tests for export / quantization fixes.""" + +import os +import subprocess +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +import torch.nn as nn + +from quantllm.core.export_router import ExportRouterMixin +from quantllm.kernels.cpu import CPUInt4Linear, CPUInt8Linear +from quantllm.kernels.triton import is_triton_available +from quantllm.kernels.triton.quantized_linear import _pytorch_dequant_matmul +from quantllm.quant import llama_cpp +from quantllm.quant.hqq import HQQLinear + + +def _outlier_matrix(out_features=256, in_features=512, seed=0): + g = torch.Generator().manual_seed(seed) + w = torch.randn(out_features, in_features, generator=g) + mask = torch.rand(out_features, in_features, generator=g) < 0.01 + w[mask] *= 10.0 + return w + + +# --------------------------------------------------------------------------- HQQ + + +class TestHQQFixes: + def test_non_divisible_in_features(self): + layer = HQQLinear(576, 64, bits=4, group_size=128) + w = torch.randn(64, 576) + layer.quantize(w) + deq = layer.dequantize(torch.float32) + assert deq.shape == (64, 576) + assert ((deq - w) ** 2).mean() < 0.05 + out = layer(torch.randn(3, 576)) + assert out.shape == (3, 64) + + def test_group_larger_than_in_features(self): + layer = HQQLinear(48, 32, bits=4, group_size=128) + w = torch.randn(32, 48) + layer.quantize(w) + assert layer.dequantize().shape == (32, 48) + + @pytest.mark.parametrize("bits", [2, 4]) + def test_optimize_lowers_reconstruction_error(self, bits): + w = _outlier_matrix() + errors = {} + for optimize in (False, True): + layer = HQQLinear(512, 256, bits=bits, group_size=64, optimize=optimize) + layer.quantize(w) + deq = layer.dequantize(torch.float32) + errors[optimize] = ((deq - w) ** 2).mean().item() + assert errors[True] < errors[False] + + def test_bf16_forward_dtype(self): + layer = HQQLinear(64, 32, bits=4, group_size=32, bias=True) + layer.quantize(torch.randn(32, 64)) + layer.bias = nn.Parameter(torch.randn(32, dtype=torch.float16)) + out = layer(torch.randn(2, 64, dtype=torch.bfloat16)) + assert out.dtype == torch.bfloat16 + + +# --------------------------------------------------------------------------- CPU / Triton kernels + + +class TestKernelFixes: + def test_cpu_int4_group_ge_in_features(self): + layer = CPUInt4Linear(48, 32, bits=4, group_size=128) + assert layer.qweight.shape == (32, 1, 64) + w = torch.randn(32, 48) + layer.quantize(w) + deq = layer._dequantize_weights() + assert deq.shape == (32, 48) + assert ((deq - w) ** 2).mean() < 0.05 + + def test_cpu_int4_non_divisible(self): + layer = CPUInt4Linear(576, 16, bits=4, group_size=128) + layer.quantize(torch.randn(16, 576)) + assert layer._dequantize_weights().shape == (16, 576) + + def test_cpu_int4_bf16_forward(self): + layer = CPUInt4Linear(64, 32, bits=4, group_size=32, bias=True) + layer.quantize(torch.randn(32, 64)) + layer.bias = nn.Parameter(torch.randn(32)) + out = layer(torch.randn(2, 64, dtype=torch.bfloat16)) + assert out.dtype == torch.bfloat16 + + def test_cpu_int8_bf16_forward(self): + layer = CPUInt8Linear(64, 32, bias=True) + layer.quantize(torch.randn(32, 64)) + layer.bias = nn.Parameter(torch.randn(32)) + out = layer(torch.randn(2, 64, dtype=torch.bfloat16)) + assert out.dtype == torch.bfloat16 + + def test_pytorch_fallback_casts_to_input_dtype(self): + qweight = torch.randint(-8, 7, (64, 16), dtype=torch.int8) + scales = torch.rand(2, 16) + zeros = torch.zeros(2, 16) + bias = torch.randn(16) + out = _pytorch_dequant_matmul(torch.randn(3, 64, dtype=torch.bfloat16), qweight, scales, zeros, bias, 32) + assert out.dtype == torch.bfloat16 + + def test_triton_requires_cuda(self): + with mock.patch("torch.cuda.is_available", return_value=False): + assert is_triton_available() is False + + +# --------------------------------------------------------------------------- Export routing + + +class TestMLXQuantParsing: + @pytest.mark.parametrize("name", [None, "", "f16", "bf16", "fp16", "f32", "none", "BF16"]) + def test_full_precision_means_no_quantization(self, name): + assert ExportRouterMixin._mlx_quant_bits(name) is None + + @pytest.mark.parametrize( + "name,bits", + [("q4_k_m", 4), ("4bit", 4), ("q8_0", 8), ("8bit", 8), ("q2_k", 4), ("q3_k_m", 4), ("q5_k_m", 8), ("q6_k", 8)], + ) + def test_bit_mapping(self, name, bits): + assert ExportRouterMixin._mlx_quant_bits(name) == bits + + +class TestGGUFOuttype: + @pytest.mark.parametrize( + "name,expected", + [ + ("f16", "f16"), + ("F16", "f16"), + ("float16", "f16"), + ("bf16", "bf16"), + ("bfloat16", "bf16"), + ("torch.bfloat16", "bf16"), + ("f32", "f32"), + ("float32", "f32"), + ("q4_k_m", None), + ("Q8_0", None), + ], + ) + def test_mapping(self, name, expected): + assert llama_cpp.gguf_outtype(name) == expected + + +def _make_converter_env(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "config.json").write_text("{}") + converter = tmp_path / "convert_hf_to_gguf.py" + converter.write_text("") + return model_dir, converter + + +def _fake_popen(lines, returncode=0, captured=None): + def factory(cmd, **kwargs): + if captured is not None: + captured["cmd"] = cmd + captured["kwargs"] = kwargs + proc = mock.MagicMock() + text = "".join(lines) + proc.stdout = mock.MagicMock() + proc.stdout.readline.side_effect = list(lines) + [""] + proc.communicate.return_value = (text, None) + proc.returncode = returncode + proc.wait.return_value = returncode + return proc + + return factory + + +class TestConvertToGGUF: + @pytest.mark.parametrize("print_output", [False, True]) + def test_merged_stream_and_output_dir(self, tmp_path, print_output): + model_dir, converter = _make_converter_env(tmp_path) + out_dir = tmp_path / "out" + out_dir.mkdir() + captured = {} + with mock.patch.object(subprocess, "Popen", side_effect=_fake_popen(["INFO: Loading model\n"], 0, captured)): + files, _ = llama_cpp.convert_to_gguf( + model_name="m", + input_folder=str(model_dir), + model_dtype="bfloat16", + quantization_type="bf16", + converter_location=str(converter), + print_output=print_output, + output_dir=str(out_dir), + ) + assert captured["kwargs"]["stderr"] is subprocess.STDOUT + cmd = captured["cmd"] + assert cmd[cmd.index("--outtype") + 1] == "bf16" + outfile = cmd[cmd.index("--outfile") + 1] + assert os.path.isabs(outfile) + assert os.path.dirname(outfile) == os.path.abspath(str(out_dir)) + assert files == [outfile] + + def test_failure_reports_merged_output(self, tmp_path): + model_dir, converter = _make_converter_env(tmp_path) + with mock.patch.object(subprocess, "Popen", side_effect=_fake_popen(["Traceback: boom from stderr\n"], 1)): + with pytest.raises(RuntimeError, match="boom from stderr"): + llama_cpp.convert_to_gguf( + model_name="m", + input_folder=str(model_dir), + model_dtype="f16", + quantization_type="f16", + converter_location=str(converter), + print_output=False, + ) + + +class TestQuantizeGGUFImatrix: + @pytest.mark.parametrize("qtype", ["iq3_xxs", "iq2_xxs", "iq1_s"]) + def test_requires_imatrix(self, tmp_path, qtype): + with pytest.raises(ValueError, match="importance matrix"): + llama_cpp.quantize_gguf(str(tmp_path / "in.gguf"), str(tmp_path / "out.gguf"), qtype) + + def test_imatrix_flag_passed(self, tmp_path): + inp = tmp_path / "in.gguf" + inp.write_bytes(b"x" * 16) + imatrix = tmp_path / "imatrix.dat" + imatrix.write_bytes(b"x") + captured = {} + + def fake_popen(cmd, **kwargs): + captured["cmd"] = cmd + proc = mock.MagicMock() + proc.stdout = iter([]) + proc.returncode = 0 + return proc + + with mock.patch.object(subprocess, "Popen", side_effect=fake_popen): + llama_cpp.quantize_gguf( + str(inp), + str(tmp_path / "out.gguf"), + "iq2_xxs", + quantizer_location="llama-quantize", + print_output=False, + imatrix=str(imatrix), + ) + cmd = captured["cmd"] + assert cmd[1:3] == ["--imatrix", str(imatrix)] + assert "IQ2_XXS" in cmd + + +class TestWindowsBinaryResolution: + def _setup(self, tmp_path, rel_parts): + llama_dir = tmp_path / "llama.cpp" + exe = llama_dir.joinpath(*rel_parts) + exe.parent.mkdir(parents=True, exist_ok=True) + exe.write_text("") + (llama_dir / "convert_hf_to_gguf.py").write_text("") + return llama_dir, exe + + def test_exe_name(self): + with mock.patch.object(llama_cpp.os, "name", "nt"): + assert llama_cpp._exe_name("llama-quantize") == "llama-quantize.exe" + with mock.patch.object(llama_cpp.os, "name", "posix"): + assert llama_cpp._exe_name("llama-quantize") == "llama-quantize" + + @pytest.mark.parametrize( + "rel_parts", + [ + ("build", "bin", "Release", "llama-quantize.exe"), + ("build", "bin", "llama-quantize.exe"), + ("llama-quantize.exe",), + ], + ) + def test_finds_msvc_and_root_exe(self, tmp_path, rel_parts): + llama_dir, exe = self._setup(tmp_path, rel_parts) + with ( + mock.patch.object(llama_cpp, "get_llama_cpp_dir", return_value=str(llama_dir)), + mock.patch.object(llama_cpp.os, "name", "nt"), + mock.patch.object(llama_cpp.os, "access", return_value=True), + ): + quantizer, converter = llama_cpp.check_llama_cpp() + assert os.path.normpath(quantizer) == os.path.normpath(str(exe)) + assert converter.endswith("convert_hf_to_gguf.py") + + def test_cmake_failure_shows_output_tail(self, tmp_path): + result = SimpleNamespace(returncode=2, stdout="line1\nCMake Error: no compiler found\n") + with mock.patch.object(subprocess, "run", return_value=result): + with pytest.raises(RuntimeError, match="no compiler found"): + llama_cpp._run_build_step(["cmake", "."], print_output=False, step="CMake configure") + + +class TestStaging: + def test_stage_in_memory_model_dequantizes_bnb(self, tmp_path): + saved = {} + + class FakeModel: + def save_pretrained(self, path, **kwargs): + saved["path"] = path + open(os.path.join(path, "config.json"), "w").close() + + dequantized = FakeModel() + router = ExportRouterMixin.__new__(ExportRouterMixin) + router.model = object() + router.tokenizer = None + router._is_bnb_quantized = lambda: True # type: ignore[assignment] + router._dequantize_model = lambda: dequantized # type: ignore[assignment] + + staging = router._stage_in_memory_model() + try: + assert saved["path"] == staging + assert os.path.exists(os.path.join(staging, "config.json")) + finally: + import shutil + + shutil.rmtree(staging, ignore_errors=True) + + def test_has_in_memory_changes(self): + router = ExportRouterMixin.__new__(ExportRouterMixin) + assert router._has_in_memory_changes() is False + router._lora_merged = True # type: ignore[attr-defined] + assert router._has_in_memory_changes() is True diff --git a/tests/test_finetune_fixes.py b/tests/test_finetune_fixes.py new file mode 100644 index 0000000..d08d432 --- /dev/null +++ b/tests/test_finetune_fixes.py @@ -0,0 +1,276 @@ +"""Regression tests for fine-tuning / training-data fixes (CPU only, tiny mocks).""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn + +from quantllm.core.training import ( + _find_text_column, + format_training_texts, + load_training_data, +) +from quantllm.core.turbo_model import TurboModel + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +class _FakeDataset: + """Minimal stand-in for ``datasets.Dataset`` (column_names, [0], map).""" + + def __init__(self, rows): + self.rows = rows + self.column_names = list(rows[0].keys()) + + def __getitem__(self, idx): + return self.rows[idx] + + def map(self, fn, batched=True, remove_columns=None, **kwargs): + batch = {k: [r[k] for r in self.rows] for k in self.column_names} + return fn(batch) + + +class _CharTokenizer: + eos_token = "" + pad_token = "" + eos_token_id = 2 + pad_token_id = 0 + chat_template = None + + def __init__(self): + self.texts = None + + def __call__(self, texts, truncation=True, max_length=8, padding=False, add_special_tokens=True, **kw): + self.texts = list(texts) + input_ids, attention_mask = [], [] + for t in texts: + ids = [10 + (ord(c) % 50) for c in t][:max_length] + mask = [1] * len(ids) + if padding == "max_length": + pad = max_length - len(ids) + ids = ids + [self.pad_token_id] * pad + mask = mask + [0] * pad + input_ids.append(ids) + attention_mask.append(mask) + return {"input_ids": input_ids, "attention_mask": attention_mask} + + +ALPACA_ROWS = [ + {"instruction": "Translate", "input": "bonjour", "output": "hello"}, + {"instruction": "Say hi", "input": "", "output": "hi"}, +] + + +# --------------------------------------------------------------------------- +# training.py: column detection + label masking +# --------------------------------------------------------------------------- + + +def test_alpaca_prefers_instruction_output_over_input(): + ds = _FakeDataset(ALPACA_ROWS) + assert _find_text_column(ds) == "__instruction_output__" + + +def test_messages_column_detected_before_text(): + ds = _FakeDataset([{"messages": [{"role": "user", "content": "x"}], "text": "x"}]) + assert _find_text_column(ds) == "__messages__:messages" + + +def test_load_training_data_includes_input_and_masks_padding(): + tok = _CharTokenizer() + result = load_training_data(_FakeDataset(ALPACA_ROWS), tok, max_length=200) + + assert "### Input:\nbonjour" in tok.texts[0] + assert "### Input:" not in tok.texts[1] # empty input is omitted + + for ids, mask, labels in zip(result["input_ids"], result["attention_mask"], result["labels"]): + for i, m, lab in zip(ids, mask, labels): + assert lab == (i if m else -100) + assert -100 in result["labels"][1] + # labels must not alias input_ids + assert result["labels"][1] is not result["input_ids"][1] + assert tok.pad_token_id in result["input_ids"][1] + + +def test_format_training_texts_appends_eos_and_uses_chat_template(): + tok = _CharTokenizer() + texts, templated = format_training_texts({k: [r[k] for r in ALPACA_ROWS] for k in ALPACA_ROWS[0]}, tok) + assert templated is False + assert all(t.endswith("") for t in texts) + assert "bonjour" in texts[0] + + tok.chat_template = "{{ bos_token }}..." + tok.apply_chat_template = Mock(return_value="Translate\n\nbonjourhello") + texts, templated = format_training_texts({k: [r[k] for r in ALPACA_ROWS] for k in ALPACA_ROWS[0]}, tok) + assert templated is True + assert texts[0].count("") == 1 # not doubled + msgs = tok.apply_chat_template.call_args_list[0][0][0] + assert msgs[0] == {"role": "user", "content": "Translate\n\nbonjour"} + assert msgs[1] == {"role": "assistant", "content": "hello"} + + +def test_format_training_texts_sharegpt_conversations(): + tok = _CharTokenizer() + batch = {"conversations": [[{"from": "human", "value": "Q"}, {"from": "gpt", "value": "A"}]]} + texts, templated = format_training_texts(batch, tok) + assert templated is False + assert texts == ["User: Q\nAssistant: A\n"] + + +# --------------------------------------------------------------------------- +# finetune_mixin.py +# --------------------------------------------------------------------------- + + +def _instance(monkeypatch, model, tok=None, quantized=False): + monkeypatch.setattr(TurboModel, "_architecture_registry", {}) + monkeypatch.setattr(TurboModel, "_model_class_registry", {}) + inst = TurboModel.__new__(TurboModel) + inst.model = model + inst.tokenizer = tok or _CharTokenizer() + inst.config = SimpleNamespace( + bits=4, + effective_loading_bits=16, + dtype=torch.float32, + gradient_checkpointing=False, + compile_model=False, + batch_size=1, + gradient_accumulation_steps=1, + max_seq_length=32, + ) + inst._is_quantized_override = quantized + inst._is_finetuned = False + inst._lora_applied = False + inst.verbose = False + return inst + + +class _Proj(nn.Module): + def __init__(self): + super().__init__() + for name in ("qkv_proj", "o_proj", "gate_up_proj", "down_proj"): + setattr(self, name, nn.Linear(2, 2)) + + +class _Phi3Like(nn.Module): + def __init__(self, model_type="phi3"): + super().__init__() + self.config = SimpleNamespace(model_type=model_type, _name_or_path="org/tiny", use_cache=False) + self.layer = _Proj() + + +def test_target_modules_phi3(monkeypatch): + inst = _instance(monkeypatch, _Phi3Like("phi3")) + assert inst._get_lora_target_modules() == ["qkv_proj", "o_proj", "gate_up_proj", "down_proj"] + + +def test_target_modules_unknown_arch_uses_all_linear(monkeypatch): + inst = _instance(monkeypatch, SimpleNamespace(config=SimpleNamespace(model_type="brand_new_arch"))) + assert inst._get_lora_target_modules() == "all-linear" + + +def test_target_modules_known_type_but_missing_modules_uses_all_linear(monkeypatch): + # falcon table but the model actually has phi3-style names + inst = _instance(monkeypatch, _Phi3Like("falcon")) + assert inst._get_lora_target_modules() == "all-linear" + + +def test_optimizer_choice_follows_bnb_quantization(monkeypatch): + bnb = SimpleNamespace(config=SimpleNamespace(quantization_config=SimpleNamespace(quant_method="bitsandbytes"))) + inst = _instance(monkeypatch, bnb) + assert inst._is_bnb_quantized() is True + hqq = SimpleNamespace(config=SimpleNamespace(quantization_config=SimpleNamespace(quant_method="hqq"))) + inst.model = hqq + assert inst._is_bnb_quantized() is False + + +def test_pad_token_made_distinct_from_eos(monkeypatch): + tok = SimpleNamespace(eos_token="", pad_token="", unk_token="") + inst = _instance(monkeypatch, SimpleNamespace(), tok=tok) + inst._ensure_distinct_pad_token() + assert tok.pad_token == "" + + +def test_pad_token_added_when_no_unk(monkeypatch): + tok = Mock(eos_token="", pad_token=None, unk_token=None) + tok.__len__ = Mock(return_value=11) + model = SimpleNamespace(resize_token_embeddings=Mock()) + inst = _instance(monkeypatch, model, tok=tok) + inst._ensure_distinct_pad_token() + tok.add_special_tokens.assert_called_once_with({"pad_token": "<|pad|>"}) + model.resize_token_embeddings.assert_called_once_with(11) + + +class _LoraModel(nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(model_type="phi3", _name_or_path="org/tiny", use_cache=False) + self.base = nn.Linear(2, 2) + self.lora_A = nn.Linear(2, 1, bias=False) + self.is_gradient_checkpointing = False + + def num_parameters(self): + return sum(p.numel() for p in self.parameters()) + + +def test_second_finetune_keeps_adapters_trainable_and_saves(monkeypatch): + pytest.importorskip("peft") + import peft + import transformers + + prepare = Mock(side_effect=lambda m, **kw: m) + monkeypatch.setattr(peft, "prepare_model_for_kbit_training", prepare) + monkeypatch.setattr(peft, "get_peft_model", Mock(side_effect=lambda m, cfg: m)) + monkeypatch.setattr(peft, "LoraConfig", Mock()) + + captured = {} + + class FakeArgs: + def __init__(self, **kwargs): + captured["args"] = kwargs + + class FakeTrainer: + def __init__(self, model=None, args=None, train_dataset=None, data_collator=None, processing_class=None): + captured["trainer_kwargs"] = {"processing_class": processing_class} + self.state = SimpleNamespace(log_history=[{"loss": 1.0}]) + self.save_model = Mock() + captured["trainer"] = self + + def train(self): + return SimpleNamespace(training_loss=0.5) + + monkeypatch.setattr(transformers, "TrainingArguments", FakeArgs) + monkeypatch.setattr(transformers, "Trainer", FakeTrainer) + monkeypatch.setattr(transformers, "DataCollatorForLanguageModeling", Mock()) + + model = _LoraModel() + tok = _CharTokenizer() + tok.save_pretrained = Mock() + inst = _instance(monkeypatch, model, tok=tok, quantized=True) + inst._prepare_dataset = Mock(return_value=[]) + + out1 = inst.finetune([], output_dir="out_dir", epochs=1) + assert prepare.call_count == 1 + assert out1["log_history"] == [{"loss": 1.0}] + captured["trainer"].save_model.assert_called_once_with("out_dir") + tok.save_pretrained.assert_called_once_with("out_dir") + assert captured["trainer_kwargs"]["processing_class"] is tok + # not bnb-quantized -> no paged optimizer + assert captured["args"]["optim"] == "adamw_torch" + # back in inference mode + assert model.training is False + assert model.config.use_cache is True + + # Simulate prepare_model_for_kbit_training freezing everything + for p in model.parameters(): + p.requires_grad_(False) + + inst.finetune([], output_dir="out_dir", epochs=1) + assert prepare.call_count == 1 # not re-run on the PeftModel + assert model.lora_A.weight.requires_grad is True + assert model.base.weight.requires_grad is False diff --git a/tests/test_generation.py b/tests/test_generation.py index 2ad8005..e22089a 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -234,3 +234,180 @@ def test_generate_handles_stop_strings(monkeypatch): result = instance.generate("Hello", stop_strings=["END"]) assert "END" not in result + + +# --------------------------------------------------------------------------- +# Regression tests for generation fixes +# --------------------------------------------------------------------------- + + +def _instance(monkeypatch, model=None, tok=None, config=None): + monkeypatch.setattr(TurboModel, "_architecture_registry", {}) + monkeypatch.setattr(TurboModel, "_model_class_registry", {}) + instance = TurboModel.__new__(TurboModel) + instance.model = model or _fake_model() + instance.tokenizer = tok or _tokenizer() + instance.config = config or _smart_config() + instance._is_quantized_override = False + instance._is_finetuned = False + instance._lora_applied = False + instance.verbose = False + return instance + + +class _BosTokenizer: + """Tokenizer that maps the literal "" to BOS and prepends BOS when asked.""" + + bos_token = "" + eos_token = "" + bos_token_id = 1 + eos_token_id = 2 + pad_token_id = 0 + chat_template = "{{ bos_token }}..." + truncation_side = "right" + + def __init__(self): + self.calls = [] + + def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=False): + assert tokenize is False + return "[INST] " + messages[-1]["content"] + " [/INST]" + + def __call__(self, text, return_tensors="pt", add_special_tokens=True, **kwargs): + self.calls.append({"add_special_tokens": add_special_tokens, "truncation_side": self.truncation_side, **kwargs}) + ids = [] + if add_special_tokens: + ids.append(self.bos_token_id) + rest = text + if rest.startswith(""): + ids.append(self.bos_token_id) + rest = rest[3:] + ids.extend(10 + (ord(c) % 50) for c in rest) + input_ids = torch.tensor([ids], dtype=torch.long) + return {"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids)} + + def decode(self, token_ids, skip_special_tokens=False): + return "ok" + + +def test_chat_does_not_double_bos(monkeypatch): + tok = _BosTokenizer() + instance = _instance(monkeypatch, tok=tok) + instance.chat([{"role": "user", "content": "Hi"}]) + input_ids = instance.model.generate.call_args[1]["input_ids"][0].tolist() + assert input_ids.count(tok.bos_token_id) == 1 + assert tok.calls[-1]["add_special_tokens"] is False + + +def test_generate_does_not_override_eos_token_id(monkeypatch): + instance = _instance(monkeypatch) + instance.generate("Hello") + call_kwargs = instance.model.generate.call_args[1] + assert "eos_token_id" not in call_kwargs + # generation_config has no pad -> tokenizer pad is filled in + assert call_kwargs["pad_token_id"] == 0 + assert "attention_mask" in call_kwargs + + +def test_generate_keeps_generation_config_pad(monkeypatch): + model = _fake_model() + model.generation_config = SimpleNamespace(pad_token_id=5, eos_token_id=[2, 107]) + instance = _instance(monkeypatch, model=model) + instance.generate("Hello") + call_kwargs = instance.model.generate.call_args[1] + assert "pad_token_id" not in call_kwargs + assert "eos_token_id" not in call_kwargs + + +def test_generate_clamps_max_length_and_truncates_left(monkeypatch): + tok = _BosTokenizer() + config = _smart_config() + config.max_seq_length = 10 + instance = _instance(monkeypatch, tok=tok, config=config) + instance.generate("Hello", max_new_tokens=50) + call = tok.calls[-1] + assert call["max_length"] == 1 + assert call["truncation"] is True + assert call["truncation_side"] == "left" + assert tok.truncation_side == "right" # restored + + +def test_generate_uses_input_embedding_device_and_eval(monkeypatch): + model = _fake_model() + model.eval = Mock() + model.get_input_embeddings = Mock(return_value=SimpleNamespace(weight=torch.zeros(1, device="cpu"))) + instance = _instance(monkeypatch, model=model) + instance.generate("Hello") + model.eval.assert_called_once() + model.get_input_embeddings.assert_called() + + +def test_generate_trims_at_earliest_stop_string(monkeypatch): + tok = _tokenizer() + tok.decode = lambda ids, skip_special_tokens=False: "answer B-STOP more A-STOP tail" + instance = _instance(monkeypatch, tok=tok) + result = instance.generate("Hello", stop_strings=["A-STOP", "B-STOP"]) + assert result == "answer" + call_kwargs = instance.model.generate.call_args[1] + assert call_kwargs["stop_strings"] == ["A-STOP", "B-STOP"] + assert call_kwargs["tokenizer"] is tok + + +class _StreamTokenizer: + eos_token_id = 2 + pad_token_id = 0 + VOCAB = {5: "Hello ", 6: "world ", 7: "ST", 8: "OP ", 9: "leak "} + + def __call__(self, text, return_tensors="pt", **kwargs): + input_ids = torch.tensor([[1, 3]], dtype=torch.long) + return {"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids)} + + def decode(self, token_ids, skip_special_tokens=False, **kwargs): + if isinstance(token_ids, torch.Tensor): + token_ids = token_ids.tolist() + return "".join(self.VOCAB.get(int(t), "") for t in token_ids) + + +def test_streaming_does_not_leak_stop_string(monkeypatch): + model = _fake_model() + + def streaming_generate(**kwargs): + streamer = kwargs["streamer"] + criteria = kwargs["stopping_criteria"] + ids = kwargs["input_ids"] + streamer.put(ids) # prompt (skipped) + for t in [5, 6, 7, 8, 9, 9]: + ids = torch.cat([ids, torch.tensor([[t]])], dim=1) + streamer.put(torch.tensor([t])) + if bool(criteria(ids, None).all()): + break + streamer.end() + + model.generate = Mock(side_effect=streaming_generate) + instance = _instance(monkeypatch, model=model, tok=_StreamTokenizer()) + chunks = list(instance.generate("Hi", stream=True, stop_strings=["STOP"])) + text = "".join(chunks) + assert text == "Hello world " + + +def test_streaming_propagates_thread_exception(monkeypatch): + import pytest + + model = _fake_model() + model.generate = Mock(side_effect=RuntimeError("CUDA exploded")) + instance = _instance(monkeypatch, model=model, tok=_StreamTokenizer()) + gen = instance.generate("Hi", stream=True, stream_timeout=5) + with pytest.raises(RuntimeError, match="CUDA exploded"): + list(gen) + + +def test_enable_flash_attention_noop_without_setter(): + from quantllm.core.generation import GenerationMixin + + cfg = SimpleNamespace(_attn_implementation="eager") + GenerationMixin._enable_flash_attention(SimpleNamespace(config=cfg), verbose=False) + assert cfg._attn_implementation == "eager" + + model = SimpleNamespace(set_attn_implementation=Mock(side_effect=ValueError("nope"))) + GenerationMixin._enable_flash_attention(model, verbose=False) # must not raise + model.set_attn_implementation.assert_called_once() diff --git a/tests/test_hub_push.py b/tests/test_hub_push.py new file mode 100644 index 0000000..4601890 --- /dev/null +++ b/tests/test_hub_push.py @@ -0,0 +1,159 @@ +"""Tests for HF Hub push: repo privacy, base_model sanitization, staging dirs.""" + +import os +from types import SimpleNamespace +from unittest import mock + +import pytest + +import quantllm.hub.hub_manager as hub_manager_mod +from quantllm.hub.hub_manager import QuantLLMHubManager +from quantllm.hub.model_card import ModelCardGenerator, sanitize_base_model + + +@pytest.fixture +def fake_api(monkeypatch): + api = mock.MagicMock() + api.repo_exists.return_value = False + monkeypatch.setattr(hub_manager_mod, "HfApi", mock.MagicMock(return_value=api)) + return api + + +class TestPrivateForwarded: + def test_push_private_true(self, fake_api): + manager = QuantLLMHubManager("me/repo") + manager.push(private=True) + fake_api.create_repo.assert_called_once() + kwargs = fake_api.create_repo.call_args.kwargs + assert kwargs["private"] is True + assert kwargs["exist_ok"] is True + + def test_push_uses_init_default(self, fake_api): + manager = QuantLLMHubManager("me/repo", private=True) + manager.push() + assert fake_api.create_repo.call_args.kwargs["private"] is True + + def test_push_default_public(self, fake_api): + QuantLLMHubManager("me/repo").push() + assert fake_api.create_repo.call_args.kwargs["private"] is False + + +class TestStagingDir: + def test_unique_per_manager_and_keeps_owner(self, fake_api): + a = QuantLLMHubManager("alice/model") + b = QuantLLMHubManager("bob/model") + try: + assert a.staging_dir != b.staging_dir + assert "alice" in os.path.basename(a.staging_dir) + assert "bob" in os.path.basename(b.staging_dir) + assert os.path.isdir(a.staging_dir) and os.listdir(a.staging_dir) == [] + finally: + a.cleanup() + b.cleanup() + + def test_cleaned_after_successful_push(self, fake_api): + manager = QuantLLMHubManager("me/repo") + with open(os.path.join(manager.staging_dir, "stale.bin"), "wb") as f: + f.write(b"x") + manager.push() + assert not os.path.exists(manager.staging_dir) + + def test_kept_after_failed_push(self, fake_api): + fake_api.upload_folder.side_effect = RuntimeError("boom") + manager = QuantLLMHubManager("me/repo") + try: + with pytest.raises(RuntimeError): + manager.push() + assert os.path.isdir(manager.staging_dir) + finally: + manager.cleanup() + assert not os.path.exists(manager.staging_dir) + + +class TestBaseModelSanitization: + @pytest.mark.parametrize("value", ["meta-llama/Llama-3.2-1B", "Qwen/Qwen2.5-0.5B", "org/model.v1_x"]) + def test_accepts_repo_ids(self, value): + assert sanitize_base_model(value) == value + + @pytest.mark.parametrize( + "value", + [ + "/home/user/models/llama", + "C:\\Users\\me\\models\\llama", + "./checkpoints/run1/final", + "model-without-owner", + "", + None, + ], + ) + def test_rejects_paths(self, value): + assert sanitize_base_model(value) is None + + def test_existing_local_dir_rejected(self, tmp_path, monkeypatch): + (tmp_path / "org" / "model").mkdir(parents=True) + monkeypatch.chdir(tmp_path) + assert sanitize_base_model("org/model") is None + + def test_card_omits_local_path(self): + card = ModelCardGenerator( + repo_id="me/repo", base_model="/home/user/secret/llama-ft", format="safetensors" + ).generate() + header = card.split("---", 2)[1] + assert "base_model:" not in header + assert "/home/user" not in card + assert "llama-ft" in card + + def test_card_keeps_hub_id(self): + card = ModelCardGenerator(repo_id="me/repo", base_model="Qwen/Qwen2.5-0.5B", format="gguf").generate() + assert "base_model: Qwen/Qwen2.5-0.5B" in card + + def test_card_uses_actual_gguf_filename(self): + card = ModelCardGenerator( + repo_id="me/other-name", + base_model="Qwen/Qwen2.5-0.5B", + format="gguf", + quantization="Q4_K_M", + gguf_filename="Qwen2.5-0.5B.Q4_K_M.gguf", + ).generate() + assert 'filename="Qwen2.5-0.5B.Q4_K_M.gguf"' in card + assert "other-name.Q4_K_M.gguf" not in card + + +class TestPushToHubMixin: + def _model(self, name_or_path): + from quantllm.core.push import PushToHubMixin + + class Stub(PushToHubMixin): + pass + + stub = Stub() + stub.model = SimpleNamespace(config=SimpleNamespace(_name_or_path=name_or_path)) + stub.export_push_config = {} + + def fake_export(*, format, output_path, quantization=None, **kwargs): + with open(output_path, "wb") as f: + f.write(b"GGUF") + return output_path + + stub.export = fake_export + return stub + + def test_gguf_push_private_filename_and_sanitized_base(self, fake_api): + uploaded = {} + + def capture(**kwargs): + folder = kwargs["folder_path"] + uploaded["files"] = sorted(os.listdir(folder)) + with open(os.path.join(folder, "README.md"), encoding="utf-8") as f: + uploaded["readme"] = f.read() + + fake_api.upload_folder.side_effect = capture + model = self._model(os.path.join(os.sep, "home", "user", "ckpt", "my-llama")) + model.push_to_hub("me/published", format="gguf", quantization="Q4_K_M", private=True) + + assert fake_api.create_repo.call_args.kwargs["private"] is True + assert "my-llama.Q4_K_M.gguf" in uploaded["files"] + readme = uploaded["readme"] + assert 'filename="my-llama.Q4_K_M.gguf"' in readme + assert "base_model:" not in readme.split("---", 2)[1] + assert "ckpt" not in readme diff --git a/tests/test_loading_fixes.py b/tests/test_loading_fixes.py new file mode 100644 index 0000000..b1f567a --- /dev/null +++ b/tests/test_loading_fixes.py @@ -0,0 +1,526 @@ +"""CPU-only regression tests for model loading / configuration fixes.""" + +import inspect +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn + +import quantllm.core.model_loader as model_loader_module +import quantllm.core.smart_config as smart_config_module +import quantllm.core.turbo_model as turbo_model_module +from quantllm.core.flash_attention import _pytorch_attention +from quantllm.core.hardware import GPUInfo, HardwareProfile, HardwareProfiler +from quantllm.core.memory import DynamicOffloader +from quantllm.core.model_analyzer import ModelAnalyzer, ModelInfo +from quantllm.core.smart_config import SmartConfig +from quantllm.core.turbo_model import TurboModel + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _gpu(index, total_gb, cc, free_gb=None): + return GPUInfo( + index=index, + name=f"gpu{index}", + memory_total_gb=total_gb, + memory_free_gb=total_gb if free_gb is None else free_gb, + compute_capability=cc, + supports_bf16=cc >= (8, 0), + supports_flash_attention=cc >= (8, 0), + ) + + +def _profile(gpus, best_gpu_index=0): + best = next((g for g in gpus if g.index == best_gpu_index), None) + return HardwareProfile( + cpu_count=4, + cpu_name="cpu", + system_memory_gb=32.0, + num_gpus=len(gpus), + gpus=gpus, + cuda_available=bool(gpus), + cuda_version="12.1" if gpus else None, + supports_bf16=any(g.supports_bf16 for g in gpus), # deliberately "any" to catch misuse + supports_flash_attention=any(g.supports_flash_attention for g in gpus), + best_gpu_index=best.index if best else 0, + total_gpu_memory_gb=sum(g.memory_total_gb for g in gpus), + ) + + +def _small_model_info(): + return ModelInfo( + model_name="org/tiny-0.5b", + num_params=500_000_000, + size_gb=500_000_000 * 4 / (1024**3), + hidden_size=1024, + num_layers=24, + num_attention_heads=16, + num_kv_heads=16, + intermediate_size=4096, + vocab_size=32000, + max_position_embeddings=4096, + model_type="llama", + ) + + +def _patch_detect(monkeypatch, profile, info=None): + monkeypatch.setattr(HardwareProfiler, "detect", classmethod(lambda cls, force_refresh=False: profile)) + monkeypatch.setattr( + ModelAnalyzer, + "analyze", + classmethod(lambda cls, *a, **k: info or _small_model_info()), + ) + + +# --------------------------------------------------------------------------- +# model_analyzer +# --------------------------------------------------------------------------- + + +def test_moe_param_estimate_counts_experts(): + mixtral = SimpleNamespace( + hidden_size=4096, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=8, + intermediate_size=14336, + vocab_size=32000, + max_position_embeddings=32768, + model_type="mixtral", + num_local_experts=8, + ) + info = ModelAnalyzer._extract_info_from_config("mistralai/Mixtral-8x7B-v0.1", mixtral) + assert 45e9 < info.num_params < 48e9 # Mixtral 8x7B is ~46.7B total + + dense = SimpleNamespace(**{k: v for k, v in vars(mixtral).items() if k != "num_local_experts"}) + dense_info = ModelAnalyzer._extract_info_from_config("dense", dense) + assert dense_info.num_params < 8e9 + + +def test_moe_param_estimate_uses_moe_intermediate_and_shared_experts(): + base = dict( + hidden_size=2048, + num_hidden_layers=10, + num_attention_heads=16, + num_key_value_heads=16, + intermediate_size=10944, + vocab_size=100000, + model_type="deepseek_v2", + n_routed_experts=64, + moe_intermediate_size=1408, + ) + no_shared = ModelAnalyzer._extract_info_from_config("a", SimpleNamespace(**base)) + shared = ModelAnalyzer._extract_info_from_config("b", SimpleNamespace(**base, n_shared_experts=2)) + expected_shared = 2 * 3 * 2048 * 1408 * 10 + assert shared.num_params - no_shared.num_params == expected_shared + # routed experts use moe_intermediate_size, not intermediate_size + assert no_shared.num_params < 10 * 64 * 3 * 2048 * 10944 + + +@pytest.mark.parametrize( + "name,expected", + [ + ("Qwen/Qwen2.5-0.5B-Instruct", 500_000_000), + ("org/model-1.5b", 1_500_000_000), + ("org/model-13B-chat", 13_000_000_000), + ("org/unknown", 7_000_000_000), + ], +) +def test_estimate_from_name_handles_decimals(name, expected): + assert ModelAnalyzer._estimate_from_name(name).num_params == expected + + +def test_analyze_defaults_to_no_remote_code(monkeypatch): + assert inspect.signature(ModelAnalyzer.analyze).parameters["trust_remote_code"].default is False + + import transformers + + captured = {} + + def _fake_from_pretrained(name, **kwargs): + captured.update(kwargs) + raise OSError("offline") + + monkeypatch.setattr(ModelAnalyzer, "_cache", {}) + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", _fake_from_pretrained) + ModelAnalyzer.analyze("org/some-model-3b", token="tok") + assert captured["trust_remote_code"] is False + assert captured["token"] == "tok" + + +# --------------------------------------------------------------------------- +# smart_config / hardware +# --------------------------------------------------------------------------- + + +def test_detect_forwards_trust_remote_code_and_token(monkeypatch): + calls = {} + + def _analyze(cls, name, **kwargs): + calls.update(kwargs) + return _small_model_info() + + monkeypatch.setattr(HardwareProfiler, "detect", classmethod(lambda cls, force_refresh=False: _profile([]))) + monkeypatch.setattr(ModelAnalyzer, "analyze", classmethod(_analyze)) + SmartConfig.detect("org/x", trust_remote_code=True, token="abc") + assert calls == {"trust_remote_code": True, "token": "abc"} + + +def test_no_gpu_defaults_to_full_precision(monkeypatch): + _patch_detect(monkeypatch, _profile([])) + cfg = SmartConfig.detect("org/x") + assert cfg.bits == 16 + assert cfg.effective_loading_bits == 16 + assert cfg.device.type == "cpu" + assert cfg.dtype == torch.float32 + assert cfg.use_flash_attention is False + + +def test_bf16_comes_from_selected_gpu(monkeypatch): + # GPU 0: biggest but pre-Ampere; GPU 1: Ampere. any() would wrongly pick bf16. + gpus = [_gpu(0, 24.0, (7, 5)), _gpu(1, 16.0, (8, 0))] + _patch_detect(monkeypatch, _profile(gpus, best_gpu_index=0)) + + cfg = SmartConfig.detect("org/x") + assert cfg.device == torch.device("cuda:0") + assert cfg.dtype == torch.float16 + assert cfg.use_flash_attention is False + + cfg1 = SmartConfig.detect("org/x", device="cuda:1") + assert cfg1.dtype == torch.bfloat16 + assert cfg1.use_flash_attention is True + + +def test_compile_model_not_auto_enabled(monkeypatch): + _patch_detect(monkeypatch, _profile([_gpu(0, 80.0, (9, 0))])) + assert SmartConfig.detect("org/x").compile_model is False + + +def test_size_estimates_use_effective_loading_bits(monkeypatch): + _patch_detect(monkeypatch, _profile([_gpu(0, 80.0, (9, 0))])) + cfg = SmartConfig.detect("org/x", bits=3) + info = _small_model_info() + assert cfg.stats["quant_size"] == pytest.approx(info.estimated_size_at_bits(4)) + + +def test_best_gpu_index_is_cuda_ordinal(monkeypatch): + # GPU 0 fails to report info, so the list has only the device with ordinal 1. + monkeypatch.setattr(HardwareProfiler, "_cached_profile", None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr( + HardwareProfiler, + "_get_gpu_info", + staticmethod(lambda i: None if i == 0 else _gpu(1, 16.0, (8, 6))), + ) + try: + profile = HardwareProfiler.detect(force_refresh=True) + assert profile.best_gpu_index == 1 + assert profile.best_gpu.index == 1 + assert profile.device == torch.device("cuda:1") + assert profile.supports_bf16 is True + finally: + HardwareProfiler._cached_profile = None + + +def test_gpu_info_uses_mem_get_info_without_set_device(monkeypatch): + props = SimpleNamespace(total_memory=16 * 1024**3, major=8, minor=0, name="fake", multi_processor_count=10) + set_device = Mock() + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda i: props) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda i: (4 * 1024**3, 16 * 1024**3)) + monkeypatch.setattr(torch.cuda, "set_device", set_device) + info = HardwareProfiler._get_gpu_info(0) + assert info.memory_free_gb == pytest.approx(4.0) + set_device.assert_not_called() + + +# --------------------------------------------------------------------------- +# turbo() / model_loader +# --------------------------------------------------------------------------- + + +def test_turbo_pops_trust_remote_code(monkeypatch): + captured = [] + monkeypatch.setattr( + TurboModel, + "from_pretrained", + classmethod(lambda cls, model, **kwargs: captured.append(kwargs)), + ) + turbo_model_module.turbo("org/x") + turbo_model_module.turbo("org/x", trust_remote_code=True) # must not raise TypeError + assert captured[0]["trust_remote_code"] is False + assert captured[1]["trust_remote_code"] is True + + +def _call_fallback(monkeypatch, fake_auto, registered=None): + monkeypatch.setattr(TurboModel, "_architecture_registry", {}) + monkeypatch.setattr(TurboModel, "_model_class_registry", {"llama": registered} if registered else {}) + monkeypatch.setattr(model_loader_module, "AutoModelForCausalLM", fake_auto) + return TurboModel._load_model_with_fallback( + "org/llama-7b", + {"torch_dtype": torch.float32}, + trust_remote_code=False, + hf_config=SimpleNamespace(model_type="llama"), + model_type_override=None, + base_model_fallback=True, + from_config_only=False, + ) + + +@pytest.mark.parametrize( + "error", + [ + OSError("org/llama-7b is not a valid model identifier (401 gated)"), + ImportError("needs package foo"), + RuntimeError("something unrelated"), + ], +) +def test_non_architecture_errors_are_reraised(monkeypatch, error): + registered = Mock() + + class _FakeAuto: + @staticmethod + def from_pretrained(*args, **kwargs): + raise error + + with pytest.raises(type(error)) as excinfo: + _call_fallback(monkeypatch, _FakeAuto, registered=registered) + assert excinfo.value is error + registered.from_pretrained.assert_not_called() + + +def test_architecture_error_message_includes_primary(monkeypatch): + primary = ValueError("Unrecognized configuration class FooConfig") + + class _FakeAuto: + @staticmethod + def from_pretrained(*args, **kwargs): + if "config" in kwargs: + raise ValueError("still bad") + raise primary + + with pytest.raises(RuntimeError) as excinfo: + _call_fallback(monkeypatch, _FakeAuto) + assert "Unrecognized configuration class FooConfig" in str(excinfo.value) + assert excinfo.value.__cause__ is primary + + +def test_auto_attn_implementation_retries_without_it(monkeypatch): + calls = [] + + class _FakeAuto: + @staticmethod + def from_pretrained(*args, **kwargs): + calls.append(dict(kwargs)) + if "attn_implementation" in kwargs: + raise ValueError( + "Model does not support an attention implementation through scaled_dot_product_attention" + ) + return "model" + + monkeypatch.setattr(model_loader_module, "AutoModelForCausalLM", _FakeAuto) + result = TurboModel._load_model_with_fallback( + "org/x", + {"attn_implementation": "sdpa"}, + trust_remote_code=False, + hf_config=None, + model_type_override=None, + base_model_fallback=True, + from_config_only=False, + auto_attn_implementation=True, + ) + assert result == "model" + assert "attn_implementation" not in calls[-1] + + +def test_choose_attn_implementation_cpu_is_sdpa(): + cfg = SimpleNamespace(dtype=torch.float32, device=torch.device("cpu"), use_flash_attention=True) + assert TurboModel._choose_attn_implementation(cfg) == "sdpa" + cfg = SimpleNamespace(dtype="float16", device="cpu", use_flash_attention=False) + assert TurboModel._choose_attn_implementation(cfg) == "sdpa" + + +def test_quantization_kwargs_empty_without_accelerator(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + if hasattr(torch, "xpu"): + monkeypatch.setattr(torch.xpu, "is_available", lambda: False) + cfg = SimpleNamespace(bits=4, dtype=torch.float16, cpu_offload=False) + assert TurboModel._get_quantization_kwargs(cfg) == {} + + +def test_quantization_kwargs_empty_without_bitsandbytes(monkeypatch): + real_find_spec = model_loader_module.importlib.util.find_spec + monkeypatch.setattr( + model_loader_module.importlib.util, + "find_spec", + lambda name, *a: None if name == "bitsandbytes" else real_find_spec(name, *a), + ) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + cfg = SimpleNamespace(bits=4, dtype=torch.float16, cpu_offload=False) + assert TurboModel._get_quantization_kwargs(cfg) == {} + + +def test_from_pretrained_passes_attn_implementation_and_skips_compile(monkeypatch): + monkeypatch.setattr(TurboModel, "_architecture_registry", {}) + monkeypatch.setattr(TurboModel, "_model_class_registry", {}) + cfg = SimpleNamespace( + bits=16, + effective_loading_bits=16, + dtype=torch.float32, + cpu_offload=False, + device=torch.device("cpu"), + gradient_checkpointing=False, + use_flash_attention=False, + compile_model=False, + print_summary=lambda: None, + ) + monkeypatch.setattr(model_loader_module.SmartConfig, "detect", lambda *a, **k: cfg) + monkeypatch.setattr( + model_loader_module.AutoTokenizer, + "from_pretrained", + lambda *a, **k: SimpleNamespace(pad_token="x", eos_token="x", eos_token_id=0), + ) + import transformers + + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda *a, **k: SimpleNamespace(model_type="llama", quantization_config=None), + ) + seen = {} + + class _FakeAuto: + @staticmethod + def from_pretrained(*args, **kwargs): + seen.update(kwargs) + return nn.Linear(2, 2) + + monkeypatch.setattr(model_loader_module, "AutoModelForCausalLM", _FakeAuto) + loaded = TurboModel.from_pretrained("org/llama-1b", quantize=False, verbose=False) + assert seen["attn_implementation"] == "sdpa" + assert isinstance(loaded.model, nn.Linear) # not wrapped by torch.compile + + TurboModel.from_pretrained("org/llama-1b", quantize=False, verbose=False, attn_implementation="eager") + assert seen["attn_implementation"] == "eager" + + +# --------------------------------------------------------------------------- +# turbo_model: dequantize / optimize_inference +# --------------------------------------------------------------------------- + + +def _bare_turbo(model, finetuned=False, lora=False): + tm = TurboModel.__new__(TurboModel) + tm.model = model + tm.tokenizer = None + tm.config = SimpleNamespace(dtype=torch.float16, bits=4) + tm._is_quantized_override = None + tm._is_finetuned = finetuned + tm._lora_applied = lora + tm.verbose = False + return tm + + +def test_dequantize_merges_lora_and_uses_in_memory_model(monkeypatch): + inner = nn.Linear(2, 2) + inner.config = SimpleNamespace(_name_or_path="org/base") + inner.dequantize = Mock(return_value=inner) + + peft_like = SimpleNamespace(merge_and_unload=Mock(return_value=inner), config=inner.config) + reload = Mock() + monkeypatch.setattr(turbo_model_module.AutoModelForCausalLM, "from_pretrained", reload) + + tm = _bare_turbo(peft_like, finetuned=True, lora=True) + result = tm._dequantize_model() + + peft_like.merge_and_unload.assert_called_once() + inner.dequantize.assert_called_once() + reload.assert_not_called() + assert result is inner + assert result.weight.dtype == torch.float16 + assert tm._lora_applied is False + + +def test_dequantize_never_reloads_hub_after_finetuning(monkeypatch): + model = nn.Linear(2, 2) + model.config = SimpleNamespace(_name_or_path="org/base") + model.dequantize = Mock(side_effect=ValueError("no quantizer")) + monkeypatch.setattr( + TurboModel, "_manual_bnb_dequantize", staticmethod(Mock(side_effect=RuntimeError("manual failed"))) + ) + reload = Mock() + monkeypatch.setattr(turbo_model_module.AutoModelForCausalLM, "from_pretrained", reload) + + tm = _bare_turbo(model, finetuned=True) + with pytest.raises(RuntimeError, match="Fine-tuned weights"): + tm._dequantize_model() + reload.assert_not_called() + + +def test_optimize_inference_noop_without_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + tm = _bare_turbo(nn.Linear(2, 2)) + replace = Mock() + monkeypatch.setattr(TurboModel, "_replace_with_triton", replace) + tm.optimize_inference() + replace.assert_not_called() + + +# --------------------------------------------------------------------------- +# flash_attention fallback +# --------------------------------------------------------------------------- + + +def _reference_attention(q, k, v, causal): + # q/k/v: [b, s, h, d] with equal head counts + q, k, v = (t.transpose(1, 2) for t in (q, k, v)) + scores = q @ k.transpose(-2, -1) / q.shape[-1] ** 0.5 + if causal: + sq, sk = q.shape[2], k.shape[2] + mask = torch.ones(sq, sk, dtype=torch.bool).tril(diagonal=sk - sq) + scores = scores.masked_fill(~mask, float("-inf")) + return (scores.softmax(-1) @ v).transpose(1, 2) + + +def test_pytorch_attention_gqa_and_bottom_right_causal(): + torch.manual_seed(0) + q = torch.randn(2, 3, 4, 8) + k = torch.randn(2, 5, 2, 8) + v = torch.randn(2, 5, 2, 8) + out = _pytorch_attention(q, k, v, causal=True) + expected = _reference_attention(q, k.repeat_interleave(2, dim=2), v.repeat_interleave(2, dim=2), causal=True) + assert out.shape == (2, 3, 4, 8) + assert torch.allclose(out, expected, atol=1e-5) + + # single decode token attends to the whole cache + out1 = _pytorch_attention(q[:, :1], k, v, causal=True) + expected1 = _reference_attention( + q[:, :1], k.repeat_interleave(2, dim=2), v.repeat_interleave(2, dim=2), causal=False + ) + assert torch.allclose(out1, expected1, atol=1e-5) + + +# --------------------------------------------------------------------------- +# memory +# --------------------------------------------------------------------------- + + +class _TinyDecoder(nn.Module): + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)]) + self.lm_head = nn.Linear(4, 4) + + +def test_dynamic_offloader_never_selects_container(): + offloader = DynamicOffloader(_TinyDecoder(), max_gpu_layers=2, device="cpu") + assert offloader.offloadable_layers == ["model.layers.0", "model.layers.1", "model.layers.2"] + offloader.move_layer_to_gpu("model.layers.2") # must not hardcode "cuda" + assert offloader._layer_locations["model.layers.2"] == "gpu" diff --git a/tests/test_registry.py b/tests/test_registry.py index cdd24c7..198493c 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -1,7 +1,28 @@ -"""Tests for the QuantLLM model registry.""" +"""Tests for the QuantLLM model registry. + +All tests are offline-deterministic: the Hub is never contacted. +""" + +from types import SimpleNamespace +from unittest import mock import pytest +import quantllm.registry as registry_mod + +_REAL_FETCH = registry_mod.ModelRegistry._fetch_from_hub + + +@pytest.fixture(autouse=True) +def _offline_registry(monkeypatch): + """Force the built-in model list and reset all caches.""" + monkeypatch.setattr(registry_mod.ModelRegistry, "_fetch_from_hub", lambda self: []) + registry_mod._HUB_CACHE.clear() + registry_mod._registry._models = None + yield + registry_mod._HUB_CACHE.clear() + registry_mod._registry._models = None + class TestModelRegistry: def test_builtin_models_exist(self): @@ -77,3 +98,109 @@ def test_model_info(self): output = model_info("QuantLLM/TinyLlama-1.1B-GGUF") assert "Params:" in output assert "1.1" in output # 1.1B params + + +def _fake_api(models): + api = mock.MagicMock() + api.list_models.return_value = models + return api + + +class TestHubFetch: + def _fetch(self, models, monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.setattr(registry_mod, "_hub_offline", lambda: False) + api = _fake_api(models) + with mock.patch("huggingface_hub.HfApi", return_value=api): + result = _REAL_FETCH(registry_mod.ModelRegistry(org="TestOrg")) + return result, api + + def test_new_style_id_attribute(self, monkeypatch): + m = SimpleNamespace(id="TestOrg/Foo-GGUF", card_data={"quantllm_family": "Qwen", "params": 1.5}) + result, api = self._fetch([m], monkeypatch) + assert result[0]["name"] == "TestOrg/Foo-GGUF" + assert result[0]["family"] == "qwen" + assert result[0]["params"] == 1.5 + assert "gguf" in result[0]["formats"] + # one list call with card data, no per-model model_info calls + assert api.list_models.call_args.kwargs.get("cardData") is True + api.model_info.assert_not_called() + + def test_legacy_model_id_attribute(self, monkeypatch): + m = SimpleNamespace(modelId="TestOrg/Legacy", cardData={"tags": ["quantllm", "mistral"]}) + result, _ = self._fetch([m], monkeypatch) + assert result[0]["name"] == "TestOrg/Legacy" + assert result[0]["family"] == "mistral" + + def test_card_data_object_with_to_dict(self, monkeypatch): + cd = SimpleNamespace(to_dict=lambda: {"quantllm_family": "phi"}) + m = SimpleNamespace(id="TestOrg/Phi", card_data=cd) + result, _ = self._fetch([m], monkeypatch) + assert result[0]["family"] == "phi" + + def test_results_cached_in_memory(self, monkeypatch): + m = SimpleNamespace(id="TestOrg/Cached", card_data={}) + self._fetch([m], monkeypatch) + api = _fake_api([]) + with mock.patch("huggingface_hub.HfApi", return_value=api): + again = _REAL_FETCH(registry_mod.ModelRegistry(org="TestOrg")) + assert again[0]["name"] == "TestOrg/Cached" + api.list_models.assert_not_called() + + def test_hub_errors_fall_back_to_empty(self, monkeypatch): + monkeypatch.setattr(registry_mod, "_hub_offline", lambda: False) + api = mock.MagicMock() + api.list_models.side_effect = OSError("network down") + with mock.patch("huggingface_hub.HfApi", return_value=api): + assert _REAL_FETCH(registry_mod.ModelRegistry(org="TestOrg")) == [] + + def test_offline_mode_skips_hub(self, monkeypatch): + monkeypatch.setattr(registry_mod, "_hub_offline", lambda: True) + with mock.patch("huggingface_hub.HfApi") as hf_api: + assert _REAL_FETCH(registry_mod.ModelRegistry(org="TestOrg")) == [] + hf_api.assert_not_called() + + +class TestFamilyParsing: + def test_dedicated_key_wins(self): + assert registry_mod._parse_family({"quantllm_family": "Llama", "tags": ["quantllm", "gguf"]}) == "llama" + + def test_skips_generic_quantllm_tag(self): + assert registry_mod._parse_family({"tags": ["quantllm", "gemma"]}) == "gemma" + + def test_unknown_when_missing(self): + assert registry_mod._parse_family({}) == "unknown" + assert registry_mod._parse_family({"tags": ["quantllm"]}) == "unknown" + + +class TestRegisterModel: + def test_existing_card_updates_metadata_only(self): + api = mock.MagicMock() + api.file_exists.return_value = True + existing = SimpleNamespace(data=SimpleNamespace(tags=["custom"])) + with ( + mock.patch("huggingface_hub.HfApi", return_value=api), + mock.patch("huggingface_hub.ModelCard.load", return_value=existing), + mock.patch("huggingface_hub.metadata_update") as mu, + ): + ok = registry_mod.ModelRegistry().register_model("me/model", {"family": "qwen"}) + assert ok is True + api.upload_file.assert_not_called() + meta = mu.call_args.args[1] + assert meta["quantllm_family"] == "qwen" + assert meta["tags"] == ["custom", "quantllm", "qwen"] + + def test_new_repo_uploads_card(self): + api = mock.MagicMock() + api.file_exists.return_value = False + with mock.patch("huggingface_hub.HfApi", return_value=api): + ok = registry_mod.ModelRegistry().register_model("me/model", {"family": "qwen"}) + assert ok is True + content = api.upload_file.call_args.kwargs["path_or_fileobj"].decode() + assert '"quantllm_family": "qwen"' in content + + def test_failure_returns_false(self): + api = mock.MagicMock() + api.create_repo.side_effect = RuntimeError("401") + with mock.patch("huggingface_hub.HfApi", return_value=api): + assert registry_mod.ModelRegistry().register_model("me/model", {}) is False diff --git a/tests/test_server.py b/tests/test_server.py index 889ffc6..cdafcc8 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -118,3 +118,239 @@ def test_serve_imports(self): from quantllm.server.app import run_server assert callable(run_server) + + +# --------------------------------------------------------------------------- +# Fake transformers-style model/tokenizer (CPU, no downloads) +# --------------------------------------------------------------------------- + +_VOCAB = {0: "", 1: "Hello", 2: " world", 3: "STOP", 4: " after", 5: " more"} + + +class _FakeTokenizer: + pad_token_id = 0 + eos_token_id = 0 + chat_template = None + + def __init__(self): + self.calls = [] + + def __call__(self, text, return_tensors=None, add_special_tokens=True): + import torch + + self.calls.append(add_special_tokens) + return {"input_ids": torch.tensor([[1, 2, 3]])} + + def encode(self, text, add_special_tokens=True): + return text.split() + + def decode(self, ids, skip_special_tokens=True): + return "".join(_VOCAB[int(i)] for i in ids) + + def apply_chat_template(self, *args, **kwargs): + raise ValueError("no chat template") + + +class _FakeModel: + device = "cpu" + + def __init__(self, ids=(1, 2), chunks=None, delay=0.0, error=None): + self.ids = list(ids) + self.chunks = chunks + self.delay = delay + self.error = error + self.calls = [] + + def generate(self, input_ids=None, streamer=None, **kwargs): + import time as _time + + import torch + + self.calls.append(kwargs) + if self.error is not None: + raise self.error + if streamer is not None: + for chunk in self.chunks or [_VOCAB[i] for i in self.ids]: + _time.sleep(self.delay) + streamer.on_finalized_text(chunk) + return None + _time.sleep(self.delay) + return torch.cat([input_ids, torch.tensor([self.ids])], dim=1) + + +def _fake_engine(**model_kwargs): + from quantllm.server import InferenceEngine, ServerConfig + + engine = InferenceEngine(ServerConfig(model="fake-model")) + engine.tokenizer = _FakeTokenizer() + engine.model = _FakeModel(**model_kwargs) + return engine + + +def _sse_events(body: str): + events = [line[len("data: ") :] for line in body.split("\n\n") if line.startswith("data: ")] + assert body.endswith("data: [DONE]\n\n") + assert events[-1] == "[DONE]" + return [json.loads(e) for e in events[:-1]] + + +class TestServerFixes: + def test_trust_remote_code_default_off(self): + from quantllm.server import ServerConfig + + assert ServerConfig().trust_remote_code is False + assert ServerConfig().gguf_file is None + + def test_gguf_detection(self, tmp_path): + from quantllm.server import InferenceEngine, ServerConfig + + local = tmp_path / "m.gguf" + local.write_bytes(b"") + assert InferenceEngine(ServerConfig(model=str(local)))._is_local_gguf() + hub = InferenceEngine(ServerConfig(model="org/Model-7B-GGUF")) + assert hub._is_hub_gguf() and not hub._is_local_gguf() + assert not InferenceEngine(ServerConfig(model="/missing/model.gguf"))._is_local_gguf() + assert not InferenceEngine(ServerConfig(model="org/plain-model"))._is_hub_gguf() + + def test_no_eos_override_and_stop_strings(self): + engine = _fake_engine(ids=[1, 3, 4]) + text = engine.generate("hi", max_tokens=8, stop=["STOP"]) + assert text == "Hello" + kwargs = engine.model.calls[0] + assert "eos_token_id" not in kwargs + assert kwargs.get("stop_strings") == ["STOP"] + + def test_chat_template_prompt_skips_special_tokens(self): + from quantllm.server.engine import create_chat_completion + + engine = _fake_engine() + engine.tokenizer.apply_chat_template = lambda *a, **k: "hi" + create_chat_completion(engine, [{"role": "user", "content": "hi"}], max_tokens=8) + assert engine.tokenizer.calls == [False] + + def test_filter_stops_emits_prefix_across_chunks(self): + from quantllm.server.engine import _filter_stops + + assert "".join(_filter_stops(iter(["Hel", "loST", "OP after"]), ["STOP"])) == "Hello" + assert "".join(_filter_stops(iter(["abc STOP"]), ["STOP"])) == "abc " + assert "".join(_filter_stops(iter(["no", " stop"]), ["STOP"])) == "no stop" + + def test_stream_generation_error_does_not_hang(self): + import threading + + engine = _fake_engine(error=RuntimeError("boom")) + outcome = {} + + def consume(): + try: + list(engine.generate("hi", max_tokens=4, stream=True)) + except RuntimeError as e: + outcome["error"] = str(e) + + t = threading.Thread(target=consume, daemon=True) + t.start() + t.join(timeout=15) + assert not t.is_alive(), "stream hung after generate() raised" + assert outcome.get("error") == "boom" + # The model lock must be released so the next request can run. + assert engine._lock.acquire(timeout=1) + engine._lock.release() + + +class TestServerApp: + @pytest.fixture(autouse=True) + def _require_fastapi(self): + pytest.importorskip("fastapi") + pytest.importorskip("httpx") + + def _client(self, engine): + from fastapi.testclient import TestClient + + from quantllm.server.app import _build_app + + return TestClient(_build_app(engine)) + + def test_stop_as_string_and_content_parts(self): + engine = _fake_engine(ids=[1, 3, 4]) + body = { + "messages": [ + {"role": "system", "content": None}, + {"role": "user", "content": [{"type": "text", "text": "Hi "}, {"type": "text", "text": "there"}]}, + ], + "stop": "STOP", + "max_tokens": 16, + } + resp = self._client(engine).post("/v1/chat/completions", json=body) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["choices"][0]["message"]["content"] == "Hello" + assert data["choices"][0]["finish_reason"] == "stop" + assert engine.model.calls[0]["stop_strings"] == ["STOP"] + + def test_finish_reason_length(self): + engine = _fake_engine(ids=[1, 2]) + resp = self._client(engine).post("/v1/completions", json={"prompt": "x", "max_tokens": 2}) + data = resp.json() + assert data["choices"][0]["text"] == "Hello world" + assert data["choices"][0]["finish_reason"] == "length" + assert data["usage"]["completion_tokens"] == 2 + + def test_chat_stream_format(self): + engine = _fake_engine(chunks=["Hel", "loST", "OP after"]) + body = {"messages": [{"role": "user", "content": "hi"}], "stream": True, "stop": ["STOP"]} + resp = self._client(engine).post("/v1/chat/completions", json=body) + events = _sse_events(resp.text) + assert all(e["object"] == "chat.completion.chunk" for e in events) + assert events[0]["choices"][0]["delta"]["role"] == "assistant" + text = "".join(e["choices"][0]["delta"].get("content", "") for e in events) + assert text == "Hello" + assert events[-1]["choices"][0]["finish_reason"] == "stop" + + def test_completion_stream_format(self): + engine = _fake_engine(chunks=["a", " b", " c"]) + resp = self._client(engine).post("/v1/completions", json={"prompt": "x", "stream": True, "max_tokens": 3}) + events = _sse_events(resp.text) + for e in events: + assert e["object"] == "text_completion" + assert "text" in e["choices"][0] and "delta" not in e["choices"][0] + assert "".join(e["choices"][0]["text"] for e in events) == "a b c" + assert events[-1]["choices"][0]["finish_reason"] == "length" + + def test_stream_error_terminates(self): + engine = _fake_engine(error=RuntimeError("boom")) + resp = self._client(engine).post("/v1/completions", json={"prompt": "x", "stream": True}) + assert resp.text.endswith("data: [DONE]\n\n") + assert "boom" in resp.text + + @pytest.mark.parametrize("stream", [False, True]) + def test_health_responsive_during_generation(self, stream): + import asyncio + import time as _time + + import httpx + + from quantllm.server.app import _build_app + + engine = _fake_engine(chunks=["a", "b", "c", "d"], delay=0.4) + app = _build_app(engine) + + async def run(): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + done = {} + + async def gen(): + await client.post("/v1/completions", json={"prompt": "x", "stream": stream}) + done["gen"] = _time.perf_counter() + + async def health(): + await asyncio.sleep(0.2) + r = await client.get("/health") + assert r.status_code == 200 + done["health"] = _time.perf_counter() + + await asyncio.wait_for(asyncio.gather(gen(), health()), timeout=30) + return done + + done = asyncio.run(run()) + assert done["health"] < done["gen"], "/health blocked while generating" From b2d8fdb7cd51ed80f1f8bdd6330376003335071a Mon Sep 17 00:00:00 2001 From: codewithdark-git Date: Sun, 1 Nov 2026 20:39:07 +0500 Subject: [PATCH 2/2] ci: pin ruff and fix README code block formatting Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 2 +- README.md | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7016073..19d1ee3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,7 +27,7 @@ jobs: - name: Install ruff run: | python -m pip install --upgrade pip - pip install "ruff>=0.5.0" + pip install "ruff==0.15.21" # pinned: newer ruff also reformats Markdown code blocks - name: Run ruff check run: ruff check . - name: Run ruff format check diff --git a/README.md b/README.md index cd45d44..368d5f5 100644 --- a/README.md +++ b/README.md @@ -189,6 +189,7 @@ Complete end-to-end pipeline demonstrating QuantLLM on a new/public model, run o ### 1. Load & Auto-Quantize ```python from quantllm import turbo + model = turbo("Qwen/Qwen2.5-7B-Instruct") # 7.62B params, BnB 8-bit on H100 ```