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.
_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._fakequant_run_prolog_worker(self) → None

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

Bases: vllm.v1.worker.gpu_worker.Worker

determine_available_memory() → int

compile_or_warm_up_model() → float