`nemo_rl.models.policy.workers.megatron_policy_worker`#

Module Contents#

Classes#

`_RefitSource`	Live parameter and its HF selector, before any precision conversion.
`MegatronReshardParamSpecCollection`	Factories for Megatron’s canonical HF refit source specifications.
`MegatronPolicyWorkerImpl`
`MegatronPolicyWorker`

Functions#

`_should_use_router_replay`
`_model_self_packs_for_cp`	Whether the model packs sequences + CP-shards inside its own forward.
`_model_self_packs_mtp_loss_mask`	Whether a self-packing model also aligns and CP-shards MTP masks.
`_model_slices_context_parallel_inputs`	Whether the model consumes full THD input and slices CP after embedding.
`_unwrapped_chunks`	Model chunks as a flat list, whatever wrapping the caller handed us.
`_model_has_chunkwise_cp`
`_model_needs_router_padding_mask`	Resolve masks once for nonzero, unfrozen expert-bias updates.
`_model_media_placeholder_token_id`	The vocabulary id this model treats as a media placeholder, if any.
`_model_accepts_media_token_validity_mask`	Whether the model’s forward takes an explicit media-token validity mask.
`_estimate_refit_tensor_size_in_bytes`	Estimate the gathered tensor size produced by Bridge export.
`_is_quantized_refit_source`	Whether a training parameter uses TE quantized physical storage.
`_dequantize_refit_source_group`	Decode primary expert weights, grouping block-aligned storage without a cache.
`_dequantize_refit_source`	Decode quantized storage to BF16, preserving unquantized source dtype.
`_get_refit_task_source`	Return the original TE parameter when Bridge exposes an export payload.
`_collect_mtp_hf_layer_names`	Return HF layer names whose weights originate from Megatron’s MTP module.
`_collect_local_refit_hf_names`	Return HF names whose Bridge tasks expose safe direct local views.
`_meta_tensor_alloc_context`	Skip real GPU work during metadata enumeration.
`_reserved_http_server_port_for_rank`	Return this worker rank’s reserved frontend port, if one was assigned.

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. New wrappers advertise the capability through model_owns_packing. The Qwen3VL type check remains as a compatibility fallback until that upstream model exposes the capability.

nemo_rl.models.policy.workers.megatron_policy_worker._model_self_packs_mtp_loss_mask(model: Any) → bool#: Whether a self-packing model also aligns and CP-shards MTP masks.

nemo_rl.models.policy.workers.megatron_policy_worker._model_slices_context_parallel_inputs(model: Any) → bool#: Whether the model consumes full THD input and slices CP after embedding.

nemo_rl.models.policy.workers.megatron_policy_worker._unwrapped_chunks(model: Any) → list[Any]#: Model chunks as a flat list, whatever wrapping the caller handed us.

nemo_rl.models.policy.workers.megatron_policy_worker._model_has_chunkwise_cp(model: Any) → bool#

nemo_rl.models.policy.workers.megatron_policy_worker._model_needs_router_padding_mask( model: Any, *, pack_sequences: bool = False, allow_unmasked_chunkwise_cp: bool = False, ) → bool#: Resolve masks once for nonzero, unfrozen expert-bias updates.

nemo_rl.models.policy.workers.megatron_policy_worker._model_media_placeholder_token_id( model: Any, ) → Optional[int]#: The vocabulary id this model treats as a media placeholder, if any.

nemo_rl.models.policy.workers.megatron_policy_worker._model_accepts_media_token_validity_mask(model: Any) → bool#: Whether the model’s forward takes an explicit media-token validity mask.

nemo_rl.models.policy.workers.megatron_policy_worker._estimate_refit_tensor_size_in_bytes( param: torch.Tensor, *, export_dtype: torch.dtype, tp_size: int, ep_size: int, ) → int#

Estimate the gathered tensor size produced by Bridge export.

Floating-point model weights are exported at the policy dtype. Integral state (for example BatchNorm num_batches_tracked buffers) keeps its original dtype and must not be looked up in a floating-point-only table.

nemo_rl.models.policy.workers.megatron_policy_worker._is_quantized_refit_source(tensor: torch.Tensor) → bool#: Whether a training parameter uses TE quantized physical storage.

