`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.
`_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_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.
`_collect_mtp_hf_layer_names`	Return HF layer names whose weights originate from Megatron’s MTP module.
`_meta_tensor_alloc_context`	Skip real GPU work during metadata enumeration.

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

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,

skip_weight_load: bool = False,

reserved_http_server_port: Optional[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

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

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

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

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

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

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

Build the conversion-task list driving refit (BF16 or FP8 export).

For BF16 / FP8-but-fp8_param=False training: standard get_conversion_tasks. For FP8-with-fp8_param=True: Bridge’s build_export_fp8_tasks, which emits a pair of tasks per FP8 weight (the FP8 data and a *_scale_inv scale tensor).

_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, 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, torch.Tensor]]#

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 shards directly from the Megatron params — no collectives. Returned tensors are views and must not be modified in place. EP: refit_conversion_tasks already 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, ) → 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, ) → 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_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.

_build_expert_groups(param_map)#

Group this rank’s local expert params into stack-ready views.

Keyed by (prefix, proj_type) and resolved to ordered param_map views ready for torch.stack.

Megatron exposes each expert’s projection as a separate param; this bins them so _group_experts 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): [tensor_0, tensor_1, ...]} — the per-expert param_map views sorted by expert index. Example — a layer with 2 local experts (gated MoE) yields three keys: (".../experts", "gate_proj"): [view(expert 0), view(expert 1)] (".../experts", "up_proj") : [view(expert 0), view(expert 1)] (".../experts", "down_proj"): [view(expert 0), view(expert 1)]

Resolving names → views here (rather than per refit in _group_experts) costs nothing extra — param_map already owns these views and they stay valid across refits (weights are updated in place; the name→view mapping is stable), so _group_experts only has to torch.stack. The index sort matters: the views are stacked in this order, so expert 0 must precede expert 1 to match the EP Shard(0) layout the gen side expects.

_group_experts(proj, grouped_name, expert_groups)#

Stack this rank’s local experts for one projection into [E_local, ...].

Using the pre-calculated expert_groups (from _build_expert_groups) it is just calling torch.stack of all the local expert params.

build_hf_to_local_param_map( refit_info: dict, ) → 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 is sharded local tensor view, sent as-is.
grouped MoE expert: pre stacks the per-expert views into [E_local, ...] fresh each refit via _group_experts.

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

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

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.

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

*,

worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding,

skip_weight_load: bool = False,

reserved_http_server_port: Optional[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`#