You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Follow-up ready: AWQ smoothing/clip search adoption of the shard-and-reduce seam, gradient accumulation inside the lane, and non-CUDA accelerator eligibility (xpu/hpu) are implemented and validated locally; they land as a follow-up PR after feat: data-parallel block tuning via --parallel_quantization #2351 merges.
Feature Description
Background
AutoRound tunes models block by block. The calibrator caches each block's inputs and reference outputs; the block's rounding values are tuned against that cache; the tuned block is unwrapped and packed before the next one loads. The sequence holds peak memory at roughly one block, which is the point of it — and it leaves every other GPU idle while the tuning runs. At iters>0 the tune loop dominates the wall time.
The algorithms that shard across replicas in the current implementation are SignRound V1/V2 block tuning and two search phases: the iters=0 RTN/optimized-RTN weight searches and the wrap-time init-scale searches. The AWQ smoothing and clip search and the other calibration-fitting optimizers are untouched by this engine and run serial, exactly as before (AWQ's activation statistics do ride the sharded collection forward — see "Adopting other algorithms"). Diffusion-style input pools use dict-shaped calibration; with the flag set, such a run stops with that reason. Multimodal (VLM) calibration has not been tested with the parallel lane.
High-level design
One iteration of the tune loop, two ranks of N shown:
flowchart LR
A1["Shard 1: fwd + loss + bwd"] --> X(("gradient exchange"))
A2["Shard 2: fwd + loss + bwd"] --> X
X --> W1["Shard 1: update"]
X --> W2["Shard 2: update"]
Loading
Every replica leaves the exchange with the identical averaged gradient, so every replica computes the same update and the mirrors stay identical copies. After the loop ends, the mirrors are dropped and the home block unwraps and packs through the normal path.
Per block, the engine does the following.
First it decides how this block parallelizes: how many replicas (the world), on which devices, and which sample range each replica owns. One function makes that decision for the collection pass, the search phases, and the tune loop, so every phase parallelizes the same blocks the same way. With the flag set, the decision stops the run with the reasons when the configuration is unsupported: the run must use a supported accelerator home device (cuda/xpu/hpu), a single process, and a flat list of calibration samples per block, and runs with a gradient scaler or LFQ are unsupported. The decision also prices each mirror — the block's weights plus its fp32 tuning values — and keeps the devices whose free memory covers that footprint plus a margin; auto derives the world size from the devices that survive. Finally it requires the block to sit whole on the home device: a block whose weights span several devices of the home type stops with a placement error, because whole-block mirrors are the contract of this lane.
The no-grad collection forwards then run shard by shard across short-lived mirrors, each rank consuming its own samples. Statistics that forward hooks write during these passes merge back into the home block — imatrix partial sums add up, activation maxima fold by max — and hook passes with self-contained statistics run on the home alone.
The search phases divide their work across the replicas: each module's search runs once, on one replica, and the result is copied to the same module everywhere. The searches are deterministic given a module's weight and imatrix, so the results match the serial run exactly.
For tuning, the engine builds one persistent mirror per surviving device and hands each rank a contiguous slice of the calibration pool. Every iteration, the engine dispatches the forward, loss, and backward work to the replica threads; gradients are exchanged so that every replica holds the identical average; every replica computes the same update. Loss values are normalized exactly as the serial path normalizes them, so best-iteration selection and the early-stop window behave the same. The loss readback at the end of an iteration waits for the slowest replica — that readback is also the loop's synchronization point. Runs with gradient_accumulate_steps above 1 tune serial through the same gate the other requirements use; sharded accumulation follows in the follow-up PR.
When the loop ends, the engine drops the mirrors, returns the merged state to the home block, and the normal unwrap and pack path takes over.
Gradient exchange. The loop needs every replica to apply the identical update each iteration, so gradients are reduced before the step — in-process, with a launcher-free single process. The exchange choice is algorithm-gated: a pure sign-SGD run exchanges the averaged gradient's signs as int8 (every replica applies bitwise-identical updates), and any other run exchanges the full averaged values in fp32. Either way the update matches what a multi-process run computes.
All parallel machinery lives behind a TuneParallelContext object owned by the composer. The tune loop calls guarded hooks on it — mirror setup, dispatching an iteration's work to the replica threads, gradient sync, teardown — and each hook is a no-op in serial mode, so the serial path keeps its exact current behavior, the same idea as the setup_ddp_if_needed_ seam already released on main.
Adopting other algorithms
The context exposes one seam per algorithm kind, agreed in the PR thread:
Gradient-based algorithms share one inherited tune loop (SignRound V1 defines it; V2 and AdamRound subclass it with their own loss and optimizer closures), so the context lands in that one loop and every subclass parallelizes with zero changes — done in the current implementation. Diffusion rides the same loop in serial mode; its dict-shaped calibration pools stop with that reason under the flag.
Collection passes (fp reference outputs, quantized-input chain, imatrix and activation-max statistics) run through one shared collection forward that shards at the composer level — done in the current implementation. AWQ collects its activation statistics through its own forward hooks that fire during that same shared forward (additive channel sums, so shard contributions total exactly), so the collection side reuses the sharded pass; AWQ's smoothing and clip search run in the block's pre-quantize phase, serial today.
No-grad searches map to a shard-and-reduce call — each replica evaluates its slice on its own copy of the block, the engine sums the partials on the home device, and with fewer than two items the caller keeps its serial loop. AWQ's smoothing grid search is the designed first adopter (its per-point losses are sample sums, so the merged sums reproduce the serial choice); the per-layer clip search follows the same pattern with map semantics. Both are implemented and validated locally, ready as the follow-up PR.
A genuinely custom loop (rotation training, for example) writes the guarded calls directly. Adoption is future work.
Intended use cases
The target setup is a multi-GPU box and a model whose largest block, with the extra space required for quantization, fits one card. The current implementation covers the RTN/OptRTN searches and SignRound V1/V2 tuning. Blocks larger than one accelerator are future work: an exploratory multi-device-per-rank lane (PP+DP replicas spanning devices, with a ZeRO-2-lite variant) was implemented and measured — reference branches are linked in the PR thread — adding data parallelism on top of the pipeline-parallel placement (one replica across two GPUs; PP+DP on two pairs, four GPUs total) bought about 8% over PP alone, so the lane is parked until larger-block cases justify revisiting it.
Expected workflow
# explicit world
auto-round --model <model> --iters 50 --parallel_quantization 4 ...
# or auto-select from free memory and visible devices
auto-round --model <model> --iters 50 --parallel_quantization auto ...
A requested world stops with the reasons whenever engagement fails; auto keeps the devices that fit.
Infrastructure changes and risks
The engine lives in two new modules — sign_round/tune_parallel.py and sign_round/data_parallel.py (mirrors, worker threads, pool sharding, the exchange) — plus guarded hooks in the tune loop and one CLI flag.
The engine runs one Python process. It creates one block mirror per GPU, drives the mirrors from worker threads, and moves gradients between mirrors with device-to-device copies — through the host where peer-to-peer access is unavailable. The single process keeps streaming offload, block placement, and memory accounting in one address space. The trade: the replicas' backward passes share the Python interpreter lock, and each iteration waits for the slowest replica at the loss readback.
Relationship to the multi-process DDP on main
Main carries a multi-process DDP path for the LLM-compressor workflow: torchrun launches one process per rank, and each rank wraps the block in torch's DistributedDataParallel when it holds one GPU, or runs a manual per-parameter all-reduce when its block spans several GPUs. The calibration loader there processes the same batches on every rank, and the all-reduce averages the identical gradients. This engine divides the batches — each GPU owns its own slice — so the averaged gradient equals the serial full-batch gradient while the work is split across GPUs. The engine runs from the plain CLI, and a run uses one lane: with a torchrun group active, the resolver stops with that reason. The two lanes have not been benchmarked against each other.
Motivation and Use Case
AutoRound tunes models block by block on a single GPU while the remaining GPUs of a multi-GPU box sit idle, so the tuning wall grows with model size. #2010 — "[Feature] Support quant parallelism on multi-GPU" — is the open request for exactly this gap (a title-only report; the substance lives in its discussion: early measured data-parallel results and maintainer interest in the direction).
The existing multi-process torchrun DDP integration serves the LLM-Compressor workflow entry point; this lane covers the plain CLI/API path. The proposal: a single-process data-parallel lane behind --parallel_quantization — replicas of the current block on the idle GPUs tune on disjoint calibration shards and exchange gradients once per iteration, with the collection and search passes sharded the same way.
Measured on Qwen3.8-27B at world=4 on 24 GB cards it roughly halves per-block tuning walls at parity KL (protocol and table in #2351). Use cases: any multi-GPU host where serial block tuning dominates the quantization wall.
Alternatives Considered
Serial single-GPU block tuning, the status quo, leaves the remaining GPUs idle and pays the walls measured in #2351. Extending the existing multi-process DDP integration (the LLM-Compressor workflow) to the plain CLI was considered; this RFC chose a single-process in-engine lane, keeping the plain CLI/API path self-contained.
Original request: [Feature] Support quant parallelism on multi-GPU #2010 — "[Feature] Support quant parallelism on multi-GPU" (title-only report; its discussion carries the motivating measurements and maintainer feedback)
Exploratory reference for blocks exceeding one accelerator (multi-device-per-rank PP+DP lane, tested): https://github.com/avtc/auto-round/tree/feature/pp-dp-tuning — parked as future work for that case (PP+DP over PP bought about 8% while using twice the GPUs); it needs a rebase onto the merged core
Feature Description
Background
AutoRound tunes models block by block. The calibrator caches each block's inputs and reference outputs; the block's rounding values are tuned against that cache; the tuned block is unwrapped and packed before the next one loads. The sequence holds peak memory at roughly one block, which is the point of it — and it leaves every other GPU idle while the tuning runs. At iters>0 the tune loop dominates the wall time.
The algorithms that shard across replicas in the current implementation are SignRound V1/V2 block tuning and two search phases: the iters=0 RTN/optimized-RTN weight searches and the wrap-time init-scale searches. The AWQ smoothing and clip search and the other calibration-fitting optimizers are untouched by this engine and run serial, exactly as before (AWQ's activation statistics do ride the sharded collection forward — see "Adopting other algorithms"). Diffusion-style input pools use dict-shaped calibration; with the flag set, such a run stops with that reason. Multimodal (VLM) calibration has not been tested with the parallel lane.
High-level design
One iteration of the tune loop, two ranks of N shown:
flowchart LR A1["Shard 1: fwd + loss + bwd"] --> X(("gradient exchange")) A2["Shard 2: fwd + loss + bwd"] --> X X --> W1["Shard 1: update"] X --> W2["Shard 2: update"]Every replica leaves the exchange with the identical averaged gradient, so every replica computes the same update and the mirrors stay identical copies. After the loop ends, the mirrors are dropped and the home block unwraps and packs through the normal path.
Per block, the engine does the following.
First it decides how this block parallelizes: how many replicas (the world), on which devices, and which sample range each replica owns. One function makes that decision for the collection pass, the search phases, and the tune loop, so every phase parallelizes the same blocks the same way. With the flag set, the decision stops the run with the reasons when the configuration is unsupported: the run must use a supported accelerator home device (cuda/xpu/hpu), a single process, and a flat list of calibration samples per block, and runs with a gradient scaler or LFQ are unsupported. The decision also prices each mirror — the block's weights plus its fp32 tuning values — and keeps the devices whose free memory covers that footprint plus a margin;
autoderives the world size from the devices that survive. Finally it requires the block to sit whole on the home device: a block whose weights span several devices of the home type stops with a placement error, because whole-block mirrors are the contract of this lane.The no-grad collection forwards then run shard by shard across short-lived mirrors, each rank consuming its own samples. Statistics that forward hooks write during these passes merge back into the home block — imatrix partial sums add up, activation maxima fold by max — and hook passes with self-contained statistics run on the home alone.
The search phases divide their work across the replicas: each module's search runs once, on one replica, and the result is copied to the same module everywhere. The searches are deterministic given a module's weight and imatrix, so the results match the serial run exactly.
For tuning, the engine builds one persistent mirror per surviving device and hands each rank a contiguous slice of the calibration pool. Every iteration, the engine dispatches the forward, loss, and backward work to the replica threads; gradients are exchanged so that every replica holds the identical average; every replica computes the same update. Loss values are normalized exactly as the serial path normalizes them, so best-iteration selection and the early-stop window behave the same. The loss readback at the end of an iteration waits for the slowest replica — that readback is also the loop's synchronization point. Runs with
gradient_accumulate_stepsabove 1 tune serial through the same gate the other requirements use; sharded accumulation follows in the follow-up PR.When the loop ends, the engine drops the mirrors, returns the merged state to the home block, and the normal unwrap and pack path takes over.
Gradient exchange. The loop needs every replica to apply the identical update each iteration, so gradients are reduced before the step — in-process, with a launcher-free single process. The exchange choice is algorithm-gated: a pure sign-SGD run exchanges the averaged gradient's signs as int8 (every replica applies bitwise-identical updates), and any other run exchanges the full averaged values in fp32. Either way the update matches what a multi-process run computes.
All parallel machinery lives behind a
TuneParallelContextobject owned by the composer. The tune loop calls guarded hooks on it — mirror setup, dispatching an iteration's work to the replica threads, gradient sync, teardown — and each hook is a no-op in serial mode, so the serial path keeps its exact current behavior, the same idea as thesetup_ddp_if_needed_seam already released on main.Adopting other algorithms
The context exposes one seam per algorithm kind, agreed in the PR thread:
Intended use cases
The target setup is a multi-GPU box and a model whose largest block, with the extra space required for quantization, fits one card. The current implementation covers the RTN/OptRTN searches and SignRound V1/V2 tuning. Blocks larger than one accelerator are future work: an exploratory multi-device-per-rank lane (PP+DP replicas spanning devices, with a ZeRO-2-lite variant) was implemented and measured — reference branches are linked in the PR thread — adding data parallelism on top of the pipeline-parallel placement (one replica across two GPUs; PP+DP on two pairs, four GPUs total) bought about 8% over PP alone, so the lane is parked until larger-block cases justify revisiting it.
Expected workflow
A requested world stops with the reasons whenever engagement fails;
autokeeps the devices that fit.Infrastructure changes and risks
The engine lives in two new modules —
sign_round/tune_parallel.pyandsign_round/data_parallel.py(mirrors, worker threads, pool sharding, the exchange) — plus guarded hooks in the tune loop and one CLI flag.The engine runs one Python process. It creates one block mirror per GPU, drives the mirrors from worker threads, and moves gradients between mirrors with device-to-device copies — through the host where peer-to-peer access is unavailable. The single process keeps streaming offload, block placement, and memory accounting in one address space. The trade: the replicas' backward passes share the Python interpreter lock, and each iteration waits for the slowest replica at the loss readback.
Relationship to the multi-process DDP on main
Main carries a multi-process DDP path for the LLM-compressor workflow: torchrun launches one process per rank, and each rank wraps the block in torch's DistributedDataParallel when it holds one GPU, or runs a manual per-parameter all-reduce when its block spans several GPUs. The calibration loader there processes the same batches on every rank, and the all-reduce averages the identical gradients. This engine divides the batches — each GPU owns its own slice — so the averaged gradient equals the serial full-batch gradient while the work is split across GPUs. The engine runs from the plain CLI, and a run uses one lane: with a torchrun group active, the resolver stops with that reason. The two lanes have not been benchmarked against each other.
Motivation and Use Case
AutoRound tunes models block by block on a single GPU while the remaining GPUs of a multi-GPU box sit idle, so the tuning wall grows with model size. #2010 — "[Feature] Support quant parallelism on multi-GPU" — is the open request for exactly this gap (a title-only report; the substance lives in its discussion: early measured data-parallel results and maintainer interest in the direction).
The existing multi-process torchrun DDP integration serves the LLM-Compressor workflow entry point; this lane covers the plain CLI/API path. The proposal: a single-process data-parallel lane behind
--parallel_quantization— replicas of the current block on the idle GPUs tune on disjoint calibration shards and exchange gradients once per iteration, with the collection and search passes sharded the same way.Measured on Qwen3.8-27B at world=4 on 24 GB cards it roughly halves per-block tuning walls at parity KL (protocol and table in #2351). Use cases: any multi-GPU host where serial block tuning dominates the quantization wall.
Alternatives Considered
Serial single-GPU block tuning, the status quo, leaves the remaining GPUs idle and pays the walls measured in #2351. Extending the existing multi-process DDP integration (the LLM-Compressor workflow) to the plain CLI was considered; this RFC chose a single-process in-engine lane, keeping the plain CLI/API path self-contained.
Definition of Done
Additional Context