class nemo_rl.models.policy.workers.megatron_policy_worker._RefitSource#

Live parameter and its HF selector, before any precision conversion.

tensor: torch.Tensor#: None

hf_spec: megatron.bridge.models.conversion.param_mapping.LocalHFParamSpec#: None

nemo_rl.models.policy.workers.megatron_policy_worker._dequantize_refit_source_group( tensors: tuple[torch.Tensor, ...], ) → torch.Tensor#

Decode primary expert weights, grouping block-aligned storage without a cache.

TE primary weights retain rowwise data and compact scales. Pack their logical scale entries, including only aggregate metadata padding; never pad the weights or change an expert’s quantization block boundaries. Callers provide a uniform group with matching shapes, dtypes and recipes.

nemo_rl.models.policy.workers.megatron_policy_worker._dequantize_refit_source(tensor: torch.Tensor) → torch.Tensor#: Decode quantized storage to BF16, preserving unquantized source dtype.

nemo_rl.models.policy.workers.megatron_policy_worker._get_refit_task_source( task: Any, ) → Optional[torch.Tensor]#: Return the original TE parameter when Bridge exposes an export payload.

class nemo_rl.models.policy.workers.megatron_policy_worker.MegatronReshardParamSpecCollection#

Factories for Megatron’s canonical HF refit source specifications.

static _default_param_spec( tensor: torch.Tensor, hf_spec: megatron.bridge.models.conversion.param_mapping.LocalHFParamSpec, ) → nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#: Convert the current full parameter to logical weights before HF slicing.

static _grouped_expert_spec( proj: str, grouped_name: str, expert_groups: dict[tuple[str, str], list[nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec]], ) → nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec#

Prepare this rank’s experts as [E_local, *hf_shard_shape].

pre decodes [E_local, *source_shard_shape] before selecting the HF projection, including the gate/up split. Transfer metadata describes [E_global, *hf_global_shape] so NCCL can reshard these local buffers.

nemo_rl.models.policy.workers.megatron_policy_worker._collect_mtp_hf_layer_names( conversion_tasks: Optional[list], ) → set[str]#

Return HF layer names whose weights originate from Megatron’s MTP module.

This is required because, in some cases, only the Megatron-side name contains the mtp string, while the HF-side name does not.

Parameters:: conversion_tasks – Megatron-Bridge WeightConversionTask list
Returns:: Set of HF layer names, e.g. {"model.layers.61", "mtp.layers.0"}.

nemo_rl.models.policy.workers.megatron_policy_worker._collect_local_refit_hf_names( conversion_tasks: Iterable[Any], ) → set[str]#

Return HF names whose Bridge tasks expose safe direct local views.

Every task contributing to a grouped HF tensor must provide local specs. Mappings that need transpose, interleave, or other grouped-export transforms return no specs and therefore stay on the normal Bridge conversion path.

The result is expert-parallel invariant. Bridge conversion tasks are EP-local – each EP rank only enumerates its own experts.N.* – while the name stream this set is matched against is EP-global, so a purely local set would classify the same expert differently on different EP ranks. The support map is therefore AND-reduced across the EP group before filtering.

nemo_rl.models.policy.workers.megatron_policy_worker._meta_tensor_alloc_context()#

Skip real GPU work during metadata enumeration.

Bridge’s export_hf_weights does PP/TP/EP gathers to materialize full unsharded tensors, but the refit-info builders only need shape+dtype. Patch the allocators to redirect to meta and turn the collectives into no-ops. Subsequent shape-only ops on meta tensors propagate correctly, while peak memory stays at zero extra GiB.

nemo_rl.models.policy.workers.megatron_policy_worker._reserved_http_server_port_for_rank( reserved_http_server_ports: Optional[dict[int, int]], rank: int, ) → Optional[int]#: Return this worker rank’s reserved frontend port, if one was assigned.

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,

processor: Optional[Any] = None,

*,

worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding,

skip_weight_load: bool = False,

is_refit_destination: bool = False,

reserved_http_server_ports: Optional[dict[int, int]] = None,

