Load from_pretrained at target dtype; restrict quantize= to int8

[maxwbuckley]

Important

No released wheel is affected. The current PyPI release — gliner 0.2.26, uploaded 2026-03-19 — pre-dates PR #342 (merged 2026-03-31), which introduced the quantize= surface. So removing the pure-downcast aliases here only reaches users installing from main / a git SHA. No one on pip install gliner needs to change anything.

Summary

Two related changes to the precision / quantization surface, landed together because they form one coherent story:

1. Add dtype= to GLiNER.from_pretrained — load weights directly at the target floating-point precision. Casts each state-dict tensor during the safe_open read and pre-casts the model shell before load_state_dict, so a full fp32 copy is never materialized. Accepts "bf16"/"fp16"/"float32"/... or a floating-point torch.dtype; non-floating dtypes are rejected up front with a clear error that points at quantize="int8" for int paths. Int/bool buffers are left untouched.
2. Restrict quantize= to "int8" only (per @Ingvarstep's review) — drops the pure-downcast aliases (True, "fp16", "bf16") entirely; those now raise ValueError with a migration message pointing at dtype= / model.to(...). quantize="int8" is unchanged: torchao int8 weight-only on GPU, FBGEMM dynamic quant on CPU.

Why dtype= matters

Previously from_pretrained always loaded the state dict at its on-disk precision (typically fp32), copied it into an fp32 model, then required a separate quantize(...) pass to cast down. Two copies co-resident, plus a post-load cast pass. For users who just want bf16/fp16 inference, that was ~2× the peak memory and an extra pass.

This is particularly useful for cold starts and scalable serverless deployments (AWS Lambda, Cloud Run, Modal, RunPod serverless, autoscaled Kubernetes pods, etc.) where startup latency and peak memory directly drive cost and SLA:

• Shorter cold-start on every new container — one pass instead of load + cast.
• Lower peak memory lets instances fit on smaller memory tiers and reduces boot-time OOMs.
• Faster first-inference latency after scale-from-zero.

For CPU-only loads the peak-memory drop is a clean ~2×→1× fp32. For map_location="cuda", state-dict tensors stream to GPU while the shell is CPU-side, so the win is avoiding a simultaneous fp32 GPU state dict + fp32 GPU model (not literal halving of total footprint) plus the separate cast pass.

Why restrict quantize= to int8

Before this PR the quantize= surface was doing two different things under one name:

quantize=	Device	What happens	Real quant?
True / "fp16"	GPU	model.half()	No — pure downcast
True / "fp16"	CPU	quantize_dynamic(Linear, fp16)	Yes — dynamic quant
"bf16"	GPU	model.bfloat16()	No — pure downcast
"bf16"	CPU	model.bfloat16()	No — pure downcast
"int8"	GPU	torchao int8 weight-only	Yes
"int8"	CPU	quantize_dynamic(Linear, qint8) (FBGEMM)	Yes

Three rows are just .to(dtype) with an extra fp32 intermediate — exactly what dtype= now does cheaply. A fourth (CPU fp16 dynamic quant) had no documented speed benefit and was asymmetric with every other CPU/GPU combination. The remaining two int8 rows are the only genuinely distinct behavior. After this PR:

• quantize=True → ValueError (pointing at quantize='int8' or dtype=).
• quantize="fp16" | "bf16" (any device) → ValueError (pointing at dtype= or model.to(torch.X)).
• quantize="int8" → unchanged: torchao on GPU, FBGEMM dynamic quant on CPU.
• model.quantize(...) takes the same restriction and defaults to "int8" instead of "fp16".

Signature narrows from Union[bool, str] = False to Optional[str] = None. See the callout at the top for why the blast radius is negligible.

Ray Serve integration (gliner.serve)

GLiNERFactory and the gliner.serve CLI are wired through in lockstep:

• gliner/serve/server.py now passes map_location=config.device + dtype=self.torch_dtype to GLiNER.from_pretrained instead of loading at fp32 and then doing a post-load .to(device=..., dtype=...) cast. Serving cold starts get the same peak-memory win as direct users — which is exactly the autoscaled / serverless use case the core change targets.
• --quantization CLI choices narrowed from {fp16, bf16, int8, None} to {int8, None}, matching the core API. Precision stays under --dtype (which already accepted float32/float16/fp16/bfloat16/bf16). docs/serving.md CLI reference + GLINER_QUANTIZATION env-var row updated.
• Without this change, anyone passing --quantization fp16 or --quantization bf16 after the int8 restriction would have hit ValueError at server boot. The serve layer is also post-0.2.26 (PR Feature/serving #346), so no released wheel is affected.

Usage

from gliner import GLiNER
import torch

# Load directly in target precision
model = GLiNER.from_pretrained("urchade/gliner_small-v2.1", dtype="bf16")
model = GLiNER.from_pretrained("urchade/gliner_small-v2.1", dtype=torch.float16, map_location="cuda")

# Post-load, PyTorch's built-in `.to()` covers precision
model.to(torch.bfloat16)

# Real quantization (only remaining `quantize=` value)
model.quantize("int8")  # torchao on GPU, FBGEMM on CPU

Alignment with HuggingFace from_pretrained

This puts GLiNER on the same axis as HuggingFace transformers.PreTrainedModel.from_pretrained, which recently settled on the same split:

• HF's current precision knob is dtype= (renamed from the now-legacy torch_dtype=), accepting a string or torch.dtype. We use the same parameter name and input surface.
• HF separates real quantization into quantization_config=BitsAndBytesConfig(...) / AutoGPTQConfig / AutoAWQConfig. Our quantize="int8" plays the same role; restricting quantize= to int8 mirrors HF's decision not to overload the quantization knob for precision changes.

Users coming from HF will find GLiNER.from_pretrained(..., dtype="bf16") immediately familiar.

Known gap (possible follow-up, not in this PR): HF supports low_cpu_mem_usage=True / device_map=, which uses PyTorch's meta device to skip allocating the random-init shell entirely and then materializes weights via load_state_dict(assign=True). That's strictly cheaper than our pre-cast approach (~0 extra bytes for the shell vs ~1× target-dtype model size for ours). It would mean constructing the backbone on torch.device("meta") rather than via from_config, plus switching to assign=True loading — a bigger surgery than this PR. Worth chasing for the last slice of cold-start memory if it matters.

Tests

New tests/test_quantize_and_dtype.py — 50 unit tests covering:

• _parse_dtype: None passthrough; all string aliases (fp16/float16/half/bf16/bfloat16/fp32/float32/float, case-insensitive); torch.dtype passthrough; unknown strings → ValueError; non-floating torch.dtype (int8/int32/int64/bool/uint8) → ValueError pointing at quantize='int8'; non-str/non-dtype → TypeError.
• _load_state_dict: default path leaves tensors at stored dtype; dtype= casts floats on read in both the safetensors and torch.load branches; int/bool buffers preserved; idempotent when stored dtype already matches.
• model.quantize(): "int8" (case-insensitive) routes to _apply_int8_quantization on both CPU and CUDA; precision aliases (fp16/float16/half/bf16/bfloat16) on both devices → ValueError with migration message naming the correct torch.float16/torch.bfloat16 replacement; unknown strings → ValueError; non-string inputs (True, False, None, ints, floats) → TypeError; ONNX model → RuntimeError (and the ONNX check runs before the alias check).
• from_pretrained(quantize=True): rejected via the shared guard in model.quantize().

Heavier end-to-end tests (real hub checkpoint load with dtype="bf16" + inference parity) are intentionally out of scope for the unit suite — worth a manual smoke before merge.

Test plan

• pytest tests/test_quantize_and_dtype.py → 50/50 pass locally.
• ruff check / ruff format --check clean on new and modified files.
• Manual end-to-end smoke against a real hub checkpoint with dtype="bf16"; confirm inference parity with the legacy quantize="bf16" path.

🤖 Generated with Claude Code

[Ingvarstep]

Hi @maxwbuckley , thank you for your contribution, please, see my review.

[Ingvarstep]

I really like that it aligns with HF standards right now.

[Ingvarstep]

It's discussable, but maybe we should allow only int8 quantizaiton?

[maxwbuckley]

Done — restricted quantize= to "int8" only in 9f551c0.

• quantize=True / "fp16" / "bf16" on from_pretrained and model.quantize(...) now raise ValueError with a migration message pointing at dtype= / model.to(torch.X).
• Dropped the CPU fp16 dynamic-quant path along with the pure-downcast aliases — it had no documented speed benefit and was asymmetric with every other CPU/GPU combination.
• quantize= signature narrowed from Union[bool, str] = False to Optional[str] = None; model.quantize() default changed from "fp16" to "int8".
• Added 50 unit tests in tests/test_quantize_and_dtype.py covering the new validation, the dtype= load path, and _load_state_dict cast-on-read.

Since gliner 0.2.26 (current PyPI, 2026-03-19) pre-dates #342, no released wheel ships the old aliases — blast radius is main / git-SHA installs only.

[Ingvarstep]

It looks good, thank you.

[Ingvarstep]

It looks good

[maxwbuckley]

Awesome. Thanks a lot!