nemo_rl.models.policy.workers.megatron_policy_worker#

Module Contents#

Classes#

MegatronPolicyWorkerImpl

MegatronPolicyWorker

Functions#

_should_use_router_replay

_model_self_packs_for_cp

Whether the model packs sequences + CP-shards inside its own forward.

Data#

log

TokenizerType

API#

nemo_rl.models.policy.workers.megatron_policy_worker.log#

‘getLogger(…)’

nemo_rl.models.policy.workers.megatron_policy_worker.TokenizerType#

‘TypeVar(…)’

nemo_rl.models.policy.workers.megatron_policy_worker._should_use_router_replay(
*,
enabled: bool,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
stage: str,
require: bool,
) → bool#
nemo_rl.models.policy.workers.megatron_policy_worker._model_self_packs_for_cp(model: Any) → bool#

Whether the model packs sequences + CP-shards inside its own forward.

Such models (mbridge VLM wrappers) call preprocess_packed_seqs in their forward, so NeMo-RL must hand them an unpacked [B, S] batch instead of pre-packing + CP-sharding itself. The only such model today is mbridge’s Qwen3VL, which is also the only mbridge VLM that supports context parallelism; classic mcore GPTModel and other VLMs do not self-pack.

class nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorkerImpl(
config: nemo_rl.models.policy.PolicyConfig,
tokenizer: nemo_rl.models.policy.workers.megatron_policy_worker.TokenizerType,
weights_path: Optional[str] = None,
optimizer_path: Optional[str] = None,
init_optimizer: bool = True,
init_reference_model: bool = True,
*,
worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding,
**kwargs: Any,
)#

Bases: nemo_rl.models.generation.megatron.megatron_worker.MegatronGenerationMixin, nemo_rl.models.generation.megatron.megatron_worker.MegatronGenerationRefitMixin, nemo_rl.data_plane.worker_mixin.TQWorkerMixin, nemo_rl.models.policy.workers.base_policy_worker.AbstractPolicyWorker, nemo_rl.models.policy.interfaces.ColocatablePolicyInterface

_train_step_state: Optional[dict[str, Any]]#

None

__repr__()#

Customizes the actor’s prefix in the Ray logs.

This makes it easier to identify which worker is producing specific log messages.

_local_coords() → dict[str, int]#
_get_replica_group() → Optional[Any]#

Replica group = TP × CP × PP siblings within this DP rank.

Always returns the real group so :meth:_is_replica_leader (used by both fetch and write-back) gives the correct single-writer answer even at CP=1 — gating on CP=1 here is what produced the -601 ILLEGAL_CLIENT duplicate-write bug. The fetch-path broadcast-vs-independent perf choice lives inside _fetch keyed on replica_group.size().

mcore exposes per-axis groups (get_tensor_model_parallel_group, get_context_parallel_group, get_pipeline_model_parallel_group) but no single combined group. We build the combined NCCL group once on first call by enumerating coordinates that share this rank’s dp_rank.

static configure_worker(
num_gpus: int | float,
bundle_indices: Optional[tuple[int, list[int]]] = None,
) → tuple[dict[str, Any], dict[str, str], dict[str, Any], dict[str, Any]]#

Worker-controlled Ray actor configuration.

Ensures that communication via NVLS functions correctly.

Parameters:

  • num_gpus – Original GPU allocation for this worker based on the placement group

  • bundle_indices – Tuple of (node_idx, local_bundle_indices) for this server

Returns:

  • ‘resources’: Resource allocation (e.g., num_gpus)

  • ’env_vars’: Environment variables for this worker

  • ’init_kwargs’: Parameters to pass to init of the worker

  • ’runtime_env’: Additional runtime_env options (e.g., nsight config)

Return type:

tuple with complete worker configuration

enable_forward_pre_hook()#
disable_forward_pre_hook(param_sync=True)#
_forward_pre_hook_enabled() → bool#
_disable_forward_pre_hook_until_next_train_step() → None#
_copy_main_params_to_param_buffer(
zero_grad_buffer: bool = False,
) → None#
_uses_mxfp8_overlap_shared_param_buffer() → bool#
_get_model_extra_state_dict() → dict[str, Any]#
_restore_model_extra_state_dict(
extra_state: dict[str, Any],
) → None#
train(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict,
loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,
eval_mode: bool = False,
gbs: Optional[int] = None,
mbs: Optional[int] = None,
check_dim_skip_keys: Optional[Iterable[str]] = None,
) → dict[str, Any]#

Train the policy on a batch of data with a given loss function.

check_dim_skip_keys is accepted for parity with the v1/v2 DTensor workers (cross-tokenizer ride-along tensors whose dim 1 is not the student sequence axis). Megatron doesn’t run cross-tokenizer, so it must be None.

_compute_moe_grad_scale(global_valid_toks)#

Build a moe_grad_scale_func that normalizes the aux-loss gradient.