**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.checkpoint_engine.MegatronCheckpointEngineSendMixin, nemo_rl.models.policy.workers.checkpoint_engine.PolicyCheckpointEngineMixin, nemo_rl.models.policy.workers.base_policy_worker.AbstractPolicyWorker, nemo_rl.models.policy.interfaces.ColocatablePolicyInterface, nemo_rl.weight_sync.nccl_reshard_utils.RefitBuilderInterface

is_refit_destination: bool#: False

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

_remote_sparse_refit: Any#: None

_async_checkpoint_cuda_cache_active: bool#: False

__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]#

_routed_experts_dimensions() → tuple[int, int]#: Return route dimensions from the initialized Megatron model config.

_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, num_gpus_per_node: Optional[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
num_gpus_per_node – Per-node GPU count (unused here; part of the shared configure_worker contract).

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( *, param_sync: bool = False, ) → None#

_restore_first_train_step_param_sync(update_successful: bool) → None#

_copy_main_params_to_param_buffer( zero_grad_buffer: bool = False, ) → None#

Stage FP32 optimizer masters into the shared MXFP8 param buffer.

Only valid while the forward pre-hooks are still installed: this stages the buffer but does not all-gather it, and callers rely on the hook-driven (or explicit) param sync to publish the result. Staging with hooks already off would zero the aliased param/grad buffer and leave every other DP rank’s slice at zero.

Parameters:: zero_grad_buffer – Also zero the aliased gradient buffer first.

_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#

_normalize_in_loss_seq_filter( loss_fn: nemo_rl.algorithms.loss.loss_functions.ClippedPGLossFn, losses_reduced: list[dict[str, Any]], *, global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, eval_mode: bool, ) → tuple[list[dict[str, Any]], torch.Tensor, torch.Tensor]#

Normalize one complete optimizer batch after in-loss filtering.

When we use seq_logprob_error_threshold with seq_logprob_error_in_loss, we only have access to the global valid token count. After the full batch is computed.

Therefore, we need to normalize the gradient values by the global valid token count. All microbatches have now finished, so sum survivor counts over DP, broadcast them to every PP stage, and correct gradients before the optimizer clips them. CP/TP replicas must not be counted as additional samples.

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 Automodel worker (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]#

_log_gpu_mem(tag: str) → None#

Emit per-rank CUDA allocator counters for memory forensics.

The max_* figures are interval peaks: the peak counters are reset on every call, so each line reports the high-water mark since the previous boundary rather than since process start. driver_used comes from the driver and therefore includes NCCL’s own device allocations, which the torch counters do not see.

Resetting the peak counters is observable to anything else reading them, so skip the whole body when the level is off rather than sampling and discarding. Run with NRL_LOG_LEVEL=DEBUG to turn these on.

Parameters:: tag – Phase-boundary name this sample belongs to.

_restore_saved_mcore_hooks(state: dict[str, Any]) → None#

Restore the mcore hooks nulled in begin_train_step.

Restores grad_sync_func, no_sync_func and finalize_model_grads_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.

Multimodal validity-mask and model-owned packing/CP behavior match the regular train path.

_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.

_resolve_output_layer_owner() → Optional[Any]#

Return the unwrapped module owning output_layer, or None off the last PP stage.

Megatron builds output_layer only on the last pipeline stage, and that is also the only stage where the opd_full loss runs, so the earlier stages legitimately have nothing to project with. They still resolve the checkpoint through Megatron-Bridge’s read_train_state, a whole-world broadcast; see _load_opd_full_teacher_lm_head_from_path.

Returns:: The module owning output_layer, or None when this rank is not the last pipeline stage.
Raises:: AttributeError – If this rank is the last pipeline stage – or runs without torch.distributed, where pipeline placement does not apply – and still has no output layer. That is a malformed model rather than a pipeline-placement consequence.

load_opd_full_teacher_lm_head( teacher_path_config: nemo_rl.models.policy.PolicyConfig, teacher_index: int, ) → str#

Resolve and load one teacher’s LM head for full-vocabulary MOPD.

Called once per unique teacher checkpoint after the teacher worker groups exist, because each teacher’s Megatron checkpoint is only materialized by its own HF conversion. Path resolution happens here rather than on the driver: the driver process is never provisioned with the mcore extra that validate_model_paths’ module needs.

Parameters:

teacher_path_config – The teacher group’s own policy config, which carries no pretrained_checkpoint, so resolution keys off the teacher’s model name.
teacher_index – Stable per-checkpoint index (see create_teacher_worker_groups) this shard is stored under.

Returns:

The resolved Megatron checkpoint root of the teacher.

_load_opd_full_teacher_lm_head_from_path( teacher_pretrained_path: str, teacher_index: int, ) → None#

Load this rank’s shard of one teacher’s LM head from an already-resolved path.

Under student pipeline parallelism only the last stage owns an output_layer, and only there does the opd_full loss run. The earlier stages request no shard, but must still call the loader because resolving the teacher’s checkpoint iteration (Megatron-Bridge’s read_train_state) broadcasts over the whole student world – skipping it would hang the last stage. They record the checkpoint path all the same, so the evict lifecycle re-enters that broadcast in lockstep with the stage that actually reloads a shard.

Raises:: ValueError – If this teacher’s hidden size disagrees with a teacher already loaded on this rank.

_move_opd_full_teacher_lm_head(device: str) → None#: Move every cached teacher LM-head shard between CPU and GPU.

_release_opd_full_teacher_lm_head() → None#

Apply the configured lifecycle policy after a training phase.

Gated on the recorded checkpoint paths rather than on the shards themselves: every rank that took part in the load has a path, but only the last pipeline stage has a shard, and the eviction flag has to advance on all of them or the reload collective goes out of lockstep.

_stage_opd_full_teacher_lm_head_for_training() → None#: Make every teacher LM-head shard resident on GPU for a train step.

get_logprobs_with_full_payload( *, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], payload: str, payload_dtype: str, micro_batch_size: Optional[int] = None, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.TeacherFullPayloadOutputSpec]#

