nemo_rl.modelopt.models.generation.vllm_quant_patch

Module Contents

Classes

FakeQuantWorker

Functions

_tolerate_dummy_weight_nan_amax	Scope-locally make MaxCalibrator.collect zero-fill on fully-NaN inputs.
_drop_nonclass_quant_registry_keys	Drop non-class keys from ModelOpt’s QuantModuleRegistry.
_register_routed_experts_quant_module	Register ModelOpt MoE fakequant for vLLM >= 0.25’s RoutedExperts layout.
_fakequant_run_prolog_worker

API

nemo_rl.modelopt.models.generation.vllm_quant_patch._tolerate_dummy_weight_nan_amax()

Scope-locally make MaxCalibrator.collect zero-fill on fully-NaN inputs.

When this prolog runs, vLLM’s model still has uninitialized / dummy weights — the real weights only arrive later via refit from the Megatron policy worker. Cumulative BF16 matmuls on dummy weights can overflow at deeper layers (e.g. Nemotron-3-Nano’s Mamba out_proj at layer 4) and produce NaN, which then cascades to every downstream quantizer’s input during dummy calibration.

The dummy-calibration amax is meant to be discarded — the prolog sentinels every enabled quantizer’s _amax to -1.0 immediately afterwards, and Megatron’s real amax is loaded via vllm_quant_backend.input_amax_loader during refit (max(-1.0, real)=real). So a fully-NaN input here should produce zero amax rather than crash the prolog.

Scoping this monkey-patch to the prolog (instead of editing MaxCalibrator.collect in modelopt) keeps modelopt’s source pristine and limits the workaround to the single dummy-weight code path that needs it. Genuine numerical NaN at runtime — when the calibrator is no longer active — would still be caught by the production callsite.

Nonfinite dummy activations are sanitized before calibration reduce. The patch is active only inside the dummy-weight prolog, before runtime generation starts and before real amax values are loaded.

nemo_rl.modelopt.models.generation.vllm_quant_patch._drop_nonclass_quant_registry_keys() → None

Drop non-class keys from ModelOpt’s QuantModuleRegistry.

vLLM >= 0.25 turned FusedMoE/SharedFusedMoE into factory functions, but ModelOpt’s vLLM plugin still registers them as QuantModuleRegistry keys. Registry lookups run issubclass(, ) over all keys (modelopt/torch/opt/dynamic.py::_get_registered_nn_class) and raise TypeError on a function key. MoE fakequant is re-registered against the vLLM >= 0.25 module layout in _register_routed_experts_quant_module.

nemo_rl.modelopt.models.generation.vllm_quant_patch._register_routed_experts_quant_module() → None

Register ModelOpt MoE fakequant for vLLM >= 0.25’s RoutedExperts layout.

ModelOpt’s vLLM plugin targets the pre-0.25 FusedMoE class, so with vLLM >= 0.25 no MoE quantizers are inserted at all and refit crashes on the incoming expert amax keys. The expert weights and kernels now live on the nested RoutedExperts module, so an adapted _QuantFusedMoEBase is registered there:

• Quantizer buffers land at ...experts.routed_experts.*_quantizer, which is exactly the parameter path model load_weights derives from incoming per-expert experts.N.<proj>.input_quantizer._amax keys via the expert mapping, so refit amax loading works unchanged.

• RoutedExperts owns w13_weight/w2_weight, keeping the base class’s B is self.w13_weight kernel-operand checks valid.

• The fakequant kernel swap wraps forward_modular/ forward_monolithic (RoutedExperts.forward asserts it must not be called) instead of the base class’s forward.

nemo_rl.modelopt.models.generation.vllm_quant_patch._fakequant_run_prolog_worker(self) → None

class nemo_rl.modelopt.models.generation.vllm_quant_patch.FakeQuantWorker

Bases: nemo_rl.models.generation.vllm.vllm_backend.NixlVllmWorker

determine_available_memory() → int

compile_or_warm_up_model() → float