Returns a callable yielding loss_scale = 1/global_valid_toks (clamped to avoid division by zero) so the MoE aux gradient is normalized consistently with the main per-token SFT loss. See the call site in train() for the full derivation.

_set_moe_grad_scale_func(func)#

Set moe_grad_scale_func on the model config for MOE aux loss scaling.

get_reference_policy_logprobs(
*,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
micro_batch_size: Optional[int] = None,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.ReferenceLogprobOutputSpec]#
_split_step_state_init(
loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,
gbs: Optional[int],
mbs: Optional[int],
) → dict[str, Any]#
_assert_step_open() → dict[str, Any]#
_restore_saved_grad_sync_func(state: dict[str, Any]) → None#

Restore the mcore hooks nulled in begin_train_step.

Restores both grad_sync_func and no_sync_func from the saved values on the open-step state. Idempotent on those values; safe to call from the happy-path finish/abort or from a try/except cleanup in train_microbatch / finish_train_step when those raise mid-body. See begin_train_step for why .config is read via getattr-by-string.

begin_train_step(
loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,
gbs: Optional[int] = None,
mbs: Optional[int] = None,
) → None#
train_microbatch(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
) → None#

One DP slice of data → one forward_backward_func invocation.

Wrapped in self.model.no_sync() so the mcore DDP hooks accumulate param.main_grad locally on each rank without dispatching a per-call DP reduce. The single true reduce is done explicitly in finish_train_step. Returns nothing: gradients land in param.main_grad and per-microbatch metrics accumulate in the open-step state until finish_train_step surfaces them.

_train_microbatch_body(
state: dict[str, Any],
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
) → None#
finish_train_step() → dict[str, Any]#
_finish_train_step_body(
state: dict[str, Any],
) → dict[str, Any]#
abort_train_step() → None#
get_logprobs(
*,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
micro_batch_size: Optional[int] = None,
require_router_replay: bool = True,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.LogprobOutputSpec]#

Get the logprobs of the model for a batch of data.

Uses the configured logprob_batch_size to do microbatching. Input data is assumed to be right-padded. The method internally converts to left-padded format for computation, and returns outputs in right-padded format. If micro_batch_size is provided, it will be used instead of the configured logprob_batch_size.

Returns:

a BatchedDataDict with key “logprobs” and shape [batch_size, sequence_length]. We use the convention that the logprob of the first token is 0 so that the sequence length is maintained. The logprob of input token i is specified at position i in the output logprobs tensor.

_apply_state_dict_to_model(
source_state_dict: dict,
*,
raise_if_key_missing: bool = False,
) → None#

Apply a state dict to self.model in-place.

  • Tensors with matching shape: in-place copy (parameters / buffers).

  • _extra_state keys (e.g. FP8 scale/amax) with shape mismatch or non-Tensor value: resolve the submodule and call set_extra_state(); supports DDP and Float16Module unwrap.

Parameters:

  • source_state_dict – State dict to apply (e.g. reference_state_dict or saved model_state_dict).

  • raise_if_key_missing – If True, raise when a key in self.model.state_dict() is missing from source_state_dict; if False, skip such keys.

use_reference_model()#

Context manager that temporarily swaps the reference model and active model.

On entry: Moves model to CPU, moves reference_model to CUDA. Swaps the references. Also disables top-k/top-p filtering since the reference policy’s distribution is different from the current policy, making filtered logprobs incompatible. On exit: Restores original references and re-flips cuda/cpu, restores sampling_params.

get_topk_logits(
*,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
k: int,
micro_batch_size: Optional[int] = None,
)#

Get the top-k logits and indices for a batch of data.

The major difference from get_logprobs is that we compute top-k logits and indices for each position in the sequence.

Returns:

  • topk_logits: Tensor of top-k logits for each position in the sequence

  • topk_indices: Tensor of top-k indices for each position in the sequence

Return type:

BatchedDataDict containing

prepare_refit_info() → None#

Prepare state dict metadata for weight refitting and IPC streaming.

_collect_mtp_metrics(metrics: dict[str, Any]) → None#

Add Multi-Token Prediction metrics to metrics when MTP is enabled.

get_mtp_metrics is imported lazily (not a module global) so cloudpickle does not pull an unpicklable torch ConfigModuleInstance into the worker actor’s serialization.

_set_mtp_grad_scale_func(func)#

Set mtp_grad_scale_func on the model config for MTP loss scaling.

_get_model_config()#

Get the underlying model config (handle Float16Module wrapper).

_calculate_refit_param_info() → list[tuple[str, int]]#

Calculate parameter information for refit.

Each task contains:

  • param_name: Local parameter name without module prefixes

  • mapping: MegatronParamMapping instance for weight transformation

  • pp_rank: Pipeline-parallel rank owning the parameter

  • vp_stage: Virtual-pipeline stage index

  • megatron_module: Reference to Megatron model/submodule

  • param_weight: Target parameter tensor for converted weight

