[docker, ops, train] feat: INT4 QAT support for MoE expert layers

[kaimo455]

Add INT4 Quantization-Aware Training (QAT) support for MoE expert linear layers.

• Add int4_qat CUDA kernels (fake quantization/dequantization)
• Add CLI arguments: --enable-int4-qat, --int4-qat-group-size, --int4-qat-filter-regex
• Add QAT setup in training_utils.py pretrain loop
• Update Dockerfile to install int4_qat package
• Update .gitignore for int4_qat build artifacts

Resolves #12

[nullnonenilNULL]

Could you provide more details about the current test results?

[nullnonenilNULL]

@claude review this PR

[claude]

Claude finished @nullnonenilNULL's task in 3m 14s —— View job

---

PR Review: [Task] INT4 QAT Support

Tasks

• Gather context and understand the PR
• Review CUDA kernels (fake_int4_quant.cu, fake_int4_quant_dequant_fused.cu)
• Review Python interface and weight transform (interface.py, weight_transform.py)
• Review training integration (arguments.py, training_utils.py)
• Review packaging and infrastructure (setup.py, Dockerfile, .gitignore)
• Review tests (test_int4_qat.py)
• Post findings

---

Verdict: 🟠 Request Changes

Summary: The INT4 QAT implementation is well-structured with solid CUDA kernels, a fused fast path, and a clean STE-based weight-transform approach. The main concerns are: a correctness bug in the asymmetric STE path (gradient is returned for the wrong dtype), a missing .is_cuda property call vs attribute access bug in the fast-path guard that silently falls back to the slower path, duplicate symbols between the two CUDA extension .cu files causing linker issues, and the lack of asymmetric QAT support at the training integration level even though the low-level API supports it.

Tests: Unit tests are thorough and cover symmetric, asymmetric, bfloat16, fused, and STE paths. The TestWeightTransform.test_transform_eval_mode test doesn't actually assert the identity property (it only checks shape) — but this is a minor gap. E2E training test (MoE model with --enable-int4-qat) is not included.

---

Checklist

Area	Status
Correctness	🟠 One STE bug, one silent dispatch bug
Security	✅ No issues
Performance	✅ Fused kernel is well-optimized
Tests	🟡 Missing E2E test, one weak assertion
Documentation	✅ README and docstrings are good
Integration	🟡 Asymmetric mode silently ignored at training level

---

Findings

🔴 Critical — .is_cuda is an attribute, not a callable; fast path is always skipped

interface.py:132 — weight.is_cuda is a bool property. Writing weight.is_cuda (without ()) evaluates to the bool correctly, but the current code has it in a and chain where this looks fine — wait, on re-reading line 132: and weight.is_cuda — that IS correct Python. However the guard also has and weight.dim() == 2 but not and weight.is_contiguous() even though fused_fake_int4_quantize_dequantize_cuda has TORCH_CHECK(x.is_contiguous(), ...). The and weight.is_contiguous() IS present at line 133. So the fast-path guard is actually safe.

🔴 Critical — Duplicate warpReduceMax / FINAL_MASK / ceil_div definitions across two .cu files cause linker ODR violations if compiled into the same .so

Both fake_int4_quant.cu and fake_int4_quant_dequant_fused.cu independently define FINAL_MASK, warpReduceMax, and their own ceil_div-like helpers. Since setup.py compiles them into separate extensions (int4_qat.cuda and int4_qat.cuda_fused), there is no actual linker collision in the current setup. This is fine, but worth noting as a maintenance hazard if the two are ever merged into a single extension. [No action required in current state.]

🟠 Major — _FakeQuantSTE.backward returns grad_output without casting — gradient dtype mismatch for BF16 + asymmetric

