Migrate Diffusion Quantization to Pre-Export - #286
kevchengcodes wants to merge 4 commits into
Conversation
…re uncompressed by default.
…are quantized at the torch level only.
| _INT4_PER_BLOCK32 = { | ||
| "dtype": "int4", | ||
| "qscheme": "symmetric_with_clipping", | ||
| "granularity": {"type": "per_block", "block_size": 32, "axis": 1}, |
There was a problem hiding this comment.
coreai-opt automatically picks the right weight quantization axis for per-channel and per-block, for the supported layers. This need not be explicitly specified.
| "granularity": {"type": "per_block", "block_size": 32, "axis": 1}, | |
| "granularity": {"type": "per_block", "block_size": 32}, |
| _INT8_PER_CHANNEL = { | ||
| "dtype": "int8", | ||
| "qscheme": "symmetric_with_clipping", | ||
| "granularity": {"type": "per_channel", "axis": 0}, |
There was a problem hiding this comment.
same for per-channel weights
| "granularity": {"type": "per_channel", "axis": 0}, | |
| "granularity": {"type": "per_channel"}, |
| @pytest.mark.parametrize("name", [n for n, p in PRESETS.items() if p["config"] is not None]) | ||
| def test_presets_satisfy_the_coreai_opt_schema(name: str) -> None: | ||
| """Catches a schema typo without downloading any weights.""" | ||
| coreai_opt_quantization = pytest.importorskip("coreai_opt.quantization") |
There was a problem hiding this comment.
I don't think we should skip this test if we cannot import coreai_opt.quantization. Its a big issue if we cannot import coreai_opt.quantization and should be raised back to the developer
| TextEncoderWrapper, | ||
| quant_weight_owner, | ||
| ) | ||
| from coreai_models.diffusion.pipeline import _quantize_component_weights, _resolve_compression |
There was a problem hiding this comment.
The presets are meant to be constant templates, but exporting mutates them -- we should do return copy.deepcopy(config) in _resolve_compression (or deepcopy in _quantize_component_weights before passing) - this is what you already do in the test also :)
6 participants
Summary
All diffusion models (FLUX2, SD-1.5,2.1,3.5, and WAN) currently would quantize weights AFTER export at the MLIR level. This is not the preferred path for
coreai-opt. This PR simply switches the weight quantization to BEFORE export at the Torch module level.We see near-parity or better with this switch, and significantly quicker export times.
Changes
apply_mlir_quantizationquantization_configTesting
FLUX2 was the only diffusion model that is 4bit compressed by default. The other 4 models were uncompressed by default. However, all models produced reasonably identical results to their respective baseline version on the main branch.