feat: data-parallel block tuning via --parallel_quantization

[avtc]

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>=2 pin 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):

• Each tuning block is mirrored onto the plan devices; replicas tune on disjoint calibration shards, so the effective batch matches the serial run's data coverage.
• Gradients are exchanged via bf16 transport plus a sign-cast exchange for sign-SGD steps (bitwise-identical across replicas); the reduce is an in-process halving-doubling allreduce — no torchrun, no process spawning, so it composes with in-memory block replay.
• The no-grad collection forwards, the iters=0 RTN/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.
• Loss parity: the wrapped starting state is bit-identical to the serial run (checksum-verified) and the per-iteration gradient math is equivalent up to cross-device reduction order and bf16 transport; end-to-end quality is at parity (table below).
• Per-iteration loss lines differ between serial and parallel runs: serial draws each iteration's batch from the full calibration pool while replicas draw from their disjoint shards. Both are reproducible per mode for a fixed seed; compare modes at the artifact level.
• An explicit N raises with the blocking reasons if it cannot engage; auto shrinks to the devices that fit.
• Each replica holds the full block mirror (bf16 weights + fp32 rounding values, gradients, best-MSE snapshot): ~12-13 GiB per GPU for 27B-class blocks (W4A16 g128, calibration 8x8192) at world=4 on 24 GB cards; host RAM peaks ~20-28 GB. Oversized blocks are declined by the VRAM guard with per-device reasons. A memory-sharded variant is future work.

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-in AR_PERF_COUNTERS diagnostics.

Measured — Qwen3.8-27B (hybrid linear/full attention), W4A16 g128, iters=50, same seed and recipe, 4x RTX 3090:

config	KL (protocol below)	per-block tuning wall
serial (1 GPU)	0.02203	~40-47 s
--parallel_quantization 4	0.021815	~25-30 s

KL 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). At iters>0 the 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

• My code has been tested locally.
• Documentation has been updated as needed.
• New or updated tests are included where applicable.
• The CUDA CI has passed. You can trigger it by commenting /azp run Unit-Test-CUDA-AutoRound.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[wenhuach21]

Thanks for this PR. This is exactly what we need right now.

[wenhuach21]

Is there a better way to minimize changes to the algorithm code? Our goal is to make it easy to implement different algorithms while handling the rest in the compressor/composer.
@yiliu30 Could you share how LLMC handles this? We can learn from their approach if it is better.

[avtc]

@wenhuach21

The DDP machinery inside sign_round/quantizer.py today is pure orchestration - plan/eligibility, ReplicaGroup build and warm-up, shard samplers, per-replica step closures, the sync call, threaded step, teardown, counters - none of it algorithm.

Proposal: a composer-owned TuneParallelContext (algorithm-agnostic module beside data_parallel.py; the composer already owns the DDP world). The quantizer keeps serial code plus neutral calls:

  accel = self.tune_parallel        # None -> fully serial
  if accel is None:
      loss = step(sampler.next_batch())
      opt_step()
  else:
      losses = accel.run_step(step)   # step runs on each replica with its shard
      accel.sync_grads()              # sign-cast exchange, momentum-aware
      accel.step(opt_step)
  ...
  accel.teardown()

step(batch) -> loss stays pure algorithm (forward/loss/backward on the current block); the context owns replica execution, mirrors, warm-up, transport, shard samplers, and counters. Same split as LLMC's AutoRound modifier - a DP-unaware algorithm with the exchange delegated to the engine (setup_ddp_if_needed_, #1897) - plus one addition: in-process replicas get no autograd-hook sync, so the engine owns replica execution too, not just the exchange.

With this seam, any gradient-based algorithm gets DP by writing step + an optimizer step; additive state (GPTQ-style Hessian sums) can join later via a sum-reduce on the same hook; data-dependent searches (AWQ clip) stay serial. Net diff effect: quantizer.py shrinks, the new module is mostly moved code. Validation: mechanical extraction pinned by the existing checksum-parity tests (bit-identical wrapped state at iter 0) plus a ddp4-vs-serial KL parity rerun before pushing; CLI/env surface unchanged.

Does this split work for you, or do you see the seams differently? Happy to align first and then extract in one pass.

[wenhuach21]

This looks much better! It would be even better if we could learn from PyTorch DDP. As shown in the DDP example, only the model- and dataset-specific parts need to be changed, while the AutoRound orchestrator could handle the distributed code. This way, algorithm developers do not need to be aware of any DDP logic. However, I don't know the details, so if this is infeasible, let's follow your current design.

BTW, I’ll be on leave for a few days. If you’d like to merge this PR soon, please ask other reviewers to take a look as well.

Thanks for your great effort on this!

[wenhuach21]

My major concern is the amount of changes required in the algorithm code. I wonder if we can learn from how torchrun or similar tools handle this and find a way to decouple the ddp code and move them to compressor/composer to support different algorithms.
For DP, the main logic is simply to run the same code on different CUDA devices and synchronize the gradients, I think it's possible.

[wenhuach21]

This function name is not very clear.

[avtc]

Will rename to _transport_segment - it moves one peer's gradient segment onto the exchange device in the configured transport dtype (bf16 by default; int8 on the sign-exchange path). It will land together with the next change set.

[wenhuach21]

When device_map is auto or a string like "1,2,3,4", we enable parallelization if there are multiple devices. The current behavior is a naive pipeline parallelism implementation, while this PR introduces data parallelism. We may also support tensor parallelism in the future.

If we want to support DP and PP at the same time, we can refer to how lm-eval handles this: https://github.com/eleutherai/lm-evaluation-harness#parallelism-modes

[avtc]

@wenhuach21 thanks for your attention to my PR. Before publishing it I have started to investigate different approaches to unlock parallel quantization of larger models that do not fit single block/layer into single GPU during quantization, and have already implemented PP+DP: https://github.com/avtc/auto-round/tree/feature/pp-dp-tuning (data parallel over pipeline parallel) and Zero-2-lite https://github.com/avtc/auto-round/tree/feature/zero2-lite-tuning (weights are cloned, but everything else is sharded). Also have tried idea of overlapping micro-batches https://github.com/avtc/auto-round/tree/feature/microbatch-pipeline . They all was implemented to the state where they work and produce correct artifacts (KLD measured) at on Qwen3.8-27B.

On my setup 8x3090 + p2p over PCI-e Gen4 x8, I had results for Qwen3.8-27B w4a16g128:

	iters	block time,s	note
Serial, PP2	10	24.5	baseline
PP2 DP2	10	22.7	9% improve
Serial, PP2, Micro-batch 4	10	27.5	12% degrade
Serial, PP2	50	48.4	baseline
Zero2-lite, 4gpu	50	72	48% degrade, bounded by p2p bandwidth on my setup

• Micro-batch 4 - full calibration batch is split into 4 micro-batches pipelined through the K stages (GPipe/1F1B schedule).
• ZeRO-2-lite re-materializes the full fp32 tune state on every replica (twice per iteration - forward + checkpoint recompute - plus grad deposits), so its per-iteration comm traffic scales with the tune state.
• I have thought about tensor parallel idea, but have not investigated it.

So I have stopped further investigations and made current PR live with DP restricted to models that fits single block/layer to single device during quantization (i suppose on beefy GPUs it may work with large models quite well). While I have started to investigate ways to improve serial path performance and ability to quantize models like hy3 (on 4x3090 or 8x3090) with SignRoundV2 and iters>0, and will soon publish a PR regarding that.

[avtc]

Regarding device_map is auto or a string like "1,2,3,4" - in pp-dp branch I have chosen a path to preserve device_map for pipeline parallelism, and mirroring that map to other visible devices for dp.
For example --device_map 1,2 with 6 visible devices will lead to three data parallel groups 1,2 3,4 5,6 preserving map for each group.