[high risk]Speedup 2.5-6X moe tuning and pin ngram to cpu by wenhuach21 · Pull Request #2303 · intel/auto-round

wenhuach21 · 2026-09-04T07:29:21Z

Description

This PR introduces a new grouped-GEMM MoE experts forward path intended to significantly speed up MoE tuning, and adds mechanisms to keep very large ngram/PLE embedding tables off accelerators (CPU-pinned or GPU-sharded) to prevent VRAM/RAM blowups during per-block dispatch and tuning.

Changes:
Add a CUDA-focused grouped experts implementation (linear_grouped) with batched fake-quantization to reduce per-expert overhead during tuning.Add ngram embedding placement controls (AR_NGRAM_DEVICE) plus helpers to pin/shard ngram tables and to avoid relocating pinned/self-managed subtrees.Extend disk-stream materialization to handle wildcard shard concatenation and broader converter-driven weight reconstruction for newer checkpoint layouts.

Type of Change

Bug fix

Related Issues

Fixes or relates to #

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.

TODO

check v2 has imatrix

for more information, see https://pre-commit.ci

…into fix_moe_loading

for more information, see https://pre-commit.ci

…into suport_fuse_moe

wenhuach21 · 2026-09-04T07:34:29Z

main branch

026-09-04 15:16:39 INFO replace_modules.py L390: Found 40 modules to replace
Replacing modules: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 40/40 [00:10<00:00, 3.83it/s]
2026-09-04 15:16:49 INFO replace_modules.py L417: Replaced 40 modules
2026-09-04 15:16:49 INFO device.py L1482: [Memory Monitor] After applying custom replacements: 'peak_ram': 1.52GB
[transformers] loss_type=None was set in the config but it is unrecognized. Using the default loss: ForCausalLMLoss.
2026-09-04 15:16:53 INFO replace_modules.py L128: Experts (after replacement/skip) [model.language_model.layers.0.mlp.experts] (SequentialQwen3_5MoeExperts):
SequentialQwen3_5MoeExperts(
(0-255): 256 x Qwen3_5MoeMLP(
(gate_proj): Linear(in_features=2048, out_features=512, bias=False)
(up_proj): Linear(in_features=2048, out_features=512, bias=False)
(down_proj): Linear(in_features=512, out_features=2048, bias=False)
(act_fn): SiLUActivation()
)
)
2026-09-04 15:16:53 INFO base.py L988: Using predefined ignore_layers: model.language_model.layers.[0-39].mlp.gate
2026-09-04 15:16:53 INFO layer_config_resolver.py L401: Ignored layers: model.language_model.layers.[0-39].mlp.gate
2026-09-04 15:16:55 INFO layer_config_resolver.py L401: Ignored layers: model.language_model.layers.[0-39].mlp.gate
2026-09-04 15:16:58 INFO base.py L1329: torch.compile is enabled
2026-09-04 15:17:03 INFO orchestrator.py L570: start to cache block inputs
2026-09-04 15:17:03 INFO mllm.py L86: Using MLLM template: qwen3_5_moe
2026-09-04 15:17:03 INFO calib_dataset.py L1113: Preprocessing calibration dataset in a subprocess to avoid memory leaks...
2026-09-04 15:17:16 INFO device.py L1448: 'peak_ram': 3.08GB
2026-09-04 15:17:16 INFO orchestrator.py L602: caching done
Quantizing model.language_model.layers.0: 0%| | 0/40 [00:02<?, ?it/s]/home/wenhuach/miniforge3/envs/autoround/lib/python3.13/site-packages/torch/_dynamo/variables/functions.py:2590: UserWarning: Dynamo detected a call to a functools.lru_cache-wrapped function at 'fused_norm_gate.py:565'. Dynamo ignores the cache wrapper and directly traces the wrapped function. Silent incorrectness is only a potential risk, not something we have observed. Enable TORCH_LOGS=+dynamo for a DEBUG stack trace.

This call originates from:
File "/home/wenhuach/miniforge3/envs/autoround/lib/python3.13/site-packages/fla/modules/fused_norm_gate.py", line 565, in layer_norm_gated_bwd
NS = min(get_multiprocessor_count(x.device.index), T)
~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^