Returns:

List of (parameter_name, size_in_bytes) tuples.

_iter_params_with_optional_kv_scales(
kv_scales: Optional[dict[str, float]] = None,
) → Iterator[tuple[str, torch.Tensor]]#

Yield exported HF parameters and optionally append FP8 KV/Q scale tensors.

This helper is used by both IPC-based streaming and collective broadcast so that the logic for adding KV scales stays consistent in one place.

stream_weights_via_ipc_zmq(
buffer_size_bytes: int = 0,
kv_scales: Optional[dict[str, float]] = None,
) → None#

Stream model weights to peer process via ZMQ IPC socket.

broadcast_weights_for_collective(
kv_scales: Optional[dict[str, float]] = None,
) → None#

Broadcast the weights for collective communication.

_use_real_quant_refit() → bool#
prepare_for_lp_inference()#
prepare_for_training(*args, **kwargs)#
finish_inference() → None#

Offload model params to CPU after inference. Only used in PPO.

_clear_fp8_caches()#

Clear FP8 workspace caches and release fragmented GPU memory.

The main memory issue in the train→offload→refit→generate cycle is CUDA allocator fragmentation, not leaked FP8 tensors. This method clears per-module _fp8_workspaces buffers (scratch memory references). The caller is responsible for running gc.collect() + empty_cache() once all references have been dropped.

For anti-fragmentation, configure PYTORCH_CUDA_ALLOC_CONF in the recipe YAML:

  • “max_split_size_mb:512” — fast, prevents large-block splitting

  • “expandable_segments:True” — most effective but ~5x slower weight transfer

offload_before_refit()#

Offload the optimizer and buffers to the CPU.

offload_after_refit()#

Offload as much as possible on the CPU.

move_model(
model: torch.nn.Module,
device: str,
move_params: bool = True,
move_grads: bool = True,
) → torch.nn.Module#
move_optimizer(device: str)#
save_checkpoint(
weights_path: str,
optimizer_path: Optional[str] = None,
**kwargs,
)#

Save a training checkpoint.

With async_save=True, this method returns after D2H staging. The actual disk write continues in a background persistent worker process. Callers must call finalize_async_save() before renaming the directory or starting another save.

With async_save=False (default), this blocks until the write is complete.

Parameters:

  • weights_path – The specific directory path where the checkpoint will be saved.

  • optimizer_path – If not None, optimizer and scheduler states are saved if they exist.

finalize_async_save()#

Block until the in-flight async write completes and run finalize_fns.

Safe to call when async_save is disabled (no-op). Does NOT terminate the persistent worker — it stays alive for the next save.

abstractmethod load_checkpoint(
weights_path: str,
optimizer_path: Optional[str] = None,
)#

Load a training checkpoint.

Parameters:

  • weights_path – The exact directory path from which to load the checkpoint.

  • optimizer_path – If not None, attempts to load optimizer and scheduler states if self.optimizer and self.scheduler are initialized.

check_tensor_parallel_attributes() → dict[str, Any]#

Check tensor parallel attributes on model parameters.

Returns:

  • tp_params: List of parameter names that have tensor_model_parallel=True

  • non_tp_params: List of parameter names that have tensor_model_parallel=False

  • total_params: Total number of parameters checked

  • tp_size: Tensor parallel size from config

Return type:

Dictionary containing information about tensor parallel parameters

calibrate_qkv_fp8_scales(
*,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
micro_batch_size: Optional[int] = None,
percentile: float = 99.9,
margin: float = 1.05,
include_q: bool = False,
) → dict[str, Any]#

One-shot calibration of Q/K/V activation scales (for FP8 KV cache).

  • Captures each layer’s query_key_value output through forward hooks, splits Q/K/V, and computes percentile amax.

  • In parallel (DP/TP/PP) environments, first computes local percentiles, then takes max across all ranks for conservativeness.

  • By default only returns and saves K/V scales, optionally returns Q.

Parameters:

  • data – Representative sample batch for calibration, following get_logprobs input conventions.

  • micro_batch_size – Micro batch size during calibration; if None, reuses logprob_batch_size.

  • percentile – Percentile for amax (e.g. 99.9).

  • margin – Margin factor, e.g. 1.05.

  • save_path – If provided, rank0 will save results as JSON.

  • include_q – Whether to also return Q scale (usually only K/V needed).

Returns:

“fp8”, “percentile”: float, “margin”: float, “layers”: { layer_name: {“k_scale”: float, “v_scale”: float[, “q_scale”: float] } } }

Return type:

{ “format”

class nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker(
config: nemo_rl.models.policy.PolicyConfig,
tokenizer: nemo_rl.models.policy.workers.megatron_policy_worker.TokenizerType,
weights_path: Optional[str] = None,
optimizer_path: Optional[str] = None,
init_optimizer: bool = True,
init_reference_model: bool = True,
*,
worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding,
**kwargs: Any,
)#

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