nemo_rl.models.megatron.opd_full_capture#

Capture the pre-LM-head hidden states of a frozen MOPD teacher.

Full-vocabulary MOPD ships these to the student, which projects them with a teacher LM-head shard. Moving hidden_size floats per token instead of vocab_size keeps teacher and student parallelism decoupled.

The capture is a forward pre-hook on output_layer: GPTModel.forward ends in logits, _ = self.output_layer(hidden_states, ...), so the hook’s first argument is the post-final_layernorm hidden state. Hooking decoder instead would miss final_layernorm, and decoder.final_layernorm is absent under some transformer specs.

Module Contents#

Classes#

OutputLayerInputCapture

Capture the tensor fed into a Megatron model’s output_layer.

Functions#

get_opd_full_capture_context

Build the capture context for one teacher forward pass.

API#

class nemo_rl.models.megatron.opd_full_capture.OutputLayerInputCapture(model: torch.nn.Module)#

Capture the tensor fed into a Megatron model’s output_layer.

The tensor is [S_local, B, H] (Megatron sequence-first). Sequence parallelism shards S_local across the TP group, so

Meth:

get_hidden_states gathers it back.

Initialization

property is_available: bool#

Whether this rank owns an output_layer to hook (last PP stage).

_pre_hook(
module: torch.nn.Module,
args: tuple[Any, ...],
) → None#
capture_context() → Iterator[None]#

Register the pre-hook for the duration of one forward pass.

get_hidden_states() → Optional[torch.Tensor]#

Return the captured hidden states in full [S, B, H] layout.

Returns:

The detached hidden states with sequence parallelism undone, or None when this rank ran no output layer (non-final pipeline stage) or the forward never reached it.

Raises:

ValueError – If the captured tensor is not rank 3.

nemo_rl.models.megatron.opd_full_capture.get_opd_full_capture_context(
model: torch.nn.Module,
enabled: bool,
) → tuple[ContextManager[Any], Optional[nemo_rl.models.megatron.opd_full_capture.OutputLayerInputCapture]]#

Build the capture context for one teacher forward pass.

Parameters:

  • model – The (possibly wrapped) Megatron model about to run forward.

  • enabled – Whether full-vocabulary hidden-state capture is requested.

Returns:

Tuple of the context manager to wrap the forward in, and the capture object to read afterwards (None when disabled or unavailable).