Repository navigation
Conversation
This comment has been minimized.
This comment has been minimized.
849e278 to
230317f
Compare
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
230317f to
b0533cc
Compare
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
Thanks for this PR. This is exactly what we need right now. |
|
|
||
| run_deferred_wrap_searches(block, None) | ||
|
|
||
| def _collect_tuning_params(mod): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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!
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
|
My major concern is the amount of changes required in the algorithm code. I wonder if we can learn from how |
| offset += n | ||
|
|
||
|
|
||
| def _xchg(seg: torch.Tensor, dev: torch.device, dtype: torch.dtype, transport: str) -> torch.Tensor: |
There was a problem hiding this comment.
This function name is not very clear.
There was a problem hiding this comment.
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.
| - **Valid Values**: comma-separated device indices | ||
|
|
||
| ```bash | ||
| export AR_TUNE_DDP_DEVICES=0,1,2,3 |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
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.
b0533cc to
837728c
Compare
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
@yiliu30 @wenhuach21 please review. |
|
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 |
7ca13fd to
837728c
Compare
|
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. |
Adds a DDP-style wrapper that runs the algorithm's unchanged per-block
tuning loop once per GPU. A rank is a thread of the same process, 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.
The engine (auto_round/algorithms/parallel/) owns:
- module: rank 0 tunes the home block, every other rank a copy of the
block on its own GPU; after the loop the copies are dropped and the
home block goes through the normal unwrap/pack path;
- data: contiguous per-rank shards of the calibration pool on the ranks'
GPUs; make_index_sampler() returns a sampler over the rank's shard;
- gradients: sync_gradients() (already called for torchrun) averages
them with an in-process fp32 halving-doubling all-reduce; AdamRound
keeps a GradScaler per rank and syncs before scaler.step;
- loss: all_reduce_loss(loss, weight) gives every rank the global loss,
so best-iteration selection and early stopping agree.
The first torch.compile of each block shape runs rank by rank; later
blocks of the same shape run concurrently. A failure in one rank aborts
the others' collectives and is re-raised.
The no-grad calibration forwards are sharded across the GPUs (also at
iters=0); quantizers declare the statistics their hooks write and their
reduction (calibration_stats), and a hook pass with undeclared statistics
stays serial. SignRoundV2's wrap-time init-scale searches are split
across warm ranks.
The search summary lines ([batched-search] wrap search and neuqi
anchors, [rtn-batch]) now count the per-module searches in the
per-device walls and report the phase's elapsed time as device-wall;
under ranks one merged line covers the block.
CLI: --parallel_quantization {off,auto,N}, --parallel_batch
{split,per_replica}. Env: AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES.
Algorithm code: SignRound takes its sampler from make_index_sampler()
and reduces the iteration loss with all_reduce_loss(); OptimizedRTN and
SignRoundV2 declare calibration_stats; BaseQuantizer gains the
calibration_stats attribute and parallel_constraints(). The hooks are
identities in a serial run.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: avtc <tarasenkov@gmail.com>
With a valid-token mask and gradient_accumulate_steps == 1, SignRound's loss was a mean over all elements (masked ones included) and was then divided again by the valid-token count: normalized twice. It is now a sum over the valid elements divided by the valid-token count, as the accumulation path already does. Serial runs: the printed masked loss changes scale. SignRound's tuned weights do not change (the update uses the gradient sign), and best-iteration selection does not change for equal-length samples (the factor is constant). For AdamRound the gradient changes by that constant factor, to which Adam's update is invariant except for eps. Data-parallel ranks: each rank holds about 1/world of the valid tokens, so the doubly normalized loss was not additive across ranks and every rank printed about world x the serial value. With one normalization the token-weighted all_reduce_loss reproduces the serial loss (CPU test: the printed losses of a two-rank run equal the serial run's). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: avtc <tarasenkov@gmail.com>
2dd3132 to
5dc253f
Compare
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
| - **用法**:在弱 CPU 主机上限制并发分片前向线程数。 | ||
|
|
||
| ```bash | ||
| AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES=2 python -m auto_round --model ... --parallel_quantization 8 |
There was a problem hiding this comment.
There are different types of parallelism. The current implementation supports DP, while the original implementation uses naive PP. We may add TP in the future, so I’d suggest using --parallel_type dp with the number of devices specified separately via --devices. Alternatively, --data_parallel 8 might be a bit clearer and more intuitive.
There was a problem hiding this comment.
renamed to --data_parallel
| 5. **自定义工作目录**:将 `AR_WORK_SPACE` 设置为磁盘空间充足的目录 | ||
|
|
||
| ### AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES | ||
| - **描述**:限制携带 hook 的分片采集前向并发使用的设备数(hook 前向会使编译后的执行器在图分段之间退回 Python 绑定,可能在 GIL 上形成串行瓶颈)。仅在 hook 前向出现串行瓶颈的主机上才需要调低;在验证机台上至 world=8 未复现。大于已启用 world 的值等效于不限制。 |
There was a problem hiding this comment.
I haven’t reviewed the code, as I’m no longer focused on this project. If the other reviewers have no further comments, please go ahead, and I’ll help merge the PR once it's ready.
Review feedback: the flag name should name the parallelism type, leaving room for future axes beside the pipeline placement --device_map already provides. Same value space (off | auto | N) and world policy; the flag has not shipped in a release, so the rename is clean with no alias. Drops the internal --ddp_world shorthand from the parser as well.
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Copilot review overview
4 open issues
rows.max()is evaluated (and converted to a Pythonint) even on the hot path, androwscan be… · New Iterating overindices(a tensor that is typically on the inputs’ device) and callingint(i)… · New When any stale accelerator kwarg is detected, nested structures are moved withto_device(...)… · New The warning text says 'batch dim %d' but the value passed issize(the batch dimension length),… · New
What changed in this PR
Adds single-process, in-process “DDP-style” data-parallel block tuning behind --data_parallel, including sharded calibration forwards, in-thread collectives, and rank-aware sampling/loss reduction.
Changes:
- Introduces a thread-rank SPMD runtime, mirror pooling, plan resolution, and sharded no-grad collection forwards for data-parallel tuning.
- Extends quantizers/composer plumbing to declare and merge calibration hook statistics and to reduce iteration loss across ranks.
- Adds CLI flags/docs and adds CPU/CUDA unit tests for parity, mirror planning, and search summary logging.
| File | Description |
|---|---|
| test/unit/test_cuda/quantization/test_ddp_parity.py | Adds CUDA parity tests asserting serial vs in-process DDP bit-exact results. |
| test/unit/test_cuda/quantization/test_ddp_mirror.py | Adds CUDA tests for mirror-plan rejection when a block spans devices. |
| test/unit/test_cpu/quantization/test_tune_parallel_context.py | Adds CPU tests pinning collection-sharding behavior/fallbacks. |
| test/unit/test_cpu/quantization/test_spmd_ranks.py | Adds CPU tests for thread ranks, collectives, samplers, and parity runs. |
| test/unit/test_cpu/quantization/test_search_dispatch.py | Adds tests for the new batched-search summary logging/merge behavior. |
| test/unit/common/compressors/test_reference_output.py | Updates mocks for changed composer hook API (now returns handles + merge map). |
| docs/step_by_step_CN.md | Documents --data_parallel and --parallel_batch behavior/limitations (CN). |
| docs/step_by_step.md | Documents --data_parallel and --parallel_batch behavior/limitations (EN). |
| docs/environments_CN.md | Documents AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES (CN). |
| docs/environments.md | Documents AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES (EN). |
| auto_round/wrapper.py | Shards deferred wrap searches across thread ranks when data-parallel is active. |
| auto_round/utils/distributed.py | Adds in-process rank hooks: sampler factory and weighted loss all-reduce. |
| auto_round/utils/device_manager.py | Adds per-thread device scoping for thread rank execution. |
| auto_round/utils/device.py | Extends dynamo cache bumping and memory monitor to include “extra” devices. |
| auto_round/envs.py | Adds AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES parsing and helper for non-negative ints. |
| auto_round/data_type/utils.py | Avoids large cross-device transfers by relocating imatrix earlier. |
| auto_round/context/compress.py | Adds parallel_policy on compress context for engine policy injection. |
| auto_round/compressors/utils.py | Adds block-forward kwarg device audit + extends IndexSampler for shard-local RNG/indices. |
| auto_round/compressors/orchestrator.py | Adds perf-phase tick reshuffle and logs zero-shot data-parallel no-op info. |
| auto_round/cli/parser.py | Adds --data_parallel and --parallel_batch CLI flags. |
| auto_round/cli/main.py | Resolves/validates --data_parallel into a ParallelPolicy set on compress context. |
| auto_round/algorithms/quantization/sign_roundv2/quantizer.py | Declares calibration stats for sharded hook reduction. |
| auto_round/algorithms/quantization/sign_round/quantizer.py | Uses rank-aware sampler and weighted loss all-reduce; fixes masked-loss normalization. |
| auto_round/algorithms/quantization/search_dispatch.py | Adds SearchSummary + supports summary collection for rank-sharded searching. |
| auto_round/algorithms/quantization/rtn/quantizer.py | Declares imatrix + imatrix_cnt as calibration stats for sharded merges. |
| auto_round/algorithms/quantization/rtn/batched_search.py | Migrates RTN batched search logging to SearchSummary. |
| auto_round/algorithms/quantization/base.py | Adds calibration_stats + parallel_constraints declaration and DP pool pull behavior. |
| auto_round/algorithms/parallel/tune_parallel.py | Implements collection-phase parallel context and sharded collection forward routing. |
| auto_round/algorithms/parallel/spmd.py | Implements thread-ranks, collectives, gradient all-reduce, and rank sharded wrap searches. |
| auto_round/algorithms/parallel/mirror_sync.py | Syncs quantized home state/stats to resident mirrors for sharded post-quant steps. |
| auto_round/algorithms/parallel/data_parallel.py | Implements parallel policy, plan resolution, mirror pool, sharded forwards, and stat merging. |
| auto_round/algorithms/parallel/__init__.py | Adds package marker and module-level description for the parallel engine. |
| auto_round/algorithms/composer.py | Integrates TuneParallelContext, sharded collection passes, pool handoff, and rank tuning entry. |
| auto_round/algorithms/block_runner.py | Makes runner rank-aware for cache device + adds device-safe cat and kwarg slicing fixes. |
Review effort: Lite
- block_runner: hoist the drawn sample ids to host ints once per draw; per-element int() reads on accelerator-resident index tensors synced once per element, at four sites in _select_batch. Device consumers (index_select over the inputs) keep the original index tensor. - block_runner: _take_rows now pins its rows to the host before the range check, restoring the host-side bounds check its docstring documents; only the final index_select moves the rows to the tensor's device. - block_runner: the fallback warning reports the batch size, not the batch dim index. - compressors/utils: the stale-kwarg device sweep relocates only the accelerator tensors the audit flagged; nested host tensors keep their device and identity instead of riding a wholesale to_device(). Signed-off-by: avtc <tarasenkov@gmail.com>
|
@wenhuach21 |
|
Copilot findings fixed/resolved |
|
Azure Pipelines: Successfully started running 5 pipeline(s). 3 pipeline(s) were filtered out due to trigger conditions. There may be pipelines that require an authorized user to comment /azp run to run. |
wenhuach21
left a comment
There was a problem hiding this comment.
I haven't really read through the code, but the current version looks good to me, as it makes only minimal changes to the algorithm code.
|
please have a review |
|
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
Got it, we will follow that. Thanks for the pr again! |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
|
Cuda tests failed with: |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
|
Please let me know when need to rebase next time - for example when you will be ready to merge - i will do it promptly |


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.autotakes all visible CUDA devices at iters=0 and the largest power of two that fits them at iters>0 (announced in the log). An explicitNis 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.pyranks and collectives,data_parallel.pyplanning, block copies and sharded calibration forwards,tune_parallel.py,mirror_sync.py).Per block, when the tune loop runs (iters>0):
make_index_sampler()gives each rank a sampler over its own shard (own RNG per rank).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 oneGradScalerper rank; the gradients are averaged beforescaler.step, so all ranks see the same inf/nan checks.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.torch.compileof 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):
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.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_sizeas the global batch of an iteration (each rank drawsbatch_size / world);per_replicalets 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_statsdeclarations:OptimizedRTNQuantizer(imatrix,imatrix_cnt),SignRoundV2Quantizer(imatrix).BaseQuantizer: thecalibration_statsattribute andparallel_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 foreps.Validation
Qwen3.8-27B (hybrid linear/full attention), W4A16 g128, iters=10, 8x RTX 3090 (24 GB):
--data_parallel 4--data_parallel 4, earlier revisions of this engineKL 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 47.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
--device_map. A pipeline+data-parallel variant is a follow-up.torch.compileunder ranks failed on GPU (CUDA error: invalid argumentduring triton autotuning in the backward); eager AdamRound works. (Serial AdamRound with compile runs out of memory on these 24 GB cards at batch 8.)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_parallelstops with that reason.Checklist Before Submitting
/azp run Unit-Test-CUDA-AutoRound.