nemo_rl.models.policy.workers.automodel_policy_worker

Module Contents

Classes

AutomodelPolicyWorkerImpl
AutomodelPolicyWorker

Functions

_refit_tensor_dtype	Preserve the dtype of non-floating-point buffers and FP32 MoE router state.
_module_params_generator	Generator that yields (name, tensor) pairs, converting DTensors to local tensors and adapting to HF format.
automodel_params_generator	Policy weights, followed by draft.<name> weights when co-training.
_assert_lora_modules_reachable	Fail loudly if a LinearLoRA has no matching state_dict weight.
_maybe_merge_lora_weight
_maybe_adapt_tensor_to_hf

API

nemo_rl.models.policy.workers.automodel_policy_worker._refit_tensor_dtype(

fqn: str,

tensor: torch.Tensor,

default_dtype: torch.dtype,

) → torch.dtype

Preserve the dtype of non-floating-point buffers and FP32 MoE router state.

Integer/bool buffers (e.g. the reduced-vocab draft’s d2t index map) must keep their dtype — token ids are not representable in bf16. FP32 e_score_correction_bias tensors are inference-critical MoE router state and must also stay FP32.

nemo_rl.models.policy.workers.automodel_policy_worker._module_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.automodel_policy_worker.automodel_params_generator(

model: torch.nn.Module,

target_dtype: torch.dtype,

draft_model: Optional[torch.nn.Module] = None,

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

Policy weights, followed by draft.<name> weights when co-training.

The draft is a plain HF module (no LoRA, no state-dict adapter), so _module_params_generator’s merge/adapt steps are no-ops for it.

nemo_rl.models.policy.workers.automodel_policy_worker._assert_lora_modules_reachable(

module_map: dict[str, torch.nn.Module],

state_dict: dict[str, Any],

) → None

Fail loudly if a LinearLoRA has no matching state_dict weight.

A LoRA module the merge cannot find reaches the inference engine as its bare base weight; nothing else reports it, and the only symptom is a slowly growing train/generation logprob mismatch many steps later.

Only LinearLoRA is covered because it is the only adapter this file merges. When the grouped-expert LoRA merge lands (#3353), extend this check to those adapter modules too so that path keeps the same loud failure instead of silently refitting base expert weights.

nemo_rl.models.policy.workers.automodel_policy_worker._maybe_merge_lora_weight(

module_map: dict[str, torch.nn.Module],

fqn: str,

tensor: torch.Tensor,

) → torch.Tensor

nemo_rl.models.policy.workers.automodel_policy_worker._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.automodel_policy_worker.AutomodelPolicyWorkerImpl(

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.AutomodelCheckpointEngineSendMixin, 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 by weight-streaming refit.

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,

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. With policy.router_replay.enabled, the rollout routed_experts are replayed (R3) unless require_router_replay is False (e.g. for the reference model).

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

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

use_legacy_packing: bool = False,

) → 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,

use_legacy_packing: bool = False,

) → 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,

*,

is_final_checkpoint: bool,

) → 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,

) → None

Initialize the AutomodelCheckpointManager for this worker.

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

Parameters:

config_updates – Automodel checkpoint fields to set during initialization.

class nemo_rl.models.policy.workers.automodel_policy_worker.AutomodelPolicyWorker(

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.models.policy.workers.automodel_policy_worker.AutomodelPolicyWorkerImpl