`nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker`#

Module Contents#

Classes#

MegatronQuantPolicyWorker

Functions#

`_quant_checkpoint_cache_suffix`	Build a short suffix for HF->Megatron checkpoints with ModelOpt state.
`_find_other_quant_checkpoint_caches`	Find valid quantized startup caches other than the selected cache.
`_warn_if_other_quant_checkpoint_caches`	Warn when a different quantization config already has a startup cache.
`_set_quantization_model_specs`	Select quantization-compatible specs across Bridge hybrid API versions.

API#

nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker._quant_checkpoint_cache_suffix( config: collections.abc.Mapping[str, object], ) → str#: Build a short suffix for HF->Megatron checkpoints with ModelOpt state.

nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker._find_other_quant_checkpoint_caches( base_pretrained_path: str, selected_pretrained_path: str, ) → list[pathlib.Path]#: Find valid quantized startup caches other than the selected cache.

nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker._warn_if_other_quant_checkpoint_caches( base_pretrained_path: str, selected_pretrained_path: str, ) → None#: Warn when a different quantization config already has a startup cache.

nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker._set_quantization_model_specs( model_config, disable_modelopt_layer_spec: bool, )#

Select quantization-compatible specs across Bridge hybrid API versions.

Recent Megatron-Bridge revisions load Nemotron-H checkpoints as a HybridModelProvider and use hybrid_stack_spec. Older revisions use the deprecated mamba_stack_spec field. Setting only the latter leaves the recent provider to infer the local ModelOpt stack when restore_modelopt_state=True; that stack contains SequentialMLP and cannot restore quantizers with both tensor and expert parallelism enabled.

class nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker(config, *args, **kwargs)#

Bases: nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorkerImpl

maybe_init_zmq() → None#: Use a longer timeout only for ModelOpt real-quant refits.

_quantize(model)#: Quantize the model if the model is not quantized yet.

_patch_validate_model_paths()#

Patch validate_model_paths to handle quantized checkpoint paths.

In cases like distillation where the teacher model is the same as the student model, we need to save an extra quantized checkpoint. This patch routes auto-converted HF checkpoints to a ModelOpt-specific cache path. It also handles pre-quantized model symlinks.

_patch_setup_model_and_optimizer()#: Patch setup_model_and_optimizer to restore modelopt state.

_restore_modelopt_state_pre_load(state, model)#

Restore ModelOpt state into the model before load_checkpoint runs.

Forwarded as the pre_load_checkpoint_hook to

Func:: setup_model_and_optimizer and :func:setup_reference_model_state via the _pre_load_checkpoint_hook instance attribute. Quantizers must exist on the model graph before load_checkpoint populates their amax/scale buffers.

hide_tensor_quantizers()#: Context manager that temporarily hides TensorQuantizer modules from module iteration.

enable_forward_pre_hook()#: Enable forward pre-hook, hiding TensorQuantizer modules.

disable_forward_pre_hook(param_sync=True)#: Disable forward pre-hook, hiding TensorQuantizer modules.

disable_quantization()#: Context manager that temporarily disables quantization.

_hide_extra_state()#

Patch model.state_dict() to exclude _extra_state keys.

ModelOpt appends quantization calibration data (amax/scale) to TE’s serialized extra state, making it larger than the non-quantized reference model’s copy. These are calibration metadata, not learned weights, and can also be resized by TE during forward passes. Filtering them out lets the base class swap/restore skip them cleanly.

use_reference_model() → collections.abc.Generator[None, None, None]#: Context manager that temporarily swaps the reference model and active model.

without_model_config()#

Temporarily remove TensorQuantizer config attributes.

Used by :meth:use_reference_model and :meth:save_checkpoint. Both of these flows traverse the module tree (e.g. for state-dict swapping or checkpoint serialization) where the unrelated config attribute on TensorQuantizer instances is detected as a model config and triggers spurious validation/serialization errors. We strip it for the duration of the call and restore it on exit.

get_quantizer_stats() → dict#

Return summary statistics for all enabled TensorQuantizers.

Useful for verifying that calibration ran and amax values are valid.

abstractmethod generate(**kwargs)#

Quantized Megatron generation is not supported.

ModelOpt unconditionally patches flash_decode_and_prefill on quantized attention modules, which breaks the Megatron generation path.

save_checkpoint(*args, **kwargs)#: Save the checkpoint.

_use_real_quant_refit() → bool#

_get_real_quant_mode() → str#: Resolve and cross-check the training and rollout quantization modes.

_iter_real_quant_refit_params( kv_scales: dict[str, float] | None = None, ) → collections.abc.Generator[tuple[str, torch.Tensor], None, None]#: Export packed NVFP4 weights and scales for real-quant vLLM rollout.

static _find_weight_quantizer( module: object, param_weight: object, ) → object | None#

Find the enabled weight quantizer that corresponds to param_weight.

Uses ModelOpt’s QuantModule.iter_weights_for_calibration to discover (weight, weight_quantizer) pairs, then matches by identity. This handles standard weight / weight_quantizer as well as custom names like gate_up_proj / gate_up_proj_weight_quantizer.

Returns the matching TensorQuantizer or None.

static _iter_hf_input_amax_names(mapping)#

static _get_enabled_input_amax(task)#

_iter_input_quantizer_amax_params(conversion_tasks, existing_names)#

_iter_params_with_optional_kv_scales(kv_scales=None)#

Pre-fold weights on-the-fly via lazy proxy tasks.

Wraps each conversion task so that reading task.param_weight returns weight_quantizer(weight) instead of the raw weight. The folded tensor is computed lazily when export_hf_weights accesses it, so only one extra weight-sized tensor exists at a time — O(1) extra memory.

Raises:: RuntimeError – If weight folding fails for a specific parameter, with context about which parameter caused the failure.

nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker`#