Skip to content

feat: support prequantized BitsAndBytes streaming with Transformers 5 - #347

Open
DevElCuy wants to merge 2 commits into
lyogavin:mainfrom
DevElCuy:feat/prequantized-bnb-streaming
Open

DevElCuy wants to merge 2 commits into
lyogavin:mainfrom
DevElCuy:feat/prequantized-bnb-streaming

Conversation

@DevElCuy

Copy link
Copy Markdown

Summary

Adds support for streaming Hugging Face checkpoints that are already quantized with BitsAndBytes, including compatibility with the Transformers 5
quantizer API.

It also adds an opt-in text-only mode for multimodal models, allowing checkpoints such as Qwen3.8-27B NF4 to run below 3 GB VRAM by leaving the unused
vision tower on the meta device.

What changed

  • Recognize BitsAndBytes quantization metadata as companion state rather than model parameters.
  • Support both:
    • Legacy create_quantized_param.
    • Transformers 5 get_weight_conversions.
  • Avoid CUDA pinning checks on CPU execution paths.
  • Add load_resident_modules=False for low-VRAM text-only inference.
  • Add benchmark options for:
    • Model, device, prompt, and shard directory.
    • AirLLM versus resident inference.
    • Prefetching.
    • Chat templates and disabled thinking.
    • Text-only multimodal execution.
  • Report generated text, throughput, process RSS, accelerator memory, and loading profiles as JSON.
  • Document an isolated ROCm setup for the Radeon RX 6400.
  • Add regression coverage for BnB metadata, Transformers 5 conversion, and text-only resident-module handling.

Motivation

AirLLM previously interpreted BnB tensors such as weight.absmax,
weight.quant_map, and weight.quant_state.bitsandbytes__nf4 as independent
model parameters. This caused streamed prequantized checkpoints to fail during
module resolution.

Transformers 5 also removed the older direct quantized-parameter creation API,
requiring AirLLM to invoke the newer weight-conversion operations.

For multimodal checkpoints, keeping an unused BF16 vision tower resident can
consume enough VRAM to prevent the streamed output head from loading. The new
text-only mode makes that memory optional while preserving the existing default
behavior.

Validation

  • 9 focused regression tests pass.
  • Python compilation passes.
  • Wheel and source distribution build successfully.
  • Tested unsloth/Qwen3-8B-unsloth-bnb-4bit on CPU and ROCm.
  • Tested Qwen3.8-27B NF4 on a 4 GB Radeon RX 6400:
    • Correctly generated Blue.
    • Peak accelerator memory: 2,540.7 MiB.
    • Peak process RSS with prefetch: 4,314.4 MiB.
    • Effective rate: 0.0578 generated tokens/s.
    • Prefetch reduced the measured two-token generation time from 42.9 s to 34.6 s.

Notes

  • No model weights or generated layer shards are included.
  • Resident multimodal modules remain enabled by default.
  • --text-only disables image and video inputs for that model instance.
  • This does not add GGUF support.
  • Hardware-specific setup and benchmark details are documented in
    ROCM_RX6400.md.

Stream Hugging Face BnB quantization state with both legacy and Transformers 5 quantizer APIs. Add CPU-safe prefetch behavior, regression coverage, and a reproducible resident-versus-streamed benchmark.
@DevElCuy
DevElCuy force-pushed the feat/prequantized-bnb-streaming branch from 8f57090 to a3a242f Compare August 24, 2026 22:47

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant