[avtc]

Description

Adds data-parallel block tuning behind one CLI flag, --data_parallel {off, auto, N}, built as a DDP-style wrapper: the engine runs the algorithm's unchanged per-block tuning loop once per GPU and takes care of the module, the data and the gradient sync. A rank is a thread of the same process (one per GPU), so the model and the calibration activations stay in one address space; the algorithm sees the same contract as on the multi-process torchrun path.

off (default) keeps the serial loop. auto takes all visible CUDA devices at iters=0 and the largest power of two that fits them at iters>0 (announced in the log). An explicit N is checked at startup: between 2 and the visible device count, and a power of two when iters>0 (the gradient all-reduce is a halving-doubling exchange).

How it works

Engine: auto_round/algorithms/parallel/ (spmd.py ranks and collectives, data_parallel.py planning, block copies and sharded calibration forwards, tune_parallel.py, mirror_sync.py).

Per block, when the tune loop runs (iters>0):

• Module. Rank 0 tunes the home block; every other rank tunes a copy of the block on its own GPU. After the loop, rank 0's block goes through the normal unwrap and pack path; the copies are dropped.
• Data. The calibration pool is split into contiguous per-rank shards placed on each rank's GPU. make_index_sampler() gives each rank a sampler over its own shard (own RNG per rank).
• Optimizer. Each rank builds its own optimizer from its replica's parameters, inside the unchanged loop.
• Gradients. sync_gradients() (the hook the loop already calls for the torchrun path) averages the gradients across ranks before the optimizer step: an in-process fp32 halving-doubling all-reduce. AdamRound keeps one GradScaler per rank; the gradients are averaged before scaler.step, so all ranks see the same inf/nan checks.
• Loss. all_reduce_loss(loss, weight) gives every rank the global loss, weighted by each rank's normalizer (element or valid-token count), so best-iteration selection and early stopping agree across ranks.
• Compile. The first torch.compile of each block shape runs rank by rank (first-time compile is not thread-safe); later blocks of the same shape run concurrently.

For every calibration pass (also at iters=0):

• The no-grad calibration forwards (reference outputs, the quantized-input chain, imatrix and activation statistics) are sharded across the GPUs. Quantizers declare the statistics their hooks write and how they reduce (calibration_stats = {"imatrix": "sum"}); the engine folds exactly those from the copies. A hook pass with undeclared statistics stays serial, so statistics are never silently partial.
• SignRoundV2's wrap-time init-scale searches are split across the ranks (each rank searches a share of the layers; the scales are exchanged). At iters=0 the per-layer searches use main's existing cross-device search.

A block that cannot run on ranks tunes serially and logs the reason: the sample count or batch does not split evenly over the ranks, or the per-rank batch exceeds its shard. LFQ, dict-shaped (diffusion-style) calibration pools and an active torchrun group are not supported. The block must sit whole on one GPU; a block spread across GPUs stops at startup with that reason. Free-VRAM estimates only warn (naming the device); an actual shortage surfaces as an out-of-memory error.

--parallel_batch {split, per_replica}: split (default) keeps --batch_size as the global batch of an iteration (each rank draws batch_size / world); per_replica lets each rank draw the full --batch_size.

Changes in algorithm code

The tuning loop itself is unchanged; the quantizer-side diff is:

• sign_round/quantizer.py: IndexSampler(...) -> make_index_sampler(...); all_reduce_loss(total_loss, num_elm) after the loss accumulation; the masked-loss fix below. AdamRound and SignRoundV2 inherit the loop: no changes.
• calibration_stats declarations: OptimizedRTNQuantizer (imatrix, imatrix_cnt), SignRoundV2Quantizer (imatrix).
• BaseQuantizer: the calibration_stats attribute and parallel_constraints() (LFQ layers are single-device).

All three hooks are identity calls in a serial run.

Serial-path change: masked loss normalization. With a valid-token mask and gradient_accumulate_steps == 1, SignRound's loss was a mean over all elements, then divided again by the valid-token count (normalized twice, and not additive across ranks: each rank printed about world x the serial value). It is now a sum over the valid elements divided by the valid-token count, as the accumulation path already does. For serial runs the printed masked loss changes scale. SignRound's tuned weights do not change (sign-based update), and best-iteration selection does not change for equal-length samples (the factor is constant). For AdamRound the gradient scale changes by that constant factor; Adam's update is invariant to it except for eps.

Validation

Qwen3.8-27B (hybrid linear/full attention), W4A16 g128, iters=10, 8x RTX 3090 (24 GB):

run	wall	KL mean vs bf16
serial, main	33.4 min	0.025723
serial, this PR (previous revision; same serial path for unmasked data)	33.4 min	0.025666
--data_parallel 4	16.9 min	0.026111
--data_parallel 4, earlier revisions of this engine	16.7 min / 18.6 min (another job started on two of the GPUs mid-run)	0.026110 / 0.026250

KL protocol: full-vocab KL against the bf16 base, openwebtext 8x8192, same samples for all runs; re-measuring the same output gives the same value.

Quality: data-parallel mean KL is 1.5-2.3% higher than serial in these runs (1.7% for the final revision against this PR's serial run); the two serial runs differ by 0.2%. Per-block tuning losses of the parallel runs match serial on the early blocks and are 9-10% higher on the last 8 blocks. With identical samples the ranks reproduce the serial weights bit-exactly (CPU, and 2 GPUs in the CUDA parity test), so the gap comes from the sample selection: each rank draws only from its own shard of the calibration pool, so iteration batches differ from serial. (The sharded calibration statistics differ from serial only in summation order; at iters=0 the KL stays within 0.15% of serial.) The previous in-algorithm DP implementation of this PR showed the same gap (0.026063 at 4 GPUs, 0.026007 at 8).

iters=0 (optimized RTN), same model, 4 GPUs (--device_map 0,1,2,3): serial 8.3 min, --data_parallel 4 7.1 min (sharded calibration forwards; the searches use main's cross-device search).

AdamRound (iters=10, eager, first two blocks): serial ~25-27 s/block, 4 ranks 13.7 s/block (19.2 s on the first block, compile turnstile), similar loss.

CPU tests: SignRound, SignRoundV2 (with the rank-split wrap searches) and AdamRound on 2 CPU ranks produce weights bit-identical to the serial run when the ranks see the serial samples; the printed masked loss on ranks equals the serial value.

Known limitations

• Blocks that do not fit one GPU together with the tuning state are not covered. Example: Hunyuan-3 (MoE) at iters>0 needs ~35 GB per block copy on 24 GB cards; it runs serial with a multi-device --device_map. A pipeline+data-parallel variant is a follow-up.
• AdamRound with torch.compile under ranks failed on GPU (CUDA error: invalid argument during triton autotuning in the backward); eager AdamRound works. (Serial AdamRound with compile runs out of memory on these 24 GB cards at batch 8.)
• LFQ is single-device.

Knobs

--data_parallel {off, auto, N}, --parallel_batch {split, per_replica}, AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES (docs/environments.md).

Type of Change

New feature

Related Issues

Implements #2010 (multi-GPU quant parallelism), single-process flavor. RFC: #2378.

The multi-process DDP integration for the LLM-Compressor workflow (most recently fixed in #1897) uses the same algorithm hooks (setup_ddp_if_needed_, sync_gradients()); a run uses one lane: with a torchrun group active, --data_parallel stops with that reason.

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!

[avtc]

Thanks! Agreed on the goal - the context is shaped for the algorithm set. Concretely:

# engine side (orchestrator/composer owns all distribution):
accel = TuneParallelContext.create(model_ctx, devices)   # mirrors, threads, transport

# gradient algorithms (SignRound family loop, or a custom loop):
batches = accel.shards(sampler)                          # ~ DistributedSampler(dataset)
loss    = accel.run_step(step, batches)                  # fan out step(batch) to replicas
accel.sync_grads()                                       # ~ autograd-hook allreduce
accel.step(opt_step)
accel.teardown()

# no-grad searches (AWQ-style grids, Hessian accumulation):
stats = accel.reduce(partial_fn, batches, op="sum")

Adoption per algorithm kind:

• the gradient algorithms already share one inherited tune loop - V1 defines it, V2 and AdamRound subclass it with their own loss/optimizer closures, diffusion rides the same engine - so the context lands in that one loop and every subclass gets data parallelism with zero changes;
• collection (fp reference outputs, quantized-input chain, imatrix/act-max stats) is already sharded today at the composer level for every algorithm using the shared capture path (RTN included) - it moves under the same context unchanged;
• no-grad searches map to the shard-and-reduce seam above (AWQ's grid, a future GPTQ Hessian);
• a genuinely custom loop (e.g. rotation training) writes the guarded calls directly.

All guarded calls are no-ops when accel is None (serial). One deliberate divergence from the tutorial: replicas are worker threads in one process, not one process per GPU - the orchestrator keeps streaming offload, block placement and memory accounting in a single address space; the multi-process flavor is already available on main, and this PR adds the in-process axis.

One seam stays explicit - gradient sync: the tutorial's invisible sync rides autograd hooks on a single backward into NCCL allreduce; here one iteration runs N replica backwards in worker threads, and the consensus is a sign-cast exchange that runs once all replica grads are ready. Making that hook-driven needs per-parameter readiness coordination across threads and pushes algorithm properties into the engine (sign-cast is exact only when momentum == 0) - so it stays one explicit guarded call.

Scope for this PR: the context, the shared-loop adoption (SignRound family) and the collection seam. AWQ/rotation adoption onto the reduce/explicit-call seams can be added later.

I am starting working on this.

[avtc]

The TuneParallelContext split discussed above has landed (single commit, rebased on latest main). Algorithm-side footprint is now one step closure plus None-guarded calls; collection sharding moved under the same context. Added an e2e test (real loop on CPU mirrors, consensus bit-identity) and a serial-vs-dp per-iteration loss-parity test with pinned draws. Server-validated: ddp4 iters=0 and iters=10 matched pre-splt.

[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.

[avtc]

renamed

[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.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[chensuyue]

@yiliu30 @wenhuach21 please review.

[yiliu30]

Hi @avtc, thanks for the PR! I'd like to understand the target scope a bit better first.

The tuning framework supports multiple devices within a single block. Are you targeting DDP + multi-device tuning per rank, or DDP + single-device tuning per rank?

Also, since this is a fairly significant feature that introduces substantial infrastructure changes, it might be helpful to start with an RFC that outlines the high-level design, intended use cases, expected workflow, potential infrastructure changes, and associated risks. This would give the broader team a chance to review the proposal early, provide feedback, and help ensure we're aligned on the overall direction before diving into the implementation details.

FYI, we did this when we enabled such feature for LLMC vllm-project/llm-compressor#2180

[avtc]

Hi @yiliu30, thanks for taking a look,

Scope: DDP + single-device tuning per rank.

Each replica, including the home copy, holds the whole block on one GPU for collection and tuning. SignRound tunes a block as one unit - every iteration runs the block forward and backward on the current shard and produces a gradient for every tuning parameter, because the reconstruction loss compares the whole block output against the cached reference. The design therefore duplicates the parameters and splits the samples. The same replicas also shard the no-grad collection passes (fp reference and quantized-input chain) and both search phases (iters=0 RTN/optimized-RTN and the wrap-time init-scale searches), so SignRound V1/V2 tuning and those searches all run under the flag.

Multi-device tuning per rank we explored and measured too (branches published in this thread in message #2351 (comment)); it is parked - the RFC's use-cases section covers where that line stands.

The RFC is up as #2378 - it follows the structure of vllm-project/llm-compressor#2180 (high-level design, use cases, workflow, infrastructure changes, risks).

One note on what is ready to follow: the RFC's "Adopting other algorithms" section describes how the AWQ smoothing and clip searches adopt the engine's shard-and-reduce seam; that adoption, gradient-accumulation support inside the lane, and non-CUDA accelerator eligibility (xpu/hpu) are implemented and validated locally and will land as a follow-up PR once this core merges, so the review can treat this PR as the base it builds on.

[wenhuach21]

please have a review, i basically have not read the code and the current version looks good to me as minimal changes is made in algoirhtm code

@wenhuach21 The Unit-Test-AutoRound failures on the latest push look environmental rather than from this branch. Evidence: the same job on this branch was green on Oct 24 (424/0/5 and 42/0/0) and red on Oct 27 (415/9/5, 41/1/0); the only change between the two commits is a flag rename the failing tests never touch. The pip freeze of the two runs differs by transformers 5.20.2 → 5.21.0, torch identical (2.14.0+cpu). All 9 failures are test_omni_model.py ('Qwen3OmniMoeThinkerModel' object has no attribute 'layers'), matching the transformers 5.21 Qwen3OmniMoE refactor; the llmc failure is an int8 view stride error from the same window. Other open PRs failed the same pipeline the same morning before this push, and the same tests pass on this branch locally with transformers 5.15. Happy to help pin or adapt if useful.

Got it, we will follow that. Thanks for the pr again!

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[avtc]

Cuda tests failed with:
❌ Timeout: Agent failed to come online within 1200 seconds.

[chensuyue]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[avtc]

Please let me know when need to rebase next time - for example when you will be ready to merge - i will do it promptly

[wenhuach21]

Please let me know when need to rebase next time - for example when you will be ready to merge - i will do it promptly

Could you please rebase it? I will merge it.

One last question: Could you confirm that your PR implements the following expected behavior? I haven't reviewed the code in detail.

When DDP is enabled, each rank should receive a different subset of the dataset through the dataloader. At the end of each epoch, the dataset should be reshuffled and repartitioned so that each rank gets a new subset of samples for the next epoch.

[wenhuach21]

After taking another look, I think your implementation is correct. However, I would appreciate it if you could double-check it as well.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[avtc]

@wenhuach21
Double-checked the dataloader behavior in code and in the CPU parity tests (test_spmd_ranks.py):

• Each rank receives a different subset: the calibration pool is split into disjoint contiguous shards, one per device, and each rank's sampler reads only its own shard (index_sampler + contiguous_shard_bounds).
• At the end of each epoch, each rank's sampler reshuffles its shard and starts over, dropping the incomplete batch, with a per-rank seed so concurrent reshuffles stay reproducible.
• The shard assignment stays fixed for the tuning run: each device owns its shard's data, so each rank draws from the same subset across epochs. The synced gradient step makes each iteration equivalent to a serial step over the union of the ranks' draws (with --parallel_batch split the global batch matches the serial batch size). This rank-constrained selection is the source of the small KL gap reported earlier; with identical draws the ranks reproduce the serial weights bit-exactly, which is how the parity tests pin it.

[avtc]

cuda tests timed out, and it's strange it still shows running there. could someone run /azp run Unit-Test-CUDA-AutoRound again or check the CI?

[wenhuach21]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[wenhuach21]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[wenhuach21]

@chensuyue @XuehaoSun could you help check the ut env?

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).