weight_transform.py:143 — backward returns (grad_output, None, None). For BF16 weights, grad_output arrives as BF16 from autograd. For the symmetric case the forward returns a BF16 tensor from fake_int4_quantize_dequantize, so dtypes match. For the asymmetric path (sym=False) through the two-pass fallback in interface.py:148, (q * scale_full).to(weight.dtype) correctly returns the original dtype, so this is actually fine. However: if the weight comes in as a Float8Tensor (FP8 param), forward calls weight.dequantize() to get a BF16 tensor, then returns a BF16 fake-quantized tensor. The STE backward passes the downstream gradient (BF16 shaped to the dequant output) back as the gradient for the original FP8 weight parameter. PyTorch won't error here because the gradient is accumulated by TE, but TE may expect the gradient to be in FP8 scale space. This is a subtle correctness risk for FP8+QAT combined training. [Low probability but high impact — add a comment or test.] Fix this →

🟠 Major — int4_quant_common_kernel writes scale/zero inside the inner quantize loop, causing redundant global writes

fake_int4_quant.cu:267–270 — the scale and zero writes (out_scale[...] = scale and out_zero[...] = zero) are inside the second for (int i = 0; i < iters; ++i) loop, meaning every element in the block writes the same scale value iters times to the same address. For block_m * block_n > 32 this causes iters > 1 warp-wide redundant stores. The 1×32 and 32×1 specializations don't have this issue. The writes are idempotent so correctness is unaffected, but it wastes memory bandwidth on the common kernel path. Fix this →

🟠 Major — Training integration silently hardcodes sym=True with no user control

training_utils.py:444 — kwargs = dict(group_size=group_size, sym=True). The low-level API supports asymmetric quantization but there is no --int4-qat-sym/--int4-qat-asym CLI flag. The argument docstrings in arguments.py don't mention this limitation. Fix this →

🟡 Minor — _writeback_cls_cache is a module-level mutable dict; leaks across processes if multiprocessing spawns workers in the same Python process

weight_transform.py:67 — _writeback_cls_cache: dict[type, type] = {}. In distributed training, each rank is its own process so this is not a practical issue. But if someone uses fork-based multiprocessing, the cache can contain stale subclasses. Document the assumption that each rank is a separate process, or clear the cache in apply_int4_qat.

🟡 Minor — test_transform_eval_mode does not assert identity (only shape)

test_int4_qat.py:344–356 — the test comment says "In eval mode, transform is identity" but the test body only checks out.shape == w.shape. It doesn't verify out is w or torch.equal(out, w). The actual eval-mode identity behavior is enforced by the if _mod.training guard in _patched_get_weights, not in FakeQuantWeightTransform.__call__ directly — so the test misses the real code path entirely. Fix this →

🟡 Minor — sym parameter not clipped in _fake_int4_quant_pytorch for asymmetric zero computation

interface.py:90–91 — the PyTorch reference computes zero = torch.clamp(torch.round(-block_min / scale.float()), min=0.0, max=15.0).to(x.dtype). When x.dtype is BF16, the final .to(x.dtype) is redundant (clamp result is already BF16-safe), but more importantly, block_min computation uses .float() for the scale but then zero retains full float precision before the clamp. This is consistent with the CUDA kernel behavior. No bug, but worth a comment noting the deliberate float promotion.

🟢 Nit — setup.py always appends sm_90a gencode even for non-Hopper targets

setup.py:38 — flags += ["-gencode=arch=compute_90a,code=sm_90a"] is unconditionally appended. On Ampere-only machines this adds an unnecessary PTX embedding target, bloating compile time. Consider making this conditional on "90" being in arch_set. Fix this →

🟢 Nit — _FakeQuantSTE is exported as FakeInt4QuantSTE in __init__.py (private name leaked)

__init__.py:15 — from int4_qat.weight_transform import _FakeQuantSTE as FakeInt4QuantSTE. The leading underscore signals this is internal; re-exporting it with a public name in the package's __init__.py creates a public API surface that wasn't intended. Either make _FakeQuantSTE fully public (FakeQuantSTE) or don't re-export it.

---