Run a teacher forward returning sampled-token logprobs and the full payload.

Full-vocabulary MOPD needs the teacher’s whole next-token distribution. Both outputs come from a single forward: adding a second entrypoint would double the teacher’s cost for the same tokens.

Unlike get_logprobs, the payload is deliberately not broadcast off the last pipeline stage – it is orders of magnitude larger than the scalar logprob column, and the caller writes it back from the stage that already owns it.

Parameters:

data – Right-padded batch to score.
payload – "hidden_states" or "logits".
payload_dtype – Torch dtype name for the transported payload.
micro_batch_size – Overrides the configured logprob batch size.

Returns:

A BatchedDataDict with logprobs [B, S] on every rank and teacher_full_payload [B, S, D] on CPU on the last pipeline stage (None elsewhere).

Raises:

ValueError – If payload is "hidden_states" but this teacher transforms its logits after output_layer, which the student’s reconstruction cannot reproduce.

_apply_state_dict_to_model( source_state_dict: dict, *, raise_if_key_missing: bool = False, non_blocking: 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.
non_blocking – Passed to the in-place copies; callers staging from pinned CPU memory must synchronize afterwards.

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_source_refit_info() → dict[str, tuple[torch.Size, torch.dtype]]#: Prepare state dict metadata for weight refitting and IPC streaming.

prepare_refit_info( state_dict_info: Optional[dict[str, Any]] = None, ) → Optional[dict[str, tuple[torch.Size, torch.dtype]]]#: Prepare refit state for this worker’s explicit source/destination role.

enable_refit_prequantize( param_names: list[str], ) → dict[str, Any]#

Quantize the listed HF params to MXFP8 on the trainer during refit.

Derives the updated metadata from the shapes cached by prepare_refit_info instead of re-running the export + quantize pass: MXFP8 keeps the value shape and adds one uint8 E8M0 scale per 32-wide block along the last dim, so no tensor needs to be materialized here.

Parameters:: param_names – fp8-eligible parameter names reported by the vLLM workers (see VllmInternalWorkerExtension.prepare_refit_info).
Returns:: the listed params become float8_e4m3fn and each gains a *_scale_from_checkpoint uint8 entry.
Return type:: Updated refit metadata

_maybe_prequantize_param( name: str, tensor: torch.Tensor, ) → Iterator[tuple[str, torch.Tensor]]#

async update_weights_from_collective( refit_timeout_s: Optional[float] = None, ) → bool#: Receive a collective refit without blocking this actor’s event loop.

_update_destination_weights_from_collective_guarded( refit_timeout_s: Optional[float], ) → bool#: Receive packed Megatron weights under the shared refit deadline.

_collect_mtp_metrics( metrics: dict[str, Any], total_num_microbatches: int, mtp_grad_norm: Optional[float], ) → 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.

Parameters:

metrics – Metrics dict to populate with MTP metrics (under “mtp_metrics”).
total_num_microbatches – Microbatches accumulated this step. The MTP loss logging helper sums the per-microbatch loss without dividing, so we pass 1/total_num_microbatches to recover the mean (mirroring the MoE path).
mtp_grad_norm – The MTP parameter group’s gradient norm, already reduced across the model-parallel group, or None when unavailable (e.g. clip_grad == 0 or mtp_detach_heads=False). Logged under “mtp_metrics” as “grad_norm”.

_set_mtp_grad_scale_func(func)#: Set mtp_grad_scale_func on the model config for MTP loss scaling.

_scale_mtp_param_grads(factor: float) → None#

Scale detached MTP parameters’ gradients by factor.

MCore tags every MTP parameter grad_norm_group='mtp' when mtp_detach_heads is enabled. In that configuration the auxiliary loss reaches no shared parameters, so its denominator can be corrected independently. main_grad is a view into the DDP gradient buffer.

Parameters:: factor – Multiplier for MTP gradients. 1.0 is a no-op.

_get_model_config()#: Get the underlying model config (handle Float16Module wrapper).

init_remote_sparse_delta_baseline( *, shard_rank: int, shard_count: int, transport: str, ) → dict[str, tuple[tuple[int, ...], torch.dtype]]#

stream_remote_sparse_weights( transport: str, targets: list[str], *, transfer_id: str, api_key_env_var: Optional[str], timeout_s: float, shard_rank: int, shard_count: int, overwrite_names: list[str], ) → dict[str, int]#

_require_remote_sparse_refit() → Any#

finish_remote_sparse_delta_sync(*, succeeded: bool) → None#

_is_fp8_export() → bool#: Return True if the train side stores weights as TE blockwise FP8.

_build_refit_conversion_tasks() → list#: Keep the general/misc export format expected by each backend.

_generation_needs_logical_weights() → bool#: Megatron’s existing non-bulk export consumes logical model weights.

_build_bulk_refit_conversion_tasks() → list[Any]#: Use standard Bridge mappings and live parameters for bulk refit.

_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.

static _iter_logical_refit_conversion_tasks( conversion_tasks: Iterable[Any], ) → Iterator[Any]#: Yield Bridge tasks whose quantized sources are materialized as BF16.

_nvfp4_pertoken_rollout_cfg() → nemo_rl.models.generation.vllm.quantization.nvfp4_pertoken_config.NvFp4PerTokenRolloutConfig | None#: Return validated per-token rollout config when the mode is enabled.

maybe_init_zmq() → None#: Allow extra time for the first quantized refit and kernel autotune.

_iter_params_with_optional_kv_scales( kv_scales: Optional[dict[str, float]] = None, conversion_tasks=None, include_draft: bool = True, ) → 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.

conversion_tasks (optional) overrides self.refit_conversion_tasks — used by the nccl_reshard_refit misc-refit path to pass a filtered subset so Bridge only does TP/EP all-gather for those tasks instead of the full model.

include_draft controls the draft.* EAGLE weights. SGLang refit sets it False: the engine keeps draft weights via enable_draft_weights_cpu_backup rather than receiving them.

_iter_local_hf_param_shards() → Iterator[tuple[str, nemo_rl.weight_sync.nccl_reshard_utils.LocalParamSpec]]#

Yield (hf_name, local_tp_shard) for this rank’s locally owned FFN params.

Used by the nccl_reshard_refit bulk path (build_hf_to_local_param_map). Only the FFN projections (gate/up/down_proj) take the bulk path, so this yields ONLY those. Others. take the misc packed_broadcast path and are skipped here (see is_nccl_reshard_param).

Unlike _iter_params_with_optional_kv_scales (PP broadcast + TP gather via export_hf_weights), this yields TP-local source specs directly from the Megatron params — no collectives. Every pre hook independently materializes the current parameter in canonical HF layout, using BF16 for quantized storage and retaining unquantized dtypes. EP: _bulk_refit_conversion_tasks holds only this rank’s local experts; PP non-local params have param_weight is None.

_iter_sglang_hf_weight_buckets( *, target_precision: str, sglang_quantization_cfg: Optional[dict] = None, buffer_size_bytes: int, ) → Iterator[list[tuple[str, torch.Tensor]]]#

Yield HF tensor buckets for SGLang refit.

Reuses the same two pieces as every other transport: the export_hf_weights walk in _iter_params_with_optional_kv_scales (without vLLM KV/Q scales or draft weights — SGLang keeps drafts engine-side via enable_draft_weights_cpu_backup) and the shared iter_named_tensor_buckets packing.

update_weights_to_sglang_colocated( *, rollout_engines: list, buffer_size_bytes: int, target_precision: str = 'bf16', sglang_quantization_cfg: Optional[dict] = None, ) → None#

Send finalized HF tensor buckets to colocated SGLang engines.

Synchronous: each chunk is awaited via ray.get inside

Func:: send_hf_buckets_via_ipc_actor_impl before the next chunk is sent, so trainer-side IPC tensors stay alive until the engine has copied them and per-chunk engine failures surface immediately. Raises RuntimeError on any chunk failure.

connect_sglang_rollout_engines_distributed( *, rollout_engines: list, engine_gpu_counts: list[int], group_name: Optional[str] = None, ) → None#

Bring up the trainer-rank-0 NCCL group for SGLang disaggregate refit.

Only trainer rank 0 broadcasts to SGLang, so only rank 0 owns the torch process group. Other ranks return immediately. Calling this again after engines recover destroys the stale group first.

update_weights_to_sglang_distributed( *, rollout_engines: list, rollout_engine_lock, buffer_size_bytes: int, target_precision: str = 'bf16', sglang_quantization_cfg: Optional[dict] = None, ) → None#

Broadcast finalized HF tensors to SGLang engines from trainer rank 0.

Non-rank-0 trainers still walk the AutoBridge iterator (Megatron gather + AutoBridge restoration is a collective), but they do not participate in the NCCL broadcast. This matches the design’s “trainer rank 0 as the only source” decision.

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, refit_timeout_s: Optional[float] = None, *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, use_legacy_packing: bool = False, ) → None#

Broadcast the weights for collective communication.

A generation rank that dies mid-broadcast leaves this call blocked in NCCL with no timeout and no error – observed as both policy workers stuck in packed_broadcast_producer -> cuda stream synchronize while the run sat wedged. The watchdog is the only way out, because the controller cannot reach this actor while its event loop is inside the collective. Disarmed unless refit_timeout_s is set, so the default path is unchanged.

_broadcast_weights_for_collective( kv_scales: Optional[dict[str, float]] = None, *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, use_legacy_packing: bool = False, ) → None#

_build_layer_to_pp_stage( pp_size: int, layer_prefix: str, ) → dict[str, int]#

Build mapping from layer group name to PP stage index.

Returns a dictionary that maps the layer group name to the PP stage index. layer_prefix is the module path before layers.N in the exported HF names (e.g. model, model.language_model, backbone)

Mirrors Megatron-LM’s get_num_layers_to_build (transformer_block.py) for the standard (non-VP, non-custom-layout) path: middle stages share num_layers - first - last evenly, while the first/last stages get their explicit counts when set.

Cases not yet supported are asserted out so failures are loud rather than silently producing wrong layer→stage mappings:

pipeline_model_parallel_layout (e.g. DeepSeek-V3)
virtual_pipeline_model_parallel_size (interleaved PP)
account_for_embedding_in_pipeline_split
account_for_loss_in_pipeline_split These cases are checked in check_nccl_reshard_refit_support function.

