nemo_rl.modelopt.models.generation.vllm_quant_moe_amax

Route per-expert fakequant amax buffers around vLLM’s expert weight loader.

The Megatron side exports one input_quantizer._amax per expert projection, e.g. model.layers.1.mixer.experts.16.up_proj.input_quantizer._amax. On the vLLM side the ModelOpt MoE quant module owns a single fused quantizer per projection group (w13_input_quantizer / w2_input_quantizer), so the incoming values fan in with max.

vLLM 0.25 loaded these through the model-level parameter dict, which NeMo-RL patches to include quantizer buffers. vLLM 0.28’s AutoWeightsLoader hands every experts.* name to RoutedExperts.load_weights instead, which rewrites the checkpoint name with the expert mapping and resolves the result with a single getattr(self, param_name). A dotted quantizer-buffer name such as w13_input_quantizer._amax cannot resolve that way and the loader raises AttributeError: Layer <experts> has no parameter 'w13_input_quantizer._amax'.

This module applies the same rewrite vLLM does (replace the mapping’s weight_name with its param_name, strip the layer prefix) and then walks the dotted path on the expert module itself, so the buffers are updated before the remaining weights reach vLLM’s loader.

Module Contents

Functions

_expert_modules	Return every module that owns an expert mapping and a vLLM layer name.
_resolve_dotted
route_moe_input_quantizer_amax	Fan per-expert input_quantizer._amax values into the fused quantizers.

Data

INPUT_QUANTIZER_AMAX_SUFFIX

API

nemo_rl.modelopt.models.generation.vllm_quant_moe_amax.INPUT_QUANTIZER_AMAX_SUFFIX

‘input_quantizer._amax’

nemo_rl.modelopt.models.generation.vllm_quant_moe_amax._expert_modules(model: torch.nn.Module) → list[Any]

Return every module that owns an expert mapping and a vLLM layer name.

nemo_rl.modelopt.models.generation.vllm_quant_moe_amax._resolve_dotted(module: Any, path: str) → Any | None

nemo_rl.modelopt.models.generation.vllm_quant_moe_amax.route_moe_input_quantizer_amax(

model: torch.nn.Module,

weights: collections.abc.Iterable[tuple[str, torch.Tensor]],

*,

reduce_fn: collections.abc.Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = torch.max,

mapper: Any | None = None,

) → list[tuple[str, torch.Tensor]]

Fan per-expert input_quantizer._amax values into the fused quantizers.

Returns the weights that were not consumed, in their original order, so the caller can hand them to vLLM’s own loader unchanged. Names that end with the amax suffix but do not belong to an expert module (dense layers, attention quantizers) pass through untouched.

mapper is the model’s hf_to_vllm_mapper (a vLLM WeightsMapper), when it has one. Refit sends checkpoint names and vLLM’s layer_name / expert mapping use vLLM names; load_weights applies the mapper before its own matching, so the same rename has to happen here or e.g. Nemotron-H’s backbone.layers.N.mixer.experts.* never matches the module’s model.layers.N.mixer.experts prefix and falls through to the loader that cannot resolve it. Names the mapper drops (returns None) are passed through untouched.

Raises:

KeyError – an expert amax name matched an expert module’s mapping but the rewritten quantizer path does not exist on that module. This is the same condition vLLM would have raised on, surfaced with the resolved path so the layout mismatch is visible.