nemo_rl.models.value.workers.megatron_value_worker#

Module Contents#

Classes#

MegatronValueWorkerImpl

Megatron-Core based value function worker for PPO.

MegatronValueWorker

Functions#

forward_step_value

Forward step for the value model.

_install_value_head_load_skip

Give the chunk a hide_loss_modules context manager that drops output_layer.*.

make_value_head_hook

Build the pre-wrap hook that installs the value head and lets it skip base loads.

_unpack_value_sequences

Unpack packed (and CP-sharded) per-token values back to [B, S].

_value_packed_loss_prepare_fn

Prepare one packed sequence’s value-head output for MseValueLossFn.

Data#

TokenizerType

API#

nemo_rl.models.value.workers.megatron_value_worker.TokenizerType#

‘TypeVar(…)’

nemo_rl.models.value.workers.megatron_value_worker.forward_step_value(
state,
global_valid_seqs,
global_valid_toks,
data_iterator,
model,
*,
loss_fn,
pack_sequences=False,
defer_fp32_logits=None,
cp_normalize=True,
num_microbatches=1,
policy_cfg=None,
)#

Forward step for the value model.

The LM head is replaced at build time by a LinearForLastLayer (hidden_size -> 1) value head (see create_value_head_hook), so model(...) returns per-token values directly — the head already gathers sequence-parallel shards.

nemo_rl.models.value.workers.megatron_value_worker._install_value_head_load_skip(
chunk: megatron.core.models.gpt.GPTModel,
) → None#

Give the chunk a hide_loss_modules context manager that drops output_layer.*.

The freshly-initialized value head is never in a base checkpoint, so Megatron-Bridge enters this context during finetune loads (and HF->Megatron conversion) to skip it. Resume loads run with finetune=False, so the trained head is restored normally.

nemo_rl.models.value.workers.megatron_value_worker.make_value_head_hook(hidden_size: int, sequence_parallel: bool)#

Build the pre-wrap hook that installs the value head and lets it skip base loads.

nemo_rl.models.value.workers.megatron_value_worker._unpack_value_sequences(
values: torch.Tensor,
cu_seqlens_padded: torch.Tensor,
unpacked_seqlen: int,
cp_group: Optional[torch.distributed.ProcessGroup] = None,
) → torch.Tensor#

Unpack packed (and CP-sharded) per-token values back to [B, S].

values is the value head output for a packed microbatch, shape [1, T // CP] (SP/TP already gathered inside the head). For each sequence delimited by cu_seqlens_padded (indices into the full, non-CP-adjusted packed layout) the local CP shard is all-gathered to the full sequence, then shifted right by one so values[t] = V(state before token t) — matching the unpacked path.

nemo_rl.models.value.workers.megatron_value_worker._value_packed_loss_prepare_fn(
logits: torch.Tensor,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
loss_fn: Optional[nemo_rl.algorithms.loss.interfaces.LossFunction] = None,
vocab_parallel_rank: Optional[int] = None,
vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
) → tuple[dict[str, torch.Tensor], nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]]#

Prepare one packed sequence’s value-head output for MseValueLossFn.

SequencePackingLossWrapper hands us a single sequence’s CP-sharded value slice [1, padded_len // CP]. All-gather across CP to [1, padded_len], shift right by one (values[t] = V(state before token t)), and truncate to the sequence’s unpadded length so it lines up with data["returns"].

class nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorkerImpl(
config: nemo_rl.models.value.config.ValueConfig,
tokenizer: nemo_rl.models.value.workers.megatron_value_worker.TokenizerType,
weights_path: Optional[str] = None,
optimizer_path: Optional[str] = None,
init_optimizer: bool = True,
*,
worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding,
**kwargs: Any,
)#

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

Megatron-Core based value function worker for PPO.

This worker wraps a Megatron-Core GPT model backbone with a value head (Linear: hidden_size -> 1) to predict per-token state values.

Supports:

  • Tensor parallelism (TP)

  • Pipeline parallelism (PP)

  • Context parallelism (CP)

  • Sequence packing

  • Activation checkpointing

  • FP8 quantization

  • MoE models

Initialization

Initialize the MegatronValueWorker.

Parameters:

  • config – Value model configuration.

  • tokenizer – HuggingFace tokenizer.

  • weights_path – Path to load finetuned weights from (optional).

  • optimizer_path – Path to load optimizer state from (optional).

  • init_optimizer – Whether to initialize the optimizer.

  • worker_sharding_annotations – Sharding topology for distributed training.

__repr__()#
static configure_worker(
num_gpus: int | float,
bundle_indices: Optional[tuple[int, list[int]]] = None,
) → tuple[dict[str, Any], dict[str, str], dict[str, Any], dict[str, Any]]#

Worker-controlled Ray actor configuration.

Mirrors MegatronPolicyWorker to ensure NVLS communication 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

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

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

Parameters:

  • data –

    BatchedDataDict containing training data with keys:

    • input_ids, input_lengths, token_mask, sample_mask, returns

  • loss_fn – Value loss function (e.g., MseValueLossFn).

  • eval_mode – If True, run forward only without parameter updates.

  • gbs – Global batch size override.

  • mbs – Micro batch size override.

Returns:

Dictionary with training metrics (global_loss, grad_norm, etc.)

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 per-token value predictions for a batch of data.

Parameters:

  • data – BatchedDataDict containing input_ids and input_lengths.

  • micro_batch_size – Override for inference micro batch size.

Returns:

BatchedDataDict with “values” key of shape [batch_size, seq_length].

prepare_for_training() → None#

Move model and optimizer to CUDA for training.

prepare_for_inference()#

Prepare model for value inference.

move_model(
model: torch.nn.Module,
device: str,
move_params: bool = True,
move_grads: bool = True,
) → torch.nn.Module#

Move model parameters and gradient buffers to the specified device.

move_optimizer(device: str)#

Move optimizer state to the specified device.

save_checkpoint(
weights_path: str,
optimizer_path: Optional[str] = None,
tokenizer_path: Optional[str] = None,
**kwargs,
)#

Save a checkpoint of the value model.

The value head is the model’s output_layer and is saved as part of the normal Megatron dist-checkpoint — no separate value-head sidecar.

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

Load a checkpoint for the value model.

finish_inference() → None#

Offload model params to CPU after inference.

finish_training() → None#

Offload model, gradients, and optimizer to CPU after training.

class nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker(
config: nemo_rl.models.value.config.ValueConfig,
tokenizer: nemo_rl.models.value.workers.megatron_value_worker.TokenizerType,
weights_path: Optional[str] = None,
optimizer_path: Optional[str] = None,
init_optimizer: bool = True,
*,
worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding,
**kwargs: Any,
)#

Bases: nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorkerImpl