[NehalBhandari]

Summary

Adds a public ExportBackend.Torch backend for fast PyTorch inference, e.g. running evals on a compressed model. finalize(backend=ExportBackend.Torch) folds weight compression into dense weights, so the model no longer reconstructs compressed weights on every forward pass. Activation fake quantization is kept, so outputs match the prepared model.

Previously, the only way to evaluate a compressed model in torch was to use the prepared model, which re-applies fake palettization / fake quantization / pruning masks to every weight on every forward pass. The private ExportBackend._TORCH existed but returned the prepared model unchanged.

Changes

• ExportBackend.Torch replaces the private _TORCH. _TORCH stays as a deprecated alias (emits DeprecationWarning), via the existing __deprecated_aliases__ mechanism.
• Eager mode (KMeansPalettizer, MagnitudePruner, eager Quantizer): new fold_compression_parametrizations helper removes the compressor's weight parametrizations with leave_parametrized=True, in eval mode.
  • A weight shared by several modules is folded once; later modules sharing the same original are unparametrized without re-applying the compression.
• Graph mode (Quantizer with execution_mode=graph): new prepare_for_torch_inference writes each fake-quantized weight in place and removes its weight fake quantization node and module.
  • A weight that the graph also reads directly (not only through its fake quantizer) keeps its fake quantization, with a warning, since folding it would change what the direct reader sees.
• Docstrings no longer recommend evaluating the prepared model instead of calling finalize; error messages that suggested ExportBackend._TORCH now suggest ExportBackend.Torch.

Results

ResNet50, ImageNet validation (50,000 images), A100, batch size 256. Forward time is the mean per batch after warm-up, excluding data loading.

┌─────────────────────────────┬───────────────────────────────────┬──────────────────────────┬─────────┐
│           Config            │ Top-1, prepared → finalize(Torch) │ Forward per batch of 256 │ Speedup │
├─────────────────────────────┼───────────────────────────────────┼──────────────────────────┼─────────┤
│ Palettization, 4-bit scalar │ 67.070% → 67.070%                 │ 104.9 → 75.1 ms          │ 1.40×   │
├─────────────────────────────┼───────────────────────────────────┼──────────────────────────┼─────────┤
│ int8 weight-only, eager     │ 76.094% → 76.094%                 │ 120.7 → 75.0 ms          │ 1.61×   │
├─────────────────────────────┼───────────────────────────────────┼──────────────────────────┼─────────┤
│ int8 weight-only, graph     │ 76.084% → 76.076%                 │ 140.0 → 73.8 ms          │ 1.90×   │
└─────────────────────────────┴───────────────────────────────────┴──────────────────────────┴─────────┘

Graph mode already applies conv-BN folding during finalize (before any backend step). Because of it, graph-mode outputs differ from the prepared model at the fp32-rounding level (~1e-6 on logits); the weight fold itself adds no difference.

Qwen3-4B, 4-bit scalar k-means palettization (group size 32), MMLU-Pro 5-shot with chat template, one A100, batch size 8.

┌────────────────────────────────────────────┬───────────────────┬─────────────────────────────────┐
│                                            │ Folded (collapse) │     Prepared (parametrized)     │
├────────────────────────────────────────────┼───────────────────┼─────────────────────────────────┤
│ First 4 batches (25 requests), measured    │ 532 s (8.9 min)   │ 13,304 s (3.7 h) — 25.0× slower │
├────────────────────────────────────────────┼───────────────────┼─────────────────────────────────┤
│ Full MMLU-Pro (12,032 requests), projected │ ~59 h (~2.5 days) │ ~1,490 h (~62 days)             │
└────────────────────────────────────────────┴───────────────────┴─────────────────────────────────┘

[dengqiaoyu]

Posting some current comments, will finish the rest later.

[dengqiaoyu]

Thanks for the changes!