torch._dynamo.utils.warn_once(msg)
/home/wenhuach/miniforge3/envs/autoround/lib/python3.13/site-packages/torch/_dynamo/variables/functions.py:2431: UserWarning: Dynamo does not know how to trace the builtin cuda_utils.get_device_properties. This function is either a Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python extension (perhaps created with pybind).
If it is a Python builtin, please file an issue on GitHub so the PyTorch team can add support for it and see the next case for a workaround.
If it is a third-party C/C++ Python extension, please either wrap it into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html for more details) or, if it is traceable, use torch.compiler.allow_in_graph.
torch._dynamo.utils.warn_once(explanation + "\n" + "\n".join(hints))
2026-09-04 15:20:10 INFO quantizer.py L545: Unquantized layers: ['mlp.shared_expert_gate']
quantized 776/777 layers in the block, loss iter 0: 1.099e-05 -> iter 175: 8.358e-07
2026-09-04 15:20:15 INFO device.py L1448: 'peak_ram': 19.44GB, 'peak_vram': 19.45GB
Quantizing model.language_model.layers.1: 2%|████▋ | 1/40 [03:35<2:20:03, 215.47s/it]2026-09-04 15:25:11 INFO quantizer.py L545: Unquantized layers: ['mlp.shared_expert_gate']
quantized 776/777 layers in the block, loss iter 0: 3.570e-06 -> iter 182: 1.313e-06
2026-09-04 15:25:17 INFO device.py L1448: 'peak_ram': 21.55GB, 'peak_vram': 24.61GB
Quantizing model.language_model.layers.2: 5%|█████████▎ | 2/40 [08:38<2:49:13, 267.20s/it]qTraceback (most recent call last):

current

(autoround) wenhuach@mlp-dgx-01:~/auto-round$ CUDA_VISIBLE_DEVICES=2 python3 -m auto_round /models/Qwen3.6-35B-A3B --output_dir "/data2/wenhuach" --tasks lambada_openai
2026-09-04 15:12:19 INFO main.py L309: start to quantize /models/Qwen3.6-35B-A3B
2026-09-04 15:12:19 INFO model.py L347: Fused-MoE checkpoint detected: building a meta skeleton and materializing weights per block (set AR_DISABLE_AUTO_META_LOAD=1 to load the whole model on CPU instead).
2026-09-04 15:12:22 INFO device.py L1446: Before applying custom replacements 'peak_ram': 1.07GB
2026-09-04 15:12:24 INFO moe_experts_interface.py L783: [MoE Prep] Unfused 40 MOE experts modules
2026-09-04 15:12:25 INFO device.py L1446: After applying custom replacements 'peak_ram': 1.26GB
2026-09-04 15:12:25 INFO replace_modules.py L102: Prepared 40 MOE modules for quantization
2026-09-04 15:12:25 INFO replace_modules.py L161: Experts (before unfuse) [model.language_model.layers.0.mlp.experts] (Qwen3_5MoeExperts):
Qwen3_5MoeExperts(
(act_fn): SiLUActivation()
)
2026-09-04 15:12:25 INFO replace_modules.py L162: Experts (after unfuse) [model.language_model.layers.0.mlp.experts] (Qwen3_5MoeExperts):
Qwen3_5MoeExperts(
(act_fn): SiLUActivation()
(0-255): 256 x _ExpertContainer(
(down_proj): Linear(in_features=512, out_features=2048, bias=False)
(gate_proj): Linear(in_features=2048, out_features=512, bias=False)
(up_proj): Linear(in_features=2048, out_features=512, bias=False)
)
)
2026-09-04 15:12:27 INFO base.py L2127: format is not set, using default auto_round format.
2026-09-04 15:12:27 WARNING logging.py L340: some layers are skipped quantization (shape not divisible by 32): model.language_model.layers.[0-39].mlp.shared_expert_gate, model.visual.blocks.[0-26].mlp.linear_fc1, model.visual.blocks.[0-26].mlp.linear_fc2
2026-09-04 15:12:28 INFO base.py L988: Using predefined ignore_layers: model.language_model.layers.[0-39].mlp.gate
2026-09-04 15:12:28 INFO layer_config_resolver.py L401: Ignored layers: model.language_model.layers.[0-39].mlp.gate
2026-09-04 15:12:29 INFO layer_config_resolver.py L401: Ignored layers: model.language_model.layers.[0-39].mlp.gate
2026-09-04 15:12:33 INFO base.py L1329: torch.compile is enabled
2026-09-04 15:12:33 INFO composer.py L219: Block-forward torch.compile is disabled because grouped MoE experts run eager; the per-layer qdq stays compiled.
2026-09-04 15:12:37 INFO orchestrator.py L570: start to cache block inputs
2026-09-04 15:12:37 INFO mllm.py L86: Using MLLM template: qwen3_5_moe
2026-09-04 15:12:37 INFO calib_dataset.py L1113: Preprocessing calibration dataset in a subprocess to avoid memory leaks...
2026-09-04 15:12:50 INFO device.py L1448: 'peak_ram': 3.13GB
2026-09-04 15:12:50 INFO orchestrator.py L602: caching done
Quantizing model.language_model.layers.0: 0%| | 0/40 [00:01<?, ?it/s]
2026-09-04 15:13:40 INFO quantizer.py L545: Unquantized layers: ['mlp.shared_expert_gate']
quantized 776/777 layers in the block, loss iter 0: 1.099e-05 -> iter 199: 8.291e-07
2026-09-04 15:13:45 INFO device.py L1448: 'peak_ram': 17.78GB, 'peak_vram': 24.52GB
Quantizing model.language_model.layers.1: 2%|████▋ | 1/40 [01:30<59:04, 90.87s/it]
2026-09-04 15:15:05 INFO quantizer.py L545: Unquantized layers: ['mlp.shared_expert_gate']
quantized 776/777 layers in the block, loss iter 0: 3.474e-06 -> iter 182: 1.286e-06
2026-09-04 15:15:09 INFO device.py L1448: 'peak_ram': 17.78GB, 'peak_vram': 27.47GB
Quantizing model.language_model.layers.2: 5%|█████████▍ | 2/40 [02:53<54:38, 86.26s/it]

