Skip to content

[enhancement]: Load torchao-quantized single-file transformers (unsloth FP8) #9798

Open

Description

@Pfannkuchensack

Is there an existing issue for this?

  • I have searched the existing issues

What should this feature add?

Load single-file transformers saved as torchao-quantized tensors, such as unsloth/Qwen-Image-2.1-FP8 (115k+ downloads). Since #9784 these files are refused at install with a clear message; before that they installed and failed at the first generation with a generic "incomplete checkpoint" error.

Format (read from the real file's header):

  • Each quantized Linear is stored as <layer>._weight_qdata (F8_E4M3, [out, in]) plus <layer>._weight_scale (F32, [out, 1], one scale per output row). There is no <layer>.weight.
  • Norms, img_in, proj_out and timestep_embedder.linear_1 stay bf16 under their plain names; gate_layer/proj are already split (diffusers layout, no ComfyUI prefix).
  • The safetensors metadata names the scheme: unsloth_format: unsloth_prequant_transformer_state_dict_v1, unsloth_metadata with quant_backend: torchao, scheme: fp8, fp8_granularity: per_row, plus a per-tensor JSON describing the Float8Tensor (block_size [1, in], PerRow activation quantization).

Likely approach: map the keys onto the scaled-fp8 layout the single-file loaders already read (._weight_qdata -> .weight, ._weight_scale -> .weight_scale). Per-output-channel scales are already supported (expand_weight_scale), so the existing fp8 path (dequantize, or keep fp8 storage) should apply unchanged. The install-time refusal in _raise_for_unsupported_comfy_quantization would then become a conversion.

Open questions:

  • unsloth also publishes Qwen-Image-2.1-INT8.safetensors and an INT8-ConvRot build; their layout was not checked (likely torchao Int8Tensor).
  • Other families: unsloth ships the same format for other models (e.g. a Qwen-Image-2.1 Turbo FP8 and text encoders). Model fingerprints that read quantized Linears' .weight (e.g. Ideogram 4's) would not match such files at all and would need to accept ._weight_qdata too.
  • Verification: compare against the bf16 transformer on the same seed, as done for the ComfyUI fp8/int8 builds.

Alternatives

Keep refusing them (current behaviour) and point users to the bf16, fp8_scaled or int8_convrot ComfyUI builds, or a GGUF.

Additional Content

Found while GPU-testing #9784 / #9787 / #9793 (Qwen-Image-2.1).

Metadata

Metadata

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions