feat: data-parallel block tuning via --parallel_quantization - #2351
Open
avtc wants to merge 1 commit into
Open
Conversation
This comment has been minimized.
This comment has been minimized.
avtc
marked this pull request as draft
September 11, 2026 20:43
avtc
force-pushed
the
feature/ddp-parallel-tuning
branch
from
September 11, 2026 22:20
849e278 to
230317f
Compare
avtc
marked this pull request as ready for review
September 11, 2026 22:34
Collaborator
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
Single-process data-parallel tuning for block-wise quantization behind one flag: --parallel_quantization {off,auto,N} (default off). auto/N>=2 pin the replica world to the largest power-of-two <= visible CUDA devices, subject to a free-VRAM mirror guard.
- each block is mirrored onto the plan devices; replicas tune on disjoint calibration shards, so the effective batch matches the serial run
- gradients exchanged via bf16 transport + sign-cast for sign-SGD steps (bitwise-identical across replicas); in-process halving-doubling allreduce, no torchrun
- collection forwards, iters=0 RTN/OptRTN searches, and V2 wrap-time init-scale searches shard across the replicas too; AWQ/GPTQ not yet adapted and run serial
- losses: the wrapped starting state is bit-identical to serial (checksum-verified across independent runs);
per-iteration loss lines differ because replicas draw from disjoint calibration shards (reproducible per mode
for a fixed seed); end-to-end quality at parity: iters=50 KL 0.021815 (world=4) vs 0.02203 serial, within
run noise
- explicit N raises with the blocking reasons if it cannot engage; auto shrinks to the devices that fit
- each replica holds a full block mirror: ~12-13 GiB/GPU for 27B-class blocks (W4A16 g128, calib 8x8192) at world=4; the VRAM guard declines oversized blocks with per-device reasons
Also: transformers-v5 fixes in the block runner (shared-cache rope kwargs arrive as per-batch (cos,sin) tuples, sliced per shard; inputs staged on the replica device); guarded excepts are fail-visible; AR_PERF_COUNTERS adds opt-in perf lines.
Knobs: AR_TUNE_DDP_DEVICES (docs/environments.md). Usage docs in docs/step_by_step.md + CN.
Tests: test/unit/test_cpu/quantization/test_ddp_core.py, test/unit/test_cuda/quantization/test_ddp_mirror.py.
Independent of the pre-existing multi-process DDP integration for the LLM-Compressor workflow (intel#1897): that one is inert from the plain CLI; this engine is single-process and additionally shards the collection passes and the zero-shot/wrap-time searches.
CI fixes (build 75613, unit suite on the previous push): the tune-loop
surgery had dropped the upstream loop-tail pair that records the iter-0
loss, so the end-of-block summary formatted None and failed 93 iters>0
tests; the aborted-block leftovers then cascaded into the four
WrapperLinear.linear_forward failures when the calibration restore sweep
met still-wrapped modules. Restores the pair verbatim, adds an AST
contract test pinning the assignment inside the loop, and removes a
duplicate P2P ordering test flagged by ruff.
P2P enablement removed after a measured no-delta A/B on the validation
rig (27B, iters=10, ddp4): per-block walls 18.4/12.6/12.2s with explicit
enablement vs 18.9/12.6/12.0s with it skipped, gradient-exchange means
sub-20ms both ways -- the fabric's cross-device copies ride the memcpy
path regardless, so explicit peer-access setup bought nothing. The
AR_TUNE_DISABLE_P2P knob, the cudart binding/signature machinery, and
their tests are gone with it.
Signed-off-by: avtc <tarasenkov@gmail.com>
avtc
force-pushed
the
feature/ddp-parallel-tuning
branch
from
September 12, 2026 10:37
230317f to
b0533cc
Compare
Collaborator
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
wenhuach21
requested review from
lvliang-intel,
n1ck-guo,
xin3he and
yiliu30
September 13, 2026 12:03
Contributor
|
Thanks for this PR. This is exactly what we need right now. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds a single-process data-parallel tuning engine behind one CLI flag:
--parallel_quantization {off, auto, N}.off(default) keeps the serial single-GPU loop unchanged;auto/N>=2pin the replica world to the largest power-of-two <= visible CUDA devices, subject to a free-VRAM mirror guard. On multi-GPU boxes the idle GPUs now do the tuning work.How it works (
auto_round/algorithms/quantization/sign_round/data_parallel.py):iters=0RTN/optimized-RTN searches, and the SignRoundV2 wrap-time init-scale searches all shard across the replicas too. AWQ scale search, GPTQ, and other calibration-fitting optimizers are not yet adapted and run serial, exactly as today.Nraises with the blocking reasons if it cannot engage;autoshrinks to the devices that fit.Also: Transformers-v5 block-runner fixes (shared-cache rope kwargs arrive as per-batch
(cos, sin)tuples, sliced per shard; inputs staged on the replica device), fail-visible CUDA peer-access enablement, opt-inAR_PERF_COUNTERSdiagnostics.Measured — Qwen3.8-27B (hybrid linear/full attention), W4A16 g128, iters=50, same seed and recipe, 4x RTX 3090:
--parallel_quantization 4KL protocol: full-vocab fp32 KL vs the bf16 base, openwebtext 8x8192, same samples/seeds for both arms. Delta is within run noise.
At
iters=0(zero-shot optimized-RTN): 64-block run 13:00 -> 6:55 (1.9x; collection 3.5 -> 1.2 s, searches 5.2 -> 2.0 s per block, world=4). Atiters>0the wrap-time init-scale search drops 5.1 -> 1.9 s per block (~2.7x, GIL-limited) and per-replica compile lowers tune forward 72 -> 42 ms/iter.Knobs:
AR_TUNE_DDP_DEVICES,AR_PERF_COUNTERS(docs/environments.md).Type of Change
New feature
Related Issues
Implements #2010 (multi-GPU quant parallelism), single-process flavor.
Independent of the pre-existing multi-process DDP integration for the LLM-Compressor workflow (most recently fixed in #1897): it is inert without an external launcher; this engine runs from the plain CLI/API and additionally shards the collection and search passes. The two lanes were not benchmarked against each other.
Checklist Before Submitting
/azp run Unit-Test-CUDA-AutoRound.