for more information, see https://pre-commit.ci

wenhuach21 · 2026-09-06T14:23:31Z

accuracy is not matched for now

for more information, see https://pre-commit.ci

…into suport_fuse_moe

for more information, see https://pre-commit.ci

…into suport_fuse_moe

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

…into suport_fuse_moe

for more information, see https://pre-commit.ci

wenhuach21 · 2026-09-08T09:31:34Z

(autoround) wenhuach@mlp-dgx-01:~/auto-round$ CUDA_VISIBLE_DEVICES=6 python3 -m auto_round /models/Qwen3.6-35B-A3B/ --tasks mmlu --format fake --output_dir /data4/wenhuach --eval_bs 4
W0908 17:27:08.914000 3454376 site-packages/torch/utils/_pytree.py:630] <enum 'KernelPreference'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
W0908 17:27:08.944000 3454376 site-packages/torch/utils/_pytree.py:630] <enum 'ScaleCalculationMode'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
2026-09-08 17:27:10 INFO main.py L309: start to quantize /models/Qwen3.6-35B-A3B
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1026/1026 [00:00<00:00, 6885.17it/s]
2026-09-08 17:27:14 INFO device.py L1591: [Memory Monitor] Before applying custom replacements: 'peak_ram': 1.3GB
2026-09-08 17:27:14 INFO replace_modules.py L421: Found 40 modules to replace
Replacing modules: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 40/40 [00:12<00:00, 3.21it/s]
2026-09-08 17:27:26 INFO replace_modules.py L448: Replaced 40 modules
2026-09-08 17:27:26 INFO device.py L1594: [Memory Monitor] After applying custom replacements: 'peak_ram': 1.42GB
2026-09-08 17:27:27 INFO replace_modules.py L160: Experts (before replacement) [model.language_model.layers.0.mlp.experts] (Qwen3_5MoeExperts):
Qwen3_5MoeExperts(
(act_fn): SiLUActivation()
)
2026-09-08 17:27:27 INFO replace_modules.py L161: Experts (after replacement) [model.language_model.layers.0.mlp.experts] (SequentialQwen3_5MoeExperts):
SequentialQwen3_5MoeExperts(
(0-255): 256 x Qwen3_5MoeMLP(
(gate_proj): Linear(in_features=2048, out_features=512, bias=False)
(up_proj): Linear(in_features=2048, out_features=512, bias=False)
(down_proj): Linear(in_features=512, out_features=2048, bias=False)
(act_fn): SiLUActivation()
)
(256): SiLUActivation()
)
2026-09-08 17:27:27 INFO base.py L999: Using predefined ignore_layers: model.language_model.layers.[0-39].mlp.gate
2026-09-08 17:27:27 INFO layer_config_resolver.py L401: Ignored layers: model.language_model.layers.[0-39].mlp.gate
2026-09-08 17:27:29 INFO layer_config_resolver.py L401: Ignored layers: model.language_model.layers.[0-39].mlp.gate
2026-09-08 17:27:32 INFO base.py L1340: torch.compile is enabled
2026-09-08 17:27:37 INFO orchestrator.py L587: start to cache block inputs
2026-09-08 17:27:37 INFO mllm.py L86: Using MLLM template: qwen3_5_moe
2026-09-08 17:27:37 INFO mllm.py L125: Multimodal model with non-MLLM calibration dataset 'NeelNanda/pile-10k' and quant_nontext_module=False: using the standard text dataloader (vision/audio towers are not being quantized, so text-only calibration through the full-model forward is sufficient).
2026-09-08 17:27:37 INFO calib_dataset.py L1113: Preprocessing calibration dataset in a subprocess to avoid memory leaks...
2026-09-08 17:27:50 INFO device.py L1560: 'peak_ram': 3.01GB
2026-09-08 17:27:50 INFO orchestrator.py L619: caching done
Quantizing model.language_model.layers.0: 0%| | 0/40 [00:02<?, ?it/s]2026-09-08 17:28:44 INFO quantizer.py L548: Unquantized layers: ['mlp.shared_expert_gate']
quantized 776/777 layers in the block, loss iter 0: 1.099e-05 -> iter 140: 8.766e-07
2026-09-08 17:28:49 INFO device.py L1560: 'peak_ram': 18.94GB, 'peak_vram': 25.18GB
Quantizing model.language_model.layers.1: 2%|████▍ | 1/40 [00:58<38:17, 58.91s/it]2026-09-08 17:29:36 INFO quantizer.py L548: Unquantized layers: ['mlp.shared_expert_gate']
quantized 776/777 layers in the block, loss iter 0: 3.543e-06 -> iter 182: 1.309e-06
2026-09-08 17:29:40 INFO device.py L1560: 'peak_ram': 22.13GB, 'peak_vram': 28.29GB
Quantizing model.language_model.layers.2: 5%|████████▉ | 2/40 [01:50<34:34, 54.58s/it]2026-09-08 17:30:21 INFO quantizer.py L548: Unquantized layers: ['mlp.shared_expert_gate']
quantized 776/777 layers in the block, loss iter 0: 7.851e-06 -> iter 191: 2.250e-06
2026-09-08 17:30:26 INFO device.py L1560: 'peak_ram': 25.28GB, 'peak_vram': 28.62GB
Quantizing model.language_model.layers.3: 8%|█████████████▎ | 3/40 [02:35<31:06, 50.44s/it]2026-09-08 17:31:10 INFO quantizer.py L548: Unquantized layers: ['mlp.shared_expert_gate']
quantized 775/776 layers in the block, loss iter 0: 9.733e-06 -> iter 173: 3.259e-06
2026-09-08 17:31:14 INFO device.py L1560: 'peak_ram': 25.37GB, 'peak_vram': 28.62GB
Quantizing model.language_model.layers.4: 10%|█████████████████▊ | 4/40 [03:24<29:48, 49.69s/it]

