[Feature] INT4 Quantization-Aware Training (QAT) for MoE Expert Linear Layers

[kaimo455]

Summary

Add INT4 Quantization-Aware Training (QAT) support for MoE expert linear layers, enabling INT4 fake quantization/dequantization during training to reduce the accuracy loss when deploying with INT4 inference.

Note: This issue is specifically scoped to INT4 QAT only. Other quantization bit-widths (e.g., INT8, FP8) are out of scope and should be tracked separately if needed.

Motivation

MoE models have a large number of expert linear layers (TEGroupedLinear), which dominate memory and compute. INT4 quantization can significantly reduce the model footprint and improve inference throughput, but post-training quantization (PTQ) often leads to unacceptable accuracy degradation for these sensitive layers. QAT mitigates this by simulating INT4 quantization effects during training, allowing the model to adapt to the lower precision.

Proposed Feature

• INT4 QAT CUDA kernels: Fake quantization and dequantization kernels that simulate INT4 rounding behavior during the forward pass while keeping weights in FP32 for gradient updates.
• Fused fake-quant/dequant kernel: A fused kernel to reduce kernel launch overhead.
• CLI arguments:
  • --enable-int4-qat: Toggle INT4 QAT on/off
  • --int4-qat-group-size: Group size for per-group quantization (default: 32)
  • --int4-qat-filter-regex: Regex filter to selectively apply QAT to specific TEGroupedLinear modules
• Training integration: Hook QAT setup into the pretrain loop in training_utils.py
• Weight transform utilities: Support for weight reshaping and quantization parameter computation.

Acceptance Criteria

• INT4 QAT CUDA kernels pass correctness tests against PyTorch reference
• Fused kernel matches the non-fused version in numerical output
• CLI arguments are wired up and functional
• QAT can be enabled/disabled without affecting normal training
• Dockerfile includes the int4_qat package build step