_prepare_source_nccl_reshard_refit_info( train_parallelism, gen_parallelism, train_world_size, gen_world_size, )#

Prepare per-layer parameter metadata for nccl_reshard-based refit.

The builder groups per-expert MoE params into backend-agnostic grouped HF entries (gate_proj/up_proj/down_proj); the gen backend maps those into its own fused layout (e.g., vLLM w13/w2) gen-side, so this train worker stays agnostic to any gen backend’s MoE-fusion layout.

prepare_nccl_reshard_refit_info( train_parallelism: Optional[dict[str, Any]] = None, gen_parallelism: Optional[dict[str, Any]] = None, train_world_size: Optional[int] = None, gen_world_size: Optional[int] = None, *, refit_info: Optional[dict[str, Any]] = None, ) → Optional[dict[str, Any]]#: Prepare NCCL-reshard state for the worker’s explicit refit role.

_build_expert_groups(param_map)#

Group this rank’s live source specs in numeric expert order.

Keyed by (prefix, proj_type) and resolved to ordered param_map specs whose raw parameters are packed and decoded together in the group’s pre hook at each refit.

Megatron exposes each expert’s projection as a separate param; this bins them so MegatronReshardParamSpecCollection._grouped_expert_spec can stack a layer’s experts into one grouped HF tensor per projection. Called from build_hf_to_local_param_map with this rank’s local param_map.

_INDIVIDUAL_EXPERT_RE captures three fields from a name like model.layers.3.mlp.experts.17.gate_proj.weight:

group 1 = prefix -> "model.layers.3.mlp.experts"
group 2 = expert index -> 17
group 3 = proj type -> "gate_proj" so the name keys into ("model.layers.3.mlp.experts", "gate_proj").

Returns {(prefix, proj): [spec_0, spec_1, ...]} — the per-expert param_map specs sorted by expert index. Example — a layer with 2 local experts (gated MoE) yields three keys: (".../experts", "gate_proj"): [spec(expert 0), spec(expert 1)] (".../experts", "up_proj") : [spec(expert 0), spec(expert 1)] (".../experts", "down_proj"): [spec(expert 0), spec(expert 1)]

Resolving names → specs here costs nothing extra and stays valid across refits because weights are updated in place. The index sort matters: source tensors are stacked in this order, so expert 0 must precede expert 1 to match the EP Shard(0) layout the gen side expects.

_build_source_hf_to_local_param_map( refit_info: dict[str, Any], ) → nemo_rl.weight_sync.nccl_reshard_utils.HFToLocalParamMap#

Build the Megatron-backend hf_to_local_param_map (HFToLocalParamMap).

Wraps this rank’s local Megatron shards into LocalParamSpecs:

direct: base describes the live parameter and HF selector; pre dequantizes if necessary and selects the HF view, preserving unquantized dtypes.
grouped MoE expert: pre packs the ordered live parameters and their scales, dequantizes the group once, then selects its HF projection as [E_local, *hf_shard_shape]. Every spec declares its wire dtype, checked against transfer metadata before any source buffers are materialized.

build_hf_to_local_param_map( refit_info: dict[str, Any], *, destination_tasks: Optional[list[Any]] = None, ) → nemo_rl.weight_sync.nccl_reshard_utils.HFToLocalParamMap#: Build the local map for this worker’s source or destination role.

async nccl_reshard_refit(kv_scales=None, refit_timeout_s=None)#

Run the refit off this actor’s event loop; see _nccl_reshard_refit_guarded.

Async purely so the blocking transfer does not occupy the loop. While it runs, this actor can still service the recovery’s init_collective – which is the whole reason the rebuild can happen at all when a generation rank goes silent.

THE DEVICE IS CARRIED ACROSS EXPLICITLY. CUDA’s current device is thread-local, so a fresh thread starts on device 0 rather than this worker’s, and NCCL on the wrong device fails with UnhandledCudaError – job 6510914 died that way on its first healthy refit, before any fault was injected. torch.cuda.current_stream() inside the transfer reads the same thread-local state, so setting the device also puts the transfer back on the intended stream.

_nccl_reshard_refit_guarded(kv_scales=None, refit_timeout_s=None)#

Transfer weights to generation workers via xferdtensor, under a deadline.

Guarded exactly like the collective producer, and for the same reason: a generation rank that dies mid-refit leaves this blocked inside NCCL with no error and no progress, and the controller cannot reach this actor to break it because its event loop is inside the transfer.