…into suport_fuse_moe

for more information, see https://pre-commit.ci

…into suport_fuse_moe

AutoRoundBot · 2026-09-08T12:23:29Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-09-08T12:23:40Z

Azure Pipelines successfully started running 1 pipeline(s).

AutoRoundBot · 2026-09-08T15:16:24Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-09-08T15:16:36Z

Azure Pipelines successfully started running 1 pipeline(s).

wenhuach21 · 2026-09-09T01:33:27Z

please have a review

wenhuach21 · 2026-09-09T05:41:04Z

nemetron still has issues, trackign here #2324

wenhuach21 · 2026-09-09T05:42:23Z

merge first, will fix issues in other prs

wenhuach21 and others added 13 commits September 3, 2026 15:22

update

5fe714d

[pre-commit.ci] auto fixes from pre-commit.com hooks

b0556ad

for more information, see https://pre-commit.ci

update

9f9331c

update

f644f23

Merge branch 'main' into fix_moe_loading

fd6f5f0

[pre-commit.ci] auto fixes from pre-commit.com hooks

c0b2cee

for more information, see https://pre-commit.ci

disable all moe replacements

84fe355

Merge branch 'fix_moe_loading' of https://github.com/intel/auto-round …

481f1a3

…into fix_moe_loading

speedup tuning

afd47ff

no speedup gain

8d0edbf

[pre-commit.ci] auto fixes from pre-commit.com hooks

1b36314

for more information, see https://pre-commit.ci

update

a1fbead

Merge branch 'suport_fuse_moe' of https://github.com/intel/auto-round …

640e31b

…into suport_fuse_moe

wenhuach21 added 2 commits September 4, 2026 15:37

update

a68ad42

better support multi-gpus

1524006

wenhuach21 changed the title ~~Speedup 3X moe tuning~~ Speedup 3X moe tuning via grouped experts Sep 4, 2026

wenhuach21 and others added 3 commits September 4, 2026 22:59

Merge branch 'main' into suport_fuse_moe

1337a48

[pre-commit.ci] auto fixes from pre-commit.com hooks

a1a5b4d

for more information, see https://pre-commit.ci

clean

ce8be47

