nemo_rl.models.value.workers.dtensor_value_worker_v2#

Module Contents#

Classes#

RightShiftLossWrapper

Wrap a LossFunction so value logits are right-shifted before loss.

DTensorValueWorkerV2Impl

DTensorValueWorkerV2

Functions#

right_shift_values

Shift values right by 1 along the sequence dim (V(s_{t+1}) -> V(s_t)).

get_train_context

Create combined context manager for training with context parallel and autocast.

API#

nemo_rl.models.value.workers.dtensor_value_worker_v2.right_shift_values(values: torch.Tensor) → torch.Tensor#

Shift values right by 1 along the sequence dim (V(s_{t+1}) -> V(s_t)).

Aligns value predictions with the Megatron value worker convention so GAE (rewards, returns), value targets, and value clipping all see the same V(s_t) semantics across backends. Preserves the input tensor shape: the first column becomes zeros and column t (t>=1) takes the value from column t-1.

class nemo_rl.models.value.workers.dtensor_value_worker_v2.RightShiftLossWrapper(
inner: nemo_rl.algorithms.loss.interfaces.LossFunction,
)#

Wrap a LossFunction so value logits are right-shifted before loss.

Initialization

__call__(*args, logits=None, **kwargs)#
__getattr__(name)#
nemo_rl.models.value.workers.dtensor_value_worker_v2.get_train_context(
cp_size: int,
cp_mesh: Any,
cp_buffers: list,
sequence_dim: int,
dtype: torch.dtype,
autocast_enabled: bool = True,
) → Generator[None, None, None]#

Create combined context manager for training with context parallel and autocast.

class nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2Impl(
config: nemo_rl.models.value.config.ValueConfig,
tokenizer: transformers.AutoTokenizer,
weights_path: Optional[str] = None,
optimizer_path: Optional[str] = None,
init_optimizer: bool = True,
**kwargs: Any,
)#

Bases: nemo_rl.models.policy.workers.base_policy_worker.AbstractPolicyWorker

__repr__() → str#

Customizes the actor’s prefix in the Ray logs.

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,
) → dict[str, Any]#

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

get_values(
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.value.interfaces.ValueOutputSpec]#

Get value predictions for a batch of data.

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

Prepare for training by loading model and optimizer to GPU.

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

Prepare for inference by loading model to GPU and setting eval mode.

finish_training() → None#

Offload value model and optimizer state after training.

finish_inference() → None#

Offload value model parameters after inference.

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

Move optimizer state to specified device.

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

Move model to specified device.

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

Move model buffers to specified device.

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

Move model to CUDA.

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

Move model to CPU.

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 value model.

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

Load a checkpoint into the value model.

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

Initialize the AutomodelCheckpointManager for this worker.

class nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2(
config: nemo_rl.models.value.config.ValueConfig,
tokenizer: transformers.AutoTokenizer,
weights_path: Optional[str] = None,
optimizer_path: Optional[str] = None,
init_optimizer: bool = True,
**kwargs: Any,
)#

Bases: nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2Impl