BOTH communicator families are handed to the watchdog. This transport moves the bulk over the per-PP-stage pp_comm_group and then broadcasts the remainder over the shared model_update_group, so the hang can be in either and nothing here can tell which. Aborting both is safe – abort() is idempotent – and the recovery rebuilds both anyway.

Disarmed unless refit_timeout_s is set, so the default path is unchanged.

kv_scales (FP8 KV cache): the per-layer k/v(/q) scales ride the misc packed-broadcast as plain scale tensors (the is_nccl_reshard_param whitelist excludes .k_scale/.v_scale/.q_scale -> misc); the gen side finalizes them via the process_weights_after_loading pass in the weight-update finalizer. No out-of-band channel needed.

_nccl_reshard_refit(kv_scales=None, refit_timeout_s=None)#

_broadcast_misc_params_packed(kv_scales=None) → None#: Broadcast misc params via the existing packed_broadcast machinery.

prepare_for_lp_inference(keep_train_buffers: bool = False) → None#

Put the model in eval mode for logprob inference.

Parameters:

keep_train_buffers –
Leave the grad buffers and the optimizer state on CUDA. Set this when a train step is already open. mcore’s _ParamAndGradBuffer.offload_to_cpu(move_grads=True) does not copy gradients anywhere — it calls grad_data.storage().resize_(0), freeing them — and the matching reload_from_cpu resizes the storage back and zero_()s it. param.main_grad stays a valid view of that storage throughout, so nothing raises: the gradients accumulated by earlier streaming chunks of this step are simply gone, leaving only the last chunk’s contribution against a 1/N normalizer computed over all of them. Keeping the buffers resident also avoids round-tripping tens of GiB per chunk.

Suppressing the offload here is sufficient only because of an mcore invariant on the other side of the detour: prepare_for_training runs before every chunk and reloads with move_grads=True, and reload_from_cpu resizes and zero_()s the grad buffer only if grad_data_size > 0, a counter set only by a matching offload_to_cpu. With the offload suppressed it stays 0, so the reload is a no-op and the accumulated gradients survive. An mcore change that dropped that guard, or that zeroed unconditionally, would silently reinstate this bug.
grad (MXFP8 overlap aliases the parameter all-gather buffer to)
when (storage. That allocation remains resident even)
pre-hooks (keep_train_buffers is false because logprob forward)
all-gather. (may need it for parameter)

_build_colocated_inference_model( config: nemo_rl.models.policy.PolicyConfig, ) → None#: Build the dedicated inference-layout model planned at setup.

pause_train_step_with_offloading() → None#

Pause the step, offloading parameters and optimizer state to CPU.

Pending dense/expert gradients stay on GPU. Counts, metrics and disabled reduction hooks are preserved until prepare_for_training resumes it.

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

sync_params_before_refit() → None#

Materialize optimizer updates before a refit reads model parameters.

overlap_param_gather defers this work to the next training forward, but refit reads the parameters first. Megatron-FSDP handles this in its own module hooks and is intentionally not handled here.

offload_before_refit()#: Offload optimizer state and buffers that are safe to release.

_clear_rope_and_moe_dispatcher_caches() → None#: Clear rotary-embedding and MoE dispatcher caches repopulated by forwards.

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.

_requires_nvrx_cuda_cache_release() → bool#: Whether checkpoint finalization must also drop cached CUDA IPC handles.

finalize_async_save()#

Finalize an async write and release unsafe colocated CUDA IPC caches.

NVRx constant-structure saves cache CUDA tensor handles in the persistent writer. That is safe while model/optimizer storage stays fixed, but a colocated policy replaces that storage during CPU offload. In that case, close the completed writer and invalidate its training-side cache; NVRx starts a fresh persistent writer lazily for the next checkpoint.

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,

processor: Optional[Any] = None,

*,

worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding,

skip_weight_load: bool = False,

is_refit_destination: bool = False,

reserved_http_server_ports: Optional[dict[int, int]] = None,

**kwargs: Any,

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

nemo_rl.models.policy.workers.megatron_policy_worker#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.policy.workers.megatron_policy_worker`#