wenhuach21 and others added 9 commits September 7, 2026 10:57

Merge branch 'main' into suport_fuse_moe

47254b8

update

a7ee60b

[pre-commit.ci] auto fixes from pre-commit.com hooks

d98505d

for more information, see https://pre-commit.ci

update

8f0af61

Merge branch 'suport_fuse_moe' of https://github.com/intel/auto-round …

37573f8

…into suport_fuse_moe

support group_mm in monkey patch

e5e67e1

[pre-commit.ci] auto fixes from pre-commit.com hooks

610b923

for more information, see https://pre-commit.ci

fix qwen4

28c72d9

Merge branch 'suport_fuse_moe' of https://github.com/intel/auto-round …

b075757

…into suport_fuse_moe

wenhuach21 and others added 4 commits September 8, 2026 17:20

Potential fix for pull request finding

e35bcd4

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

fix issue

f960c8e

Merge branch 'suport_fuse_moe' of https://github.com/intel/auto-round …

96c067a

…into suport_fuse_moe

[pre-commit.ci] auto fixes from pre-commit.com hooks

14f3a2e

for more information, see https://pre-commit.ci

wenhuach21 changed the title ~~[high risk]Speedup 2.5-4X moe tuning via grouped experts and pin ngram to cpu to fix qwen ram issue~~ [high risk]Speedup 2.5-4X moe tuning and pin ngram to cpu Sep 8, 2026

wenhuach21 requested review from lvliang-intel, n1ck-guo, xin3he and yiliu30 September 8, 2026 09:32

wenhuach21 changed the title ~~[high risk]Speedup 2.5-4X moe tuning and pin ngram to cpu~~ [high risk]Speedup 2.5-6X moe tuning and pin ngram to cpu Sep 8, 2026

wenhuach21 and others added 6 commits September 8, 2026 17:47

update log

a63adbc

Merge branch 'suport_fuse_moe' of https://github.com/intel/auto-round …

ffb592a

…into suport_fuse_moe

[pre-commit.ci] auto fixes from pre-commit.com hooks

7578ec0

for more information, see https://pre-commit.ci

update log

38868ff

Merge branch 'suport_fuse_moe' of https://github.com/intel/auto-round …

a2c2317

…into suport_fuse_moe

try to fix ut

96930d5

try to fix ut

e334b63

This was referenced Sep 9, 2026

feat: support Nemotron-H / Nemotron-Cascade-2 (#1711) #1712

Closed

[Feature]: very slow to quatnize qwen3.8-flash-next for full attention layers #2323

Open

wenhuach21 merged commit 678f8f5 into main Sep 9, 2026
48 of 51 checks passed

wenhuach21 deleted the suport_fuse_moe branch September 9, 2026 05:42

avtc mentioned this pull request Sep 14, 2026

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

Open

4 tasks

Navigation Menu

[high risk]Speedup 2.5-6X moe tuning and pin ngram to cpu - #2303

[high risk]Speedup 2.5-6X moe tuning and pin ngram to cpu#2303
wenhuach21 merged 52 commits into
mainfrom
suport_fuse_moe

wenhuach21 commented Sep 4, 2026 •

edited

Loading

Uh oh!

wenhuach21 commented Sep 4, 2026

Uh oh!

wenhuach21 commented Sep 6, 2026

Uh oh!

wenhuach21 commented Sep 8, 2026

Uh oh!

AutoRoundBot commented Sep 8, 2026

Uh oh!

azure-pipelines Bot commented Sep 8, 2026

Uh oh!

AutoRoundBot commented Sep 8, 2026

Uh oh!

azure-pipelines Bot commented Sep 8, 2026

Uh oh!

wenhuach21 commented Sep 9, 2026

Uh oh!

wenhuach21 commented Sep 9, 2026

Uh oh!

wenhuach21 commented Sep 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

wenhuach21 commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of Change

Related Issues

Checklist Before Submitting

TODO

Uh oh!

wenhuach21 commented Sep 4, 2026

main branch

current

Uh oh!

wenhuach21 commented Sep 6, 2026

Uh oh!

wenhuach21 commented Sep 8, 2026

Uh oh!

AutoRoundBot commented Sep 8, 2026

Uh oh!

azure-pipelines Bot commented Sep 8, 2026

Uh oh!

AutoRoundBot commented Sep 8, 2026

Uh oh!

azure-pipelines Bot commented Sep 8, 2026

Uh oh!

wenhuach21 commented Sep 9, 2026

Uh oh!

wenhuach21 commented Sep 9, 2026

Uh oh!

wenhuach21 commented Sep 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wenhuach21 commented Sep 4, 2026 •

edited

Loading