`nemo_rl.models.policy.workers.dtensor_policy_worker_v2`#

Module Contents#

Classes#

`DTensorPolicyWorkerV2Impl`
`DTensorPolicyWorkerV2`

Functions#

`_refit_tensor_dtype`	Preserve the FP32 dtype used by inference-critical MoE router state.
`dtensor_params_generator`	Generator that yields (name, tensor) pairs, converting DTensors to local tensors and adapting to HF format.
`_maybe_merge_lora_weight`
`_maybe_adapt_tensor_to_hf`

API#

nemo_rl.models.policy.workers.dtensor_policy_worker_v2._refit_tensor_dtype( fqn: str, tensor: torch.Tensor, default_dtype: torch.dtype, ) → torch.dtype#: Preserve the FP32 dtype used by inference-critical MoE router state.

nemo_rl.models.policy.workers.dtensor_policy_worker_v2.dtensor_params_generator( model: torch.nn.Module, target_dtype: torch.dtype, ) → Generator[tuple[str, torch.Tensor], None, None]#

Generator that yields (name, tensor) pairs, converting DTensors to local tensors and adapting to HF format.

Parameters:

model – The model whose parameters to generate.
target_dtype – The default dtype for refit tensors. Source-FP32 e_score_correction_bias tensors retain FP32.

Yields:

Tuples of (fully_qualified_name, tensor) where tensors are converted to the refit dtype and made contiguous.

nemo_rl.models.policy.workers.dtensor_policy_worker_v2._maybe_merge_lora_weight( module_map: dict[str, torch.nn.Module], fqn: str, tensor: torch.Tensor, ) → torch.Tensor#

nemo_rl.models.policy.workers.dtensor_policy_worker_v2._maybe_adapt_tensor_to_hf( model_part: torch.nn.Module, fqn: str, tensor: torch.Tensor, quantization: bool = False, ) → list[tuple[str, torch.Tensor]]#

class nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2Impl(

config: nemo_rl.models.policy.PolicyConfig,

weights_path: Optional[str] = None,

optimizer_path: Optional[str] = None,

init_optimizer: bool = True,

init_reference_model: bool = True,

**kwargs: Any,

)#

Bases: nemo_rl.data_plane.worker_mixin.TQWorkerMixin, nemo_rl.models.policy.workers.checkpoint_engine.DTensorCheckpointEngineSendMixin, nemo_rl.models.policy.workers.checkpoint_engine.PolicyCheckpointEngineMixin, nemo_rl.models.policy.workers.base_policy_worker.AbstractPolicyWorker, nemo_rl.models.policy.interfaces.ColocatablePolicyInterface

__repr__() → str#

Customizes the actor’s prefix in the Ray logs.

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

_get_replica_group() → Optional[Any]#: Replica group = flattened (cp, tp) sub-mesh — see V1 worker.

_local_coords() → dict[str, int]#

_update_moe_gate_bias_if_supported() → None#: Update the non-gradient MoE routing bias after the optimizer step.

_autocast_context() → contextlib.AbstractContextManager[Any]#: Return the worker-owned precision context for one microbatch.

set_rollout_num_gpus_per_engine(num_gpus_per_engine: int) → None#: Record the rollout engine’s TP size for later use in stream_weights_via_http.

train( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], 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.

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

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.

score( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.ScoreOutputSpec]#

get_topk_logits( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], k: int, micro_batch_size: Optional[int] = None, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#

Return per-position top-k logits and corresponding global indices.

Notes:

Return shapes are [B, S, k].
Computes top-k over the full sequence (no trimming of the last position).
If alignment with next-token targets is required, the caller should handle it.
If logits are TP-sharded DTensor, performs distributed global top-k across TP.
Supports context parallelism with proper CP gather.
Otherwise, computes local top-k on full-vocab tensor.

get_full_logits_ipc( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], micro_batch_size: Optional[int] = None, ) → dict[str, Any]#

Teacher forward; full-vocab logits exposed via persistent CUDA IPC storage.

Used by cross-tokenizer distillation; supports heterogeneous teacher TP/CP. Each microbatch writes into slot self._teacher_ipc_storage[buf_idx] and shares one cached IPC handle. Returns {"per_sample_handles": list, "dp_rank": int} where each handle carries buf_idx and sample_index_in_buf for the consumer to index the slot view, plus the TP/CP shard metadata (vocab_start_index, global_seq_start, …) the consumer uses to route shards across heterogeneous teacher/student TP/CP.

release_ipc_buffer() → None#: Free the persistent teacher-logit IPC storage. Called once at end of training/validation.

use_reference_model() → Generator[None, None, None]#

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.

_add_noise_to_weights() → None#: Add small Gaussian noise to the weights of the model. Note that this is used for testing purposes only.

return_state_dict()#

return_model_config() → dict[str, Any]#

Return the model configuration as a dictionary.

Returns:: Model configuration dictionary
Return type:: dict

prepare_refit_info() → Optional[dict[str, Any]]#: Prepare state dict metadata for weight refitting and IPC streaming.

abstractmethod 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]#: Placeholder for FP8 Q/K/V scale calibration, not implemented for DTensorPolicyWorkerV2.

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.

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

Send FSDP weights to colocated SGLang engines via Ray CUDA IPC.

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.

_checkpoint_engine_params() → Generator[tuple[str, torch.Tensor], None, None]#

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.

Guarded exactly as the Megatron worker is, and for the same reason: a generation rank that dies mid-broadcast leaves this call blocked in NCCL with no timeout and no error. 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#

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 optimizer state on CUDA because a train step is already open. This backend accumulates gradients in param.grad and never offloads them, so unlike the Megatron backend there is nothing here that could discard them; the flag only suppresses the per-chunk optimizer round trip.

prepare_for_training(*args, **kwargs) → None#

finish_inference() → None#: Offload model params to CPU after inference. Only used in PPO.

offload_before_refit() → None#: Offload the optimizer to the CPU.

offload_after_refit() → None#: Offload as much as possible on the CPU.

move_optimizer_to_device(device: str | torch.device) → None#

move_to_device( model: torch.nn.Module, device: str | torch.device, ) → torch.nn.Module#

move_buffer_to_device( model: torch.nn.Module, device: str | torch.device, ) → torch.nn.Module#

move_to_cuda(model: torch.nn.Module) → torch.nn.Module#

move_to_cpu(model: torch.nn.Module) → torch.nn.Module#

save_checkpoint( weights_path: str, optimizer_path: Optional[str] = None, tokenizer_path: Optional[str] = None, checkpointing_cfg: Optional[nemo_rl.utils.checkpoint.CheckpointingConfig] = None, ) → None#

Save a checkpoint of the model.

the optimizer states are saved only if optimizer and optimizer_path are provided.

finalize_async_save() → None#

Block until this worker’s in-flight async checkpoint writes complete.

Overrides the base no-op: this worker initializes the checkpoint manager with is_async=True, so the caller-side rename of tmp_step_N to step_N must wait for the staged writes to land.

load_checkpoint( weights_path: str, optimizer_path: Optional[str] = None, ) → None#: Load a checkpoint into the model using Automodel Checkpointer.

_init_checkpoint_manager( config_updates: Optional[dict[str, Any]] = None, checkpoint_root: Optional[str] = None, ) → None#

Initialize the AutomodelCheckpointManager for this worker.

This creates the checkpoint manager bound to this worker’s device meshes and initializes its underlying checkpointer.