Skip to content

feat: data-parallel block tuning via --data_parallel - #2351

Open
avtc wants to merge 4 commits into
intel:mainfrom
avtc:feature/ddp-parallel-tuning
Open

avtc wants to merge 4 commits into
intel:mainfrom
avtc:feature/ddp-parallel-tuning

Conversation

@avtc

@avtc avtc commented Sep 11, 2026 •

Copy link
Copy Markdown
Collaborator

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

This comment has been minimized.

@avtc
avtc marked this pull request as draft September 11, 2026 20:43
@avtc
avtc force-pushed the feature/ddp-parallel-tuning branch from 849e278 to 230317f Compare September 11, 2026 22:20
@avtc
avtc marked this pull request as ready for review September 11, 2026 22:34
@AutoRoundBot

Copy link
Copy Markdown
Collaborator

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

@avtc
avtc force-pushed the feature/ddp-parallel-tuning branch from 230317f to b0533cc Compare September 12, 2026 10:37
@AutoRoundBot

Copy link
Copy Markdown
Collaborator

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

@avtc
avtc requested a review from wenhuach21 September 12, 2026 13:22
@wenhuach21 wenhuach21 added this to the 0.16.0 milestone Sep 13, 2026
@wenhuach21

Copy link
Copy Markdown
Contributor

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


run_deferred_wrap_searches(block, None)

def _collect_tuning_params(mod):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

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.

offset += n


def _xchg(seg: torch.Tensor, dev: torch.device, dtype: torch.dtype, transport: str) -> torch.Tensor:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function name is not very clear.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

renamed

Comment thread docs/environments.md Outdated
- **Valid Values**: comma-separated device indices

```bash
export AR_TUNE_DDP_DEVICES=0,1,2,3

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@avtc
avtc force-pushed the feature/ddp-parallel-tuning branch from b0533cc to 837728c Compare September 15, 2026 19:10
@AutoRoundBot

Copy link
Copy Markdown
Collaborator

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

@chensuyue

Copy link
Copy Markdown
Contributor

@yiliu30 @wenhuach21 please review.

@yiliu30

yiliu30 commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

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
avtc marked this pull request as draft September 17, 2026 00:58
@avtc
avtc force-pushed the feature/ddp-parallel-tuning branch from 7ca13fd to 837728c Compare September 17, 2026 22:00
@avtc

avtc commented Sep 17, 2026

Copy link
Copy Markdown
Collaborator Author

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.

@chensuyue chensuyue removed this from the 0.16.0 milestone Sep 24, 2026
@avtc
avtc marked this pull request as ready for review September 25, 2026 22:32
avtc and others added 2 commits October 24, 2026 03:59
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>
@avtc
avtc force-pushed the feature/ddp-parallel-tuning branch from 2dd3132 to 5dc253f Compare October 24, 2026 01:09
@AutoRoundBot

Copy link
Copy Markdown
Collaborator

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

Comment thread docs/environments_CN.md Outdated
- **用法**:在弱 CPU 主机上限制并发分片前向线程数。

```bash
AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES=2 python -m auto_round --model ... --parallel_quantization 8

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

renamed to --data_parallel

Comment thread docs/environments_CN.md
5. **自定义工作目录**:将 `AR_WORK_SPACE` 设置为磁盘空间充足的目录

### AR_TUNE_DDP_MAX_COLLECT_FORWARD_DEVICES
- **描述**:限制携带 hook 的分片采集前向并发使用的设备数(hook 前向会使编译后的执行器在图分段之间退回 Python 绑定,可能在 GIL 上形成串行瓶颈)。仅在 hook 前向出现串行瓶颈的主机上才需要调低;在验证机台上至 world=8 未复现。大于已启用 world 的值等效于不限制。

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@avtc avtc changed the title feat: data-parallel block tuning via --parallel_quantization feat: data-parallel block tuning via --data_parallel Oct 27, 2026
@wenhuach21
wenhuach21 requested a balanced review from Copilot October 27, 2026 09:16

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread auto_round/algorithms/block_runner.py Resolved
Comment thread auto_round/algorithms/block_runner.py Resolved
Comment thread auto_round/compressors/utils.py Resolved
Comment thread auto_round/algorithms/block_runner.py Resolved
- 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>
@avtc

avtc commented Oct 27, 2026

Copy link
Copy Markdown
Collaborator Author

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

@avtc

avtc commented Oct 27, 2026

Copy link
Copy Markdown
Collaborator Author

Copilot findings fixed/resolved

@avtc
avtc marked this pull request as draft October 27, 2026 16:55
@avtc
avtc marked this pull request as ready for review October 27, 2026 16:56
@azure-pipelines

Copy link
Copy Markdown
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
wenhuach21 self-requested a review October 28, 2026 01:22

@wenhuach21 wenhuach21 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@wenhuach21

Copy link
Copy Markdown
Contributor

please have a review

@wenhuach21

Copy link
Copy Markdown
Contributor

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!

@avtc
avtc marked this pull request as draft October 28, 2026 14:10
@avtc
avtc marked this pull request as ready for review October 28, 2026 14:10
@AutoRoundBot

Copy link
Copy Markdown
Collaborator

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

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

@avtc

avtc commented Oct 28, 2026

Copy link
Copy Markdown
Collaborator Author

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

@chensuyue

Copy link
Copy Markdown
Contributor

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

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

@avtc

avtc commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants