nemo_rl.models.automodel.draft.hidden_capture#

Forward-hook capture of the policy’s per-layer hidden states for the draft.

Module Contents#

Classes#

DSparkHiddenCapture

Forward hooks capturing the policy’s per-layer hiddens for the draft.

Functions#

_resolve_layers

Locate the decoder backbone (embedding owner) and its layer container.

_layer_module

_hook_output_tensor

API#

nemo_rl.models.automodel.draft.hidden_capture._resolve_layers(
policy_model: torch.nn.Module,
) → tuple[torch.nn.Module, Any]#

Locate the decoder backbone (embedding owner) and its layer container.

The container is an nn.ModuleList for HF-style stacks or an nn.ModuleDict keyed by stringified layer index for Automodel custom backbones (e.g. Qwen3_5Moe); use _layer_module to index either uniformly.

nemo_rl.models.automodel.draft.hidden_capture._layer_module(layers: Any, layer_id: int) → torch.nn.Module#
nemo_rl.models.automodel.draft.hidden_capture._hook_output_tensor(output: Any) → torch.Tensor#
class nemo_rl.models.automodel.draft.hidden_capture.DSparkHiddenCapture(
policy_model: torch.nn.Module,
target_layer_ids: list[int],
)#

Forward hooks capturing the policy’s per-layer hiddens for the draft.

Captured tensors are detached: the DSpark loss never backprops into the policy trunk. The hooks are “armed” by a pre-forward hook on the policy root and disarmed by collect(): activation checkpointing replays layer forwards during backward (after the loss consumed the capture), and the disarmed hooks turn those replays into no-ops instead of re-running capture work (and, under TP, its collectives).

Initialization

property active: bool#
activate() → None#
deactivate() → None#
clear() → None#
collect() → torch.Tensor#

Concatenated target hidden states captured by the layer hooks.