Skip to content

Add qfloat8 offloading#930

Open
PlagueKind wants to merge 2 commits into
ostris:mainfrom
PlagueKind:ostris/feature/qfloat8offload
Open

Add qfloat8 offloading#930
PlagueKind wants to merge 2 commits into
ostris:mainfrom
PlagueKind:ostris/feature/qfloat8offload

Conversation

@PlagueKind

Copy link
Copy Markdown
Contributor

Add qfloat8 (quanto) support with layer offloading and LoRA/LoKr merge

Quanto's qfloat8 quantization was forced to float8 (torchao) in several places because the memory offload and merge-in code paths only recognized torchao's tensor subclass. This adds proper detection for quanto tensors and removes the resulting workarounds.

Changes

  • toolkit/memory_management/manager_modules.py_is_ao_quantized_tensor now also recognizes optimum.quanto.* tensor subclasses, not just torchao.*. Fixes offload/pin/bounce for qfloat8 weights, which were previously silently treated as plain float tensors.
  • toolkit/config_modules.py — removed the forced qfloat8float8 override for layer_offloading, now unnecessary.
  • toolkit/util/quantize.pyis_quantized_tensor now recognizes quanto tensors; requantize_module_weight now re-quantizes quanto modules via .freeze() when no torchao config is available (quanto modules retain their own weight_qtype/optimizer config from creation).
  • toolkit/network_mixins.py, toolkit/models/lokr.py, toolkit/lora_special.py — removed early-returns that silently skipped merge_in/merge_out entirely for quanto-quantized weights ('weight._data' in state_dict). The dequantize → merge → re-quantize logic underneath already worked generically once detection was fixed.
  • extensions_built_in/diffusion_models/krea2/krea2.py — removed the forced qtype override for assistant LoRA loading, now unnecessary since merge_in correctly re-quantizes.

Testing

  • Verified _is_ao_quantized_tensor / _move_params_to_cpu_and_pin against a real quanto QLinear: outer wrapper and inner _data/_scale storage both correctly land on CPU, not just the outer tensor.
  • Verified steady-state VRAM under offload is bounded prefetch-ring cost (PIPELINE_DEPTH × weight size), not per-layer accumulation.
  • Ran toolkit/memory_management/test_memory_manager.py (temporarily swapped to qfloat8) across all 8 offload/checkpoint configs on a 1.2B-param model with a real LoRA network — no exceptions, correct output/gradient match against non-offloaded reference (output diff ~2e-3, input grad diff ~1e-6, consistent with fp8 quant noise).
  • Verified merge_in/merge_out round-trip correctness for both FullModule and LoRAModule against real quanto QBytesTensor weights — weight changes as expected, stays quantized throughout, round-trip diff ~4e-3 (quant noise only).

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