Enable int4 quantization for parakeet

[larryliu0820]

Summary

• Add int4/int8 quantization support for Parakeet TDT model export using torchao
• Extract quantization utilities to a separate module for reusability

Changes

Quantization Support for Parakeet

Added support for quantizing encoder and decoder components with multiple configurations:

• Linear layers: 4w, 8w, 8da4w, 8da8w quantization configs
• Embedding layers: 4w, 8w quantization configs
• Packing formats: tile_packed_to_4d for optimized inference on CUDA

New CLI Arguments

Argument	Description
--qlinear_encoder	Quantization config for encoder linear layers
--qlinear_encoder_group_size	Group size for encoder quantization (default: 32)
--qlinear_encoder_packing_format	Packing format for encoder
--qlinear	Quantization config for decoder linear layers
--qlinear_group_size	Group size for decoder quantization (default: 32)
--qlinear_packing_format	Packing format for decoder
--qembedding	Quantization config for embedding layer
--qembedding_group_size	Group size for embedding quantization

Code Organization

• Extracted quantize_model_() function to examples/models/parakeet/quantize.py (mirrors optimum-executorch naming)
• Model moved to CUDA after preprocessor export when --backend cuda is specified
• Example inputs created on correct device to match model parameters

Example Usage

Int4 Linear with Tile Packing

python examples/models/parakeet/export_parakeet_tdt.py \
    --backend cuda \
    --dtype bf16 \
    --qlinear_encoder 4w \
    --qlinear_encoder_packing_format tile_packed_to_4d \
    --qlinear 4w \
    --qlinear_packing_format tile_packed_to_4d \
    --output-dir ./parakeet_int4

Int4 Linear + Int8 Embedding

python examples/models/parakeet/export_parakeet_tdt.py \
    --backend cuda \
    --dtype bf16 \
    --qlinear_encoder 4w \
    --qlinear_encoder_packing_format tile_packed_to_4d \
    --qlinear 4w \
    --qlinear_packing_format tile_packed_to_4d \
    --qembedding 8w \
    --output-dir ./parakeet_int4_emb8

Test Plan

• Export with CUDA backend and int4 quantization completes successfully
• Model runs through encoder with storage_offset tensors
• Verify full transcription accuracy matches eager mode
• Verify model size reduction with quantization

Model Sizes

bfloat16 (baseline)

-rw-r--r--. 1 dev dev 1274592908 Jan 23 07:01 aoti_cuda_blob.ptd
-rw-r--r--. 1 dev dev    4497576 Jan 23 07:01 model.pte

Encoder and decoder int4 group (32) wise quantized

-rw-r--r--. 1 dev dev 542162572 Jan 30 07:03 aoti_cuda_blob.ptd
-rw-r--r--. 1 dev dev   4125496 Jan 30 07:03 model.pte

Int4 linear + Int8 embedding

-rw-r--r--.  1 dev dev 536943756 Feb  2 18:21 aoti_cuda_blob.ptd
-rw-r--r--.  1 dev dev   4134072 Feb  2 18:21 model.pte

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/17061

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures, 6 Pending, 1 Unrelated Failure

As of commit 89808f9 with merge base 9d257c8 ():

NEW FAILURES - The following jobs have failed:

• pull / unittest-buck / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 91bb5996646dda3b886a3f7ffed6aa5448b270add1ec57f5da3750bb5382bc95 /exec failed with exit code 3
• trunk / test-arm-backend-ethos-u (test_model_smollm2-135M) / linux-job (gh)
  RuntimeError: Command docker exec -t 39bb5c89ad2a305ddb6fe4a8df89a7d6318c83182c9bda9ccb005dbfc3bbe275 /exec failed with exit code 2
• trunk / test-llama-runner-linux (bf16, portable, linux.arm64.2xlarge, executorch-ubuntu-22.04-gcc11-aarch64) / linux-job (gh)
  RuntimeError: Command docker exec -t 30716aa502558e5fd4d52ef94f3f597915b9c871eb3398cac96e76e511c94f68 /exec failed with exit code 35
• trunk / test-qnn-model (fp32, mv2) / linux-job (gh)
  RuntimeError: Command docker exec -t 07658f01ecc107ceea4f5907ba72b9032f31da03bb85ed73998479d841383b64 /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest-buck / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 3

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[mergennachin]

Also please report size reductions

• encoder quantized (4w, groupwise)
• encoder and decoder quantized (4w, groupwise)
• encoder, decoder (4w, groupwise) and embedding quantized (8w, per-channel)