feat: multi-GPU quantization improvements -- batched device-local sea…

[avtc]

Description

Multi-GPU tuning for large MoE models (validated end-to-end on a 300B-A21B MoE, 8x3090, W4A16): device-local batched quantization searches, calibration-data pool placement, atomic expert placement with per-device grouped-experts plans, and a memory-model-driven auto pick of the experts implementation. Plus four upstream bugs the new instrumentation surfaced (fixes benefit single-GPU users too).

Features

• Batched device-local quantization searches (iters=0 OptRTN and iters>0 SignRound V2). Scale searches for modules with identical shape and quantization config stack into one call per device; when a block's weights span several GPUs, the devices run their batches in parallel. In the iters=0 lane stacked batches may also run on other CUDA GPUs with free memory, including GPUs outside the device map.
• --calibration_data_device auto | off | cpu | cuda:1,cuda:2 | 1,2 (CLI and Python API). Per-block placement of the calibration input/output pools: they stay on the primary GPU when they fit beside the activation window, and are spread across the other GPUs otherwise, placed so that every GPU's peak VRAM during tuning stays close to the same level. Placement decisions come from a memory model of the block (activation window, collection retention, tuning state, and activation estimates computed per GPU from the modules whose weights live there).
• Best-params snapshots next to their weights. The running-best snapshot lives on the GPU hosting its weight instead of host RAM, removing a host round trip per saved improvement; under memory pressure it falls back to host.
• Reference outputs on the loss device. The fp reference pool was previously gathered across GPUs on every tune iteration; it now moves once to the loss device before the loop when it fits, eliminating the per-iteration gather.
• FP input pool released. The fp input pool is released at tune-loop start, freeing its VRAM for the loop.
• Background shard writes. Shard saves and the AR_RESUME_DIR marks are written in the background, overlapping with the next block's tuning.
• Atomic expert placement. The data-driven allocator places whole experts on single devices; chunks split by whole experts when no single device has budget. Previously the projections of one expert landed on different GPUs because expert detection matched only nn.ModuleList containers, and the MoE unfuse path builds a plain container.
• Per-device grouped-experts plans. The grouped-GEMM fast path ([high risk]Speedup 2.5-6X moe tuning and pin ngram to cpu #2303) required all active experts of a layer on one device; experts are now grouped by home device and run one stacked grouped GEMM per group, so the fast path stays active when a layer's experts span several GPUs.
• AR_MOE_EXPERTS_IMPL=auto (new value of the existing env var, now the default). Predicts the per-device VRAM of grouped tuning before the first MoE block and selects linear_loop for the run when grouped exceeds the free memory; explicit values take precedence.

Upstream fixes (single-GPU users benefit too)

• layer_config rebuilt the full compression-plan dict on every read (~50 ms on a 300B plan); the pack loops read it once per module, which cost ~27 s per MoE block. The reads are hoisted.
• Immediate pack ran after the block moved to CPU; it now runs while the block is still on the GPUs, matching the model-free path. Packed artifacts still land in host RAM.
• With --enable_torch_compile, dynamo/FX caches retain tracing artifacts across blocks (~5 GiB in the block-3 OOM census on the 300B lane) until the recompile limit forces an eager fallback. Blocks wrap fresh modules, so no compiled graph is reused across blocks; the block-boundary dynamo reset (when compilation is enabled) frees the dead caches and keeps retention from accumulating.
• The block memory estimator crashed on V2-wrapped modules (act_bits was read off the wrapper instead of the wrapped layer); it now reads quantization attributes through orig_layer.

Instrumentation (opt-in)

AR_PERF_COUNTERS emits a per-block phase summary; an OOM tensor census with holder attribution is installed by the CLI; pool-placement decisions are logged. The three upstream fixes were found with these tools.

Measured results (300B-A21B MoE, 80 blocks; iters=0 lane = optimized RTN, symmetric scheme; iters=20 lane = SignRound V2)

Per-block phase walls on 8x3090, serial baseline (with fixes unlocking iters=20) vs this branch:

Phase	Before	After	Speedup
quant scale search (iters=0)	61.4 s	9.4 s	6.5x
wrap scale search (iters=20)	80 s	7.7 s	10.4x
tune loop (iters=20), block 1	62.8 s	38.1 s	1.6x
tune loop (iters=20), blocks 2+	99.3 s	47.7 s	2.1x
pack (layer_config fix)	29 s	1.2 s	24x
move block to CPU (pack order)	4.5 s	1.2 s	3.8x

Full runs, same rig and model:

Configuration	main	this branch	Speedup
8x3090, iters=0	120 s/block, 2.7 h	26 s/block, 34 min	4.6x
8x3090, iters=20	CUDA OOM in the first MoE block	112 s/block, 2.5 h	unlocked
4x3090, iters=0	CUDA OOM in block 2	28.5 s/block, 38 min	unlocked
4x3090, iters=20	-	166 s/block, 3.7 h (auto linear_loop)	-

VRAM stays flat across blocks on the branch's full runs. On 5 and more GPUs the grouped-experts path stays active.

Numerical parity: iters=0 outputs are bit-identical to the serial path; at iters=20 the tuning losses match the serial path within run noise (7.305e-11 to 2.670e-11 batched, 2.671e-11 serial).

New knobs

• --calibration_data_device auto | off | cpu | cuda:1,cuda:2 | 1,2 (CLI and Python API): calibration pool placement
• AR_MOE_EXPERTS_IMPL=auto: new value of the existing env var, picks grouped vs linear_loop by predicted memory fit
• AR_DISABLE_BATCHED_SEARCH=1: serial search fallback
• AR_DISABLE_MULTIGPU_SEARCH=1: keeps search batches on the devices hosting the weights
• AR_SEARCH_BATCH_GB: VRAM cap for stacked search batches
• AR_PERF_COUNTERS=1: per-block phase timing

All documented in docs/environments.md and its CN counterpart.

Tests

About 3,100 lines of new CPU-runnable tests: search batching and dispatch, pool placement, grouped-experts plans and atomic placement (including a CUDA-gated two-device test), the auto linear_loop pick, pool release and writer ordering.

Type of Change

New feature (performance) + bug fix

Related Issues

Relates to #2303 (keeps the grouped-experts fast path active when experts span several GPUs)

Checklist Before Submitting

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

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).