[avtc]

Description

Follow-up to #2351. Under --data_parallel, AWQ (--algorithm awq, awq,auto_round) ran entirely on the home GPU while the other GPUs idled: its calibration forward, the smoothing-ratio grid search and the clip search. This PR shards all three across the same GPUs. The quantized model is bit-identical to the serial run. It also fixes a torch.compile crash in the thread ranks that #2351 left in place.

AWQ under --data_parallel

• Calibration forward. The block forward that carries AWQ's hooks (activation statistics, clip features, cached parent-module calls) runs on one copy of the block per GPU, each over a contiguous slice of the samples. Each GPU keeps its captured calls. The captures are merged in sample order, and the activation sums are re-added call by call in that order, so every statistic equals the serial one. A serial run still parks the captured calls in host memory, as before. The pass uses one-pass copies of the block, not the block's mirror pool: AWQ changes the weights right after this pass, and the pool is reused by the passes after the transform.
• Grid search. The cached parent calls are split across the GPUs. Each GPU replays its slice on its own copy of the mapping's parent module (the home GPU uses the parent itself). It returns the loss sum for every candidate plus its element count. The sums are added up, and the candidate is chosen exactly as in the serial loop, on sum / count.
• Clip search. The per-layer searches are independent, so they are spread over the GPUs, each on a copy of the layer weight.
• When it stays serial. A search falls back to the unchanged serial loop when there are too few calls, the parent module lies outside the block, a balance layer lies outside the parent, or no reference output survives. A calibration pass the engine would not shard (a non-list pool, fewer samples than GPUs) keeps the serial capture.
• Interfaces. The engine side is a block-scoped search seam (auto_round/algorithms/parallel/search_shard.py); the composer attaches it around pre_quantize_block only. Preprocessors opt in through two BasePreprocessor hooks: set_parallel_search(seam) for the searches and sharded_calibration(devices) for the calibration pass. The defaults keep everything serial, so other transforms are unaffected.
• Memory. The captured calls stay on their GPUs until the block's searches finish: about 1.5 GB more per GPU on a 27B-class block (peak 13.4 GB on the home GPU, 9.4 GB on the others, in the run below).

Fix: compiled calls in data-parallel threads

On torch < 2.14, torch.fx keeps its "symbolic tracing in progress" flag in a process-wide global (torch 2.14+ keeps it per thread, so the race below cannot happen there). The thread ranks compile some specializations of the quant functions after the serial first phase: a new shape, device or grad mode (for example the no-grad final QDQ at the end of a block). While one rank traces such a compile, a peer that enters a compiled function raises RuntimeError: Detected that you are using FX to symbolically trace a dynamo-optimized function. Reproduced on Qwen3.8-27B, --iters 10 --data_parallel 4: 3 of 5 runs failed on block 1, with or without AWQ. Each engine thread (spmd.run_ranks, run_threaded_spawn, the search seam) now runs with dynamo's error_on_nested_fx_trace off, so such a peer call runs eagerly for that one call. The tracing thread is unaffected (its dispatch modes are thread-local). The setting is applied per thread because dynamo keeps its config in context variables. On torch 2.14+ the guard is a no-op.

Measured (Qwen3.8-27B, W4A16 g128, --algorithm awq, RTN, --iters 0, --device_map 0, 128 samples, RTX 3090; times per block are medians):

	serial	--data_parallel 4
calibration forward (s/block)	10.2	1.1
grid search (s/block)	26.2	4.2
reference collection (s/block)	2.3	1.0
block total (s)	41.4	8.6
model total	45.1 min	10.1 min
quantized model	reference	identical (2013 tensors compared)

With --awq_apply_clip, the clip search went from 22.9 to 7.8 s per block before the calibration forward was sharded. On main, --data_parallel 4 left the AWQ phases serial: per-block time with clip stayed at the serial level (63-72 s against ~64 s serial). awq,auto_round with --iters 10 --data_parallel 4 and torch.compile passes with the fix (8-block smoke run); before the fix it failed on block 1.

Diagnostics. No new options. Under AR_PERF_COUNTERS there is a per-block [perf] pipeline phases line (calibration forward, pre-quantize transforms, reference and quantized-input collection, tuning, cascade), a per-block [perf] awq searches line, and a [perf] awq grid split line per grid search (documented in docs/environments.md). docs/step_by_step.md describes the sharded AWQ phases.

Tests.

• test/unit/test_cpu/quantization/test_awq_sharded_search.py (new):
  • the seam (slicing, private replicas, no-grad workers, error propagation);
  • grid search parity with the serial search (losses and selected scales, for every duo_scaling mode and the micro-batched replay);
  • clip search parity;
  • the serial fallbacks;
  • capture merge order and bit-equal statistics;
  • composer wiring.
• test/unit/test_cpu/quantization/test_spmd_ranks.py: compiled calls stay legal in rank and spawned threads while a peer traces. These tests fail without the fix.

Type of Change

New feature / Bug fix / Performance

Related Issues

Relates to #2351 (follow-up)

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.

🤖 Generated with Claude Code

[azure-pipelines]

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.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

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