feat(policy): execute bounded exact-call and shared-prefix training layouts

nemo_rl/algorithms/loss/__init__.py

@@ -34,11 +34,13 @@
 from nemo_rl.algorithms.loss.utils import (
     prepare_loss_input,
     prepare_packed_loss_input,
+    prepare_tree_loss_input,
 )
 from nemo_rl.algorithms.loss.wrapper import (
     DraftLossWrapper,
     SequencePackingFusionLossWrapper,
     SequencePackingLossWrapper,
+    TreePackingLossWrapper,
     wrap_loss_fn_with_input_preparation,
 )
 
@@ -61,8 +63,10 @@
     "PreferenceLossFn",
     "prepare_loss_input",
     "prepare_packed_loss_input",
+    "prepare_tree_loss_input",
     "SequencePackingFusionLossWrapper",
     "SequencePackingLossWrapper",
+    "TreePackingLossWrapper",
     "DraftLossWrapper",
     "wrap_loss_fn_with_input_preparation",
 ]

nemo_rl/algorithms/loss/utils.py

@@ -25,6 +25,7 @@
 from nemo_rl.algorithms.x_token.loss_utils import (
     prepare_xtoken_cross_tokenizer_loss_input,
 )
+from nemo_rl.data.packed_rollouts import TREE_ATTENTION_EDGE_TARGET_IDS
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.model_utils import (
     ChunkedDistributedCrossEntropyToFixedLogits,
@@ -33,6 +34,7 @@
     _get_tokens_on_this_cp_rank,
     allgather_cp_sharded_tensor,
     from_parallel_logits_to_logprobs_packed_sequences,
+    from_parallel_tree_logits_to_logprobs,
     get_cp_sharded_next_token_logprobs,
     get_distillation_topk_logprobs_from_logits,
     get_next_token_logprobs_from_logits,
@@ -370,6 +372,51 @@ def prepare_opd_full_loss_input(
     return loss_input
 
 
+def prepare_tree_loss_input(
+    logits: torch.Tensor,
+    data: BatchedDataDict[Any],
+    loss_fn: LossFunction,
+    packed_seq_params: Any,
+    vocab_parallel_rank: int,
+    vocab_parallel_group: torch.distributed.ProcessGroup,
+    context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+    sampling_params: Optional[TrainingSamplingParams] = None,
+    chunk_size: Optional[int] = None,
+) -> tuple[dict[str, Any], BatchedDataDict[Any]]:
+    """Prepare policy loss from logits projected only at sampled tree edges."""
+    if loss_fn.input_type != LossInputType.LOGPROB:
+        raise ValueError("tree attention currently supports logprob losses only")
+    logprobs = from_parallel_tree_logits_to_logprobs(
+        logits,
+        data[TREE_ATTENTION_EDGE_TARGET_IDS],
+        packed_seq_params,
+        vocab_start_index=vocab_parallel_rank * logits.shape[-1],
+        vocab_end_index=(vocab_parallel_rank + 1) * logits.shape[-1],
+        tp_group=vocab_parallel_group,
+        cp_group=context_parallel_group,
+        inference_only=False,
+        chunk_size=chunk_size,
+        sampling_params=sampling_params,
+    )
+    if need_top_k_or_top_p_filtering(sampling_params):
+        mask = data["token_mask"][:, 1:] * data["sample_mask"].unsqueeze(-1)
+        logprobs = mask_out_neg_inf_logprobs(logprobs, mask, "curr_logprobs")
+        if getattr(loss_fn, "reference_policy_kl_penalty", 0) != 0:
+            data["curr_logprobs_unfiltered"] = from_parallel_tree_logits_to_logprobs(
+                logits,
+                data[TREE_ATTENTION_EDGE_TARGET_IDS],
+                packed_seq_params,
+                vocab_start_index=vocab_parallel_rank * logits.shape[-1],
+                vocab_end_index=(vocab_parallel_rank + 1) * logits.shape[-1],
+                tp_group=vocab_parallel_group,
+                cp_group=context_parallel_group,
+                inference_only=False,
+                chunk_size=chunk_size,
+                sampling_params=None,
+            )
+    return {"next_token_logprobs": logprobs}, data
+
+
 def prepare_loss_input(
     logits: torch.Tensor,
     data: BatchedDataDict[Any],

nemo_rl/algorithms/loss/wrapper.py

@@ -65,7 +65,6 @@ def __init__(
             vocab_parallel_rank: Vocab parallel rank.
             vocab_parallel_group: Vocab parallel group.
             context_parallel_group: Context parallel group.
-
             vocab_parallel_rank, vocab_parallel_group, context_parallel_group are only used for megatron policy worker.
 
         Returns:
@@ -185,6 +184,49 @@ def __call__(
         return loss_accum, metrics_accum
 
 
+class TreePackingLossWrapper:
+    """Prepare and apply one batched loss over sampled edges of a packed tree."""
+
+    def __init__(
+        self,
+        loss_fn: LossFunction,
+        prepare_fn: Callable[..., Any],
+        packed_seq_params: Any,
+        vocab_parallel_rank: int,
+        vocab_parallel_group: torch.distributed.ProcessGroup,
+        context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+    ):
+        self.loss_fn = loss_fn
+        self.prepare_fn = prepare_fn
+        self.packed_seq_params = packed_seq_params
+        self.vocab_parallel_rank = vocab_parallel_rank
+        self.vocab_parallel_group = vocab_parallel_group
+        self.context_parallel_group = context_parallel_group
+
+    def __call__(
+        self,
+        next_token_logits: Tensor,
+        data: BatchedDataDict[Any],
+        global_valid_seqs: Tensor | None,
+        global_valid_toks: Tensor | None,
+    ) -> tuple[Tensor, dict[str, Any]]:
+        loss_input, prepared_data = self.prepare_fn(
+            logits=next_token_logits,
+            data=data,
+            loss_fn=self.loss_fn,
+            packed_seq_params=self.packed_seq_params,
+            vocab_parallel_rank=self.vocab_parallel_rank,
+            vocab_parallel_group=self.vocab_parallel_group,
+            context_parallel_group=self.context_parallel_group,
+        )
+        return self.loss_fn(
+            data=prepared_data,
+            global_valid_seqs=global_valid_seqs,
+            global_valid_toks=global_valid_toks,
+            **loss_input,
+        )
+
+
 class SequencePackingFusionLossWrapper:
     """Fused sequence packing loss wrapper that processes all sequences in one forward pass.
 

nemo_rl/algorithms/single_controller.py

@@ -4627,7 +4627,13 @@ async def _sync_weights(
             # the loop this is a blocking Ray call, and a wedged generation worker would
             # freeze the event loop itself -- taking the watchdog, which is an asyncio
             # task on that same loop, down with it.
-            await asyncio.to_thread(self._gen.invalidate_kv_cache)
+            if not await asyncio.to_thread(self._gen.invalidate_kv_cache):
+                raise RuntimeError(
+                    "Generation KV-cache invalidation failed after a weight update"
+                )
+        await asyncio.to_thread(
+            self._gen.set_generation_weight_version, self._trainer_version
+        )
         elapsed = time.monotonic() - t0
 
         print(f"  _sync_weights: sync done in {elapsed:.3f}s", flush=True)