feat: moe-router-bypass-batch-size

[avtc]

@Qubitium Hi, this feature adds ability to process moe weights in chunks in MoeRoutingBypass mode. Allowing to quantize large MoE models with less VRAM and devices.

Currently for each expert weight in GPTQ the hessian accumulator created during forward pass, which for example for GLM-4.5-Air is around 17Gb for up/gate moe modules for one layer. Processing expert weights in chunks require less hessian accumulator matrices and thus VRAM.

Example option usage:

moe={'routing': {'class': 'ExpertsRoutingBypass', 'batch_size': 50}}

[Qubitium]

@avtc Do you have some rough numbers for vram saving in your setup with and without the PR?

[avtc]

@Qubitium Hi,
Here are screens from quantizing a GLM-4.5-Air - layer 0 - non moe, layer 1+ - moe:
Calibration data: 2320 samples - 769K tokens.
5x3090 with first one excluded from compute: compute_device_filter=lambda devices: [d for d in devices if d.index != 0]
VramStrategy=Exclusive (i.e. data parallel).

Without batching experts, forward of moe up/gate modules on the layer 1:

Results in CUDA OOM in a while after screenshot.

With batch_size = 128, same stage:

It works, we batch 128 modules from 256 (128 x up/gate) in one forward/quantize.

With batch_size = 32, same stage for comparison:

[Qubitium]

@Qubitium Hi, Here are screens from quantizing a GLM-4.5-Air - layer 0 - non moe, layer 1+ - moe: Calibration data: 2320 samples - 769K tokens. 5x3090 with first one excluded from compute: compute_device_filter=lambda devices: [d for d in devices if d.index != 0] VramStrategy=Exclusive (i.e. data parallel).

Without batching experts, forward of moe up/gate modules on the layer 1: Results in CUDA OOM in a while after screenshot.

With batch_size = 128, same stage: It works, we batch 128 modules from 256 (128 x up/gate) in one forward/quantize.

With batch_size = 32, same stage for comparison:

Looks good! We will work on this PR after a newly awq discovered regression is fixed.

[Qubitium]

Cleanup

[Qubitium]

@avtc 5.7.0 took much longer to push out due to small bugs we keep finding/fixing in CI regression tests. Should be out soon so this can be merged.

skip