diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index bd332e54195..1640fe88e9c 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -65,6 +65,8 @@ def _aggregate_train_results(results: list[dict[str, Any]]) -> dict[str, Any]: } if "moe_metrics" in results[0]: out["moe_metrics"] = results[0]["moe_metrics"] + if "mtp_metrics" in results[0]: + out["mtp_metrics"] = results[0]["mtp_metrics"] all_mb_metrics: dict[str, list[Any]] = defaultdict(list) for r in results: for k, v in r["all_mb_metrics"].items(): diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 3135f88721c..79888f9be75 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -1241,10 +1241,40 @@ def _split_step_state_init( if not isinstance(metric_normalizations, dict): metric_normalizations = {} + model_config = self._get_model_config() + mtp_num_layers = getattr(model_config, "mtp_num_layers", None) + mtp_enabled = mtp_num_layers is not None and mtp_num_layers > 0 + mtp_detach_heads = bool(getattr(model_config, "mtp_detach_heads", False)) + mtp_loss_scaling_factor = getattr(model_config, "mtp_loss_scaling_factor", 0.1) + loss_type = getattr(loss_fn, "loss_type", LossType.TOKEN_LEVEL) + + # MTP is a token-summed auxiliary loss and therefore always needs the + # valid-token denominator. Under a sequence-level main loss, the split + # path can apply that distinct denominator only when mcore has isolated + # and tagged the detached MTP parameters. Attached-head gradients are + # already mixed into the backbone and cannot be corrected after the + # global counts become available at finish. + if ( + mtp_enabled + and mtp_loss_scaling_factor != 0 + and loss_type != LossType.TOKEN_LEVEL + and not mtp_detach_heads + ): + raise ValueError( + "MTP with a nonzero loss weight and sequence-level loss requires " + "policy.megatron_cfg.mtp_detach_heads=True on the SingleController " + "split training path because the MTP auxiliary gradient must be " + "normalized by valid tokens independently of the main loss. " + f"Got loss_type={loss_type}, mtp_num_layers={mtp_num_layers}, " + f"mtp_loss_scaling_factor={mtp_loss_scaling_factor}." + ) + return { "loss_fn": loss_fn, - "loss_type": getattr(loss_fn, "loss_type", LossType.TOKEN_LEVEL), + "loss_type": loss_type, "metric_normalizations": metric_normalizations, + "mtp_enabled": mtp_enabled, + "mtp_detach_heads": mtp_detach_heads, "gbs": gbs or self.cfg["train_global_batch_size"], "mbs": mbs or self.cfg["train_micro_batch_size"], "local_valid_seqs": torch.zeros((), dtype=torch.float64, device="cuda"), @@ -1351,6 +1381,12 @@ def begin_train_step( state = self._split_step_state_init(loss_fn=loss_fn, gbs=gbs, mbs=mbs) + # Leave this unset so mcore falls back to config.grad_scale_func and + # inherits the optimizer's dynamic loss scale (especially for fp16). + # Also clear any transient callable left by an interrupted older step. + if state["mtp_enabled"]: + self._set_mtp_grad_scale_func(None) + # Null the three mcore hooks that would fire a mid-step DP reduce: # grad_sync_func — PP scheduler's direct call on last-MB boundaries # (PP>1 path). @@ -1422,6 +1458,12 @@ def train_microbatch( # gradient finalization at all. Restore here; the caller is still # expected to invoke abort_train_step (idempotent on the saved # values) to drop ``_train_step_state``. + try: + self._set_mtp_grad_scale_func(None) + except Exception: + log.exception( + "failed to clear MTP gradient scaling after train_microbatch error" + ) try: self._restore_saved_mcore_hooks(state) except Exception: @@ -1457,6 +1499,15 @@ def _train_microbatch_body( state["local_valid_seqs"] = state["local_valid_seqs"] + call_local_seqs state["local_valid_toks"] = state["local_valid_toks"] + call_local_toks + # Match the synchronous Megatron path: derive the mask on the worker + # immediately before microbatch processing so sequence packing applies + # the same layout transformation to tokens and the MTP loss mask. The + # mask is worker-local derived data; it does not need a TQ schema field. + if state["mtp_enabled"] and "token_mask" in data and "sample_mask" in data: + data["mtp_loss_mask"] = data["token_mask"] * data["sample_mask"].unsqueeze( + -1 + ) + # The number of chunks per optimizer step is a first-class property of # this path — it decides how many times gradients are accumulated before # a single reduce — but it was previously only recoverable by calibrating @@ -1492,6 +1543,9 @@ def _train_microbatch_body( self.cfg, state["mbs"], straggler_timer=self.mcore_state.straggler_timer, + delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, ) state["total_num_microbatches"] += int(num_microbatches) @@ -1583,6 +1637,12 @@ def finish_train_step(self) -> dict[str, Any]: # the body got). Restore unconditionally so future steps run with # the right config. Leave ``_train_step_state`` for the caller's # abort_train_step to clear. + try: + self._set_mtp_grad_scale_func(None) + except Exception: + log.exception( + "failed to clear MTP gradient scaling after finish_train_step error" + ) try: self._restore_saved_mcore_hooks(state) except Exception: @@ -1616,6 +1676,17 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]: # global mean grad; for reduce_scatter (dist-opt) it's the shard. # Either way, opt.step sees the right-normalized gradient. self.model.scale_gradients(inv_n) + # The uniform rescale gives MTP the main loss's denominator. Correct + # detached, MTP-tagged parameters back to the valid-token denominator + # used by the synchronous path. For token-level loss the factor is 1. + # This runs before gradient reduction; scaling and reduction are linear. + if state["mtp_enabled"] and state["mtp_detach_heads"]: + self._scale_mtp_param_grads( + float((n_safe / global_valid_toks.clamp(min=1)).item()) + ) + # No more forward/backward calls remain in this step. Clear the + # callable before optimizer/scheduler/checkpoint state can serialize it. + self._set_mtp_grad_scale_func(None) # End-of-step gradient finalization, exactly once per optimizer step. # ``begin_train_step`` nulled ``finalize_model_grads_func`` so mcore's @@ -1670,6 +1741,11 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]: # opt.step clips internally (clip_grad config); operates on the # already-rescaled grad. Returns (success, grad_norm, num_zeros). update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step() + mtp_grad_norm = ( + self.optimizer.grad_norms_by_group.get("mtp") + if state["mtp_enabled"] + else None + ) pg_collection = get_pg_collection(self.model) update_successful = logical_and_across_model_parallel_group( @@ -1681,6 +1757,13 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]: num_zeros_in_grad = reduce_max_stat_across_model_parallel_group( num_zeros_in_grad, mp_group=pg_collection.mp ) + if state["mtp_enabled"]: + # MTP parameters live on the last PP stage. Make their independently + # clipped grad norm visible to every model-parallel rank before the + # driver selects a replica leader's result. + mtp_grad_norm = reduce_max_stat_across_model_parallel_group( + mtp_grad_norm, mp_group=pg_collection.mp + ) # Mirrors train(): without re-enabling the pre-hook __init__ removed, the # param all-gather never runs and each forward sees only its own shard. @@ -1833,6 +1916,12 @@ def _scale_metric(name: str, value: Any) -> Any: if moe_metrics: metrics["moe_metrics"] = moe_metrics + self._collect_mtp_metrics( + metrics, + state["total_num_microbatches"], + mtp_grad_norm, + ) + self._train_step_state = None return metrics @@ -1841,12 +1930,15 @@ def abort_train_step(self) -> None: state = getattr(self, "_train_step_state", None) if state is None: return - # Restore the mcore hooks first so the model is back to a normal - # state before zero_grad_buffer touches anything. - self._restore_saved_mcore_hooks(state) - self.model.zero_grad_buffer() - self.optimizer.zero_grad() - self._train_step_state = None + # Drop the step-local MTP scaler and restore the mcore hooks before + # zero_grad_buffer touches anything. + try: + self._set_mtp_grad_scale_func(None) + finally: + self._restore_saved_mcore_hooks(state) + self.model.zero_grad_buffer() + self.optimizer.zero_grad() + self._train_step_state = None @wrap_with_nvtx_name("megatron_policy_worker/get_logprobs") def get_logprobs( @@ -2228,6 +2320,26 @@ def _set_mtp_grad_scale_func(self, func): if config is not None: config.mtp_grad_scale_func = func + def _scale_mtp_param_grads(self, factor: float) -> None: + """Scale detached MTP parameters' gradients by ``factor``. + + MCore tags every MTP parameter ``grad_norm_group='mtp'`` when + ``mtp_detach_heads`` is enabled. In that configuration the auxiliary + loss reaches no shared parameters, so its denominator can be corrected + independently. ``main_grad`` is a view into the DDP gradient buffer. + + Args: + factor: Multiplier for MTP gradients. ``1.0`` is a no-op. + """ + if factor == 1.0: + return + for param in self.model.parameters(): + if getattr(param, "grad_norm_group", None) != "mtp": + continue + main_grad = getattr(param, "main_grad", None) + if main_grad is not None: + main_grad.mul_(factor) + def _get_model_config(self): """Get the underlying model config (handle Float16Module wrapper).""" model = self.model diff --git a/tests/unit/models/policy/test_megatron_split_state.py b/tests/unit/models/policy/test_megatron_split_state.py index 61eaf84a855..13513fc5ea8 100644 --- a/tests/unit/models/policy/test_megatron_split_state.py +++ b/tests/unit/models/policy/test_megatron_split_state.py @@ -41,11 +41,14 @@ - ``prepare_for_lp_inference`` offloading grad buffers mid-step, which frees (not copies) every earlier chunk's gradients while the 1/N normalizer still counts them. + - MTP masks, auxiliary-gradient scaling, grad norms, and tracker metrics being + present only in the synchronous Megatron path. """ from __future__ import annotations import logging +from typing import Any from unittest.mock import MagicMock, patch import pytest @@ -85,6 +88,15 @@ def _make_mock_model(): # hook need to be able to clear it to a real None. model.config.finalize_model_grads_func = MagicMock(name="ORIGINAL_FINALIZE") model.config.num_moe_experts = None # disable MoE branch + model.config.mtp_num_layers = None # disable MTP unless a test opts in + model.config.mtp_grad_scale_func = None + # Explicit values prevent MagicMock's auto-created attributes from making + # every test look like it uses detached heads with an arbitrary loss weight. + model.config.mtp_detach_heads = False + model.config.mtp_loss_scaling_factor = 0.1 + # Prevent MagicMock's dynamic attributes from making _get_model_config + # mistake this bare model for a Float16Module wrapper. + model.module = None # no_sync() is a context manager — return a MagicMock that supports # __enter__/__exit__ so the `with self.model.no_sync():` block works. model.no_sync = MagicMock( @@ -115,6 +127,7 @@ def _make_worker(loss_type): w.optimizer = MagicMock() # MegatronOptimizer.step returns (success, grad_norm, num_zeros) w.optimizer.step.return_value = (True, 0.5, 0) + w.optimizer.grad_norms_by_group = {} w.optimizer.param_groups = [{"lr": 1e-4, "weight_decay": 0.01}] w.scheduler = MagicMock() w.scheduler.get_lr.return_value = 1e-4 @@ -144,6 +157,9 @@ def _make_worker(loss_type): w.dtype = torch.float32 w._is_reward_model = False w._router_replay_enabled = False + w.delegate_pack_to_model = False + w.delegate_mtp_loss_mask_to_model = False + w.model_slices_context_parallel_inputs = False # Normally set from get_rank_safe() in __init__, which object.__new__ skips. # The step summary in finish_train_step reads it eagerly to decide whether # this rank prints. @@ -315,6 +331,51 @@ def test_uses_cfg_defaults_when_gbs_mbs_omitted(self, mock_module_symbols): assert w._train_step_state["gbs"] == w.cfg["train_global_batch_size"] assert w._train_step_state["mbs"] == w.cfg["train_micro_batch_size"] + def test_records_mtp_enabled_from_model_config(self, mock_module_symbols): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.model.config.mtp_num_layers = 2 + w.begin_train_step(loss_fn=w._test_loss_fn) + assert w._train_step_state["mtp_enabled"] is True + + def test_rejects_sequence_level_mtp_with_attached_heads(self, mock_module_symbols): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.SEQUENCE_LEVEL) + w.model.config.mtp_num_layers = 2 + with pytest.raises(ValueError, match="mtp_detach_heads"): + w.begin_train_step(loss_fn=w._test_loss_fn) + assert getattr(w, "_train_step_state", None) is None + + def test_allows_sequence_level_mtp_with_detached_heads(self, mock_module_symbols): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.SEQUENCE_LEVEL) + w.model.config.mtp_num_layers = 2 + w.model.config.mtp_detach_heads = True + w.begin_train_step(loss_fn=w._test_loss_fn) + assert w._train_step_state["mtp_detach_heads"] is True + + def test_allows_token_level_mtp_with_attached_heads(self, mock_module_symbols): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.model.config.mtp_num_layers = 2 + w.begin_train_step(loss_fn=w._test_loss_fn) + assert w._train_step_state["mtp_detach_heads"] is False + + def test_allows_zero_weight_sequence_level_mtp_with_attached_heads( + self, mock_module_symbols + ): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.SEQUENCE_LEVEL) + w.model.config.mtp_num_layers = 2 + w.model.config.mtp_loss_scaling_factor = 0.0 + w.begin_train_step(loss_fn=w._test_loss_fn) + assert w._train_step_state["mtp_detach_heads"] is False + # ── _assert_step_open ──────────────────────────────────────────────────── @@ -369,6 +430,45 @@ def test_invokes_megatron_forward_backward_once(self, mock_module_symbols): w.train_microbatch(_fake_batch()) assert mock_module_symbols["mfb"].call_count == 1 + @pytest.mark.parametrize( + ( + "delegate_pack_to_model", + "delegate_mtp_loss_mask_to_model", + "model_slices_context_parallel_inputs", + ), + [ + pytest.param(True, True, False, id="model-owned-packing"), + pytest.param(False, False, True, id="model-owned-cp-slicing"), + ], + ) + def test_forwards_model_owned_packing_flags( + self, + mock_module_symbols: dict[str, MagicMock], + delegate_pack_to_model: bool, + delegate_mtp_loss_mask_to_model: bool, + model_slices_context_parallel_inputs: bool, + ) -> None: + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.model.config.mtp_num_layers = 1 + w.delegate_pack_to_model = delegate_pack_to_model + w.delegate_mtp_loss_mask_to_model = delegate_mtp_loss_mask_to_model + w.model_slices_context_parallel_inputs = model_slices_context_parallel_inputs + + w.begin_train_step(loss_fn=w._test_loss_fn) + w.train_microbatch(_fake_batch()) + + kwargs = mock_module_symbols["gmi"].call_args.kwargs + assert kwargs["delegate_pack_to_model"] is delegate_pack_to_model + assert ( + kwargs["delegate_mtp_loss_mask_to_model"] is delegate_mtp_loss_mask_to_model + ) + assert ( + kwargs["model_slices_context_parallel_inputs"] + is model_slices_context_parallel_inputs + ) + def test_passes_placeholder_n_one_to_loss(self, mock_module_symbols): """The N=1 trick: loss must be called with global_valid_*=1 so it returns un-normalized sums; finish does the 1/N rescale.""" @@ -427,6 +527,37 @@ def test_does_not_call_optimizer_step(self, mock_module_symbols): w.train_microbatch(_fake_batch()) w.optimizer.step.assert_not_called() + def test_builds_mtp_mask_and_uses_main_loss_scale_fallback( + self, mock_module_symbols + ): + """The mask combines both inputs and MTP inherits the main loss scale.""" + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.model.config.mtp_num_layers = 2 + batch = _fake_batch() + batch["token_mask"][0, 5] = 0 + batch["sample_mask"][3] = 0 + w.begin_train_step(loss_fn=w._test_loss_fn) + w.train_microbatch(batch) + + mask = batch["mtp_loss_mask"] + assert mask.shape == (8, 257) + assert mask[3].sum().item() == pytest.approx(0.0) + assert mask[0].sum().item() == pytest.approx(256.0) + assert mask[1].sum().item() == pytest.approx(257.0) + assert w.model.config.mtp_grad_scale_func is None + + def test_skips_mtp_mask_and_scale_when_disabled(self, mock_module_symbols): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + batch = _fake_batch() + w.begin_train_step(loss_fn=w._test_loss_fn) + w.train_microbatch(batch) + assert "mtp_loss_mask" not in batch + assert w.model.config.mtp_grad_scale_func is None + # ── finish_train_step ──────────────────────────────────────────────────── @@ -532,6 +663,60 @@ def test_picks_global_valid_seqs_for_sequence_level_loss(self, mock_module_symbo arg = w.model.scale_gradients.call_args.args[0] assert arg == pytest.approx(1.0 / 8.0, rel=1e-4) + @staticmethod + def _mtp_params() -> tuple[MagicMock, MagicMock]: + """Create one MTP-tagged parameter and one untagged parameter.""" + mtp_param = MagicMock() + mtp_param.grad_norm_group = "mtp" + mtp_param.main_grad = torch.ones(4) + other_param = MagicMock() + other_param.grad_norm_group = None + other_param.main_grad = torch.ones(4) + return mtp_param, other_param + + def test_mtp_grads_use_token_denominator_under_sequence_level_loss( + self, mock_module_symbols + ): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.SEQUENCE_LEVEL) + w.model.config.mtp_num_layers = 2 + w.model.config.mtp_detach_heads = True + mtp_param, other_param = self._mtp_params() + w.model.parameters = MagicMock(return_value=[mtp_param, other_param]) + w.begin_train_step(loss_fn=w._test_loss_fn) + w.train_microbatch(_fake_batch()) + grad_seen_by_finalize: list[torch.Tensor] = [] + w._train_step_state["saved_finalize_model_grads_func"] = ( + lambda models, num_tokens: grad_seen_by_finalize.append( + mtp_param.main_grad.clone() + ) + ) + w.finish_train_step() + + # The main loss gets 1/8; MTP additionally gets 8/2048, for net 1/2048. + assert w.model.scale_gradients.call_args.args[0] == pytest.approx(1.0 / 8.0) + expected_mtp_grad = torch.full((4,), 8.0 / 2048.0) + torch.testing.assert_close(mtp_param.main_grad, expected_mtp_grad) + torch.testing.assert_close(grad_seen_by_finalize[0], expected_mtp_grad) + torch.testing.assert_close(other_param.main_grad, torch.ones(4)) + + def test_mtp_grads_need_no_correction_under_token_level_loss( + self, mock_module_symbols + ): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.model.config.mtp_num_layers = 2 + w.model.config.mtp_detach_heads = True + mtp_param, _ = self._mtp_params() + w.model.parameters = MagicMock(return_value=[mtp_param]) + w.begin_train_step(loss_fn=w._test_loss_fn) + w.train_microbatch(_fake_batch()) + w.finish_train_step() + + torch.testing.assert_close(mtp_param.main_grad, torch.ones(4)) + def test_restores_grad_sync_func(self, mock_module_symbols): from nemo_rl.algorithms.loss.interfaces import LossType @@ -624,6 +809,37 @@ def test_moe_branch_uses_total_num_microbatches_for_scale( kwargs = mock_module_symbols["moe"].call_args.kwargs assert kwargs["loss_scale"] == pytest.approx(1.0 / 6.0, rel=1e-6) + def test_collects_mtp_metrics_and_reduced_grad_norm(self, mock_module_symbols): + """Finish surfaces the same post-#3194 MTP payload as sync train().""" + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.model.config.mtp_num_layers = 2 + w.optimizer.grad_norms_by_group = {"mtp": 1.25} + + def _collect( + metrics: dict[str, Any], + total_num_microbatches: int, + mtp_grad_norm: float | None, + ) -> None: + assert total_num_microbatches == 2 + metrics["mtp_metrics"] = { + "mtp_1_loss": 0.5, + "grad_norm": mtp_grad_norm, + } + + w._collect_mtp_metrics = MagicMock(side_effect=_collect) + w.begin_train_step(loss_fn=w._test_loss_fn) + w.train_microbatch(_fake_batch()) + metrics = w.finish_train_step() + + args = w._collect_mtp_metrics.call_args.args + assert args[1] == 2 # fixture: two pipeline microbatches per chunk + assert args[2] == pytest.approx(1.25) + assert metrics["mtp_metrics"]["mtp_1_loss"] == pytest.approx(0.5) + assert metrics["mtp_metrics"]["grad_norm"] == pytest.approx(1.25) + assert w.model.config.mtp_grad_scale_func is None + def test_loss_advertised_normalizers_applied(self, mock_module_symbols): """finish scales each metric by the denominator the loss advertised: TOKENS → 1/global_valid_toks, SEQUENCES → 1/global_valid_seqs, @@ -775,6 +991,21 @@ def test_can_begin_new_step_after_abort(self, mock_module_symbols): assert w._train_step_state is not None assert float(w._train_step_state["local_valid_seqs"].item()) == 0.0 + def test_clears_stale_mtp_gradient_scale_at_step_boundaries( + self, mock_module_symbols + ): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.model.config.mtp_num_layers = 1 + w.model.config.mtp_grad_scale_func = lambda: torch.tensor(7.0) + w.begin_train_step(loss_fn=w._test_loss_fn) + assert w.model.config.mtp_grad_scale_func is None + w.train_microbatch(_fake_batch()) + assert w.model.config.mtp_grad_scale_func is None + w.abort_train_step() + assert w.model.config.mtp_grad_scale_func is None + # ── grad_sync_func full lifecycle (integration of begin → finish/abort) ─ diff --git a/tests/unit/models/policy/test_split_api_wrappers.py b/tests/unit/models/policy/test_split_api_wrappers.py index b5db49e5601..79d1bef7a6c 100644 --- a/tests/unit/models/policy/test_split_api_wrappers.py +++ b/tests/unit/models/policy/test_split_api_wrappers.py @@ -228,6 +228,31 @@ def _result(leader: bool) -> dict: # _aggregate_train_results surfaces global_loss under "loss" assert out["loss"] == 1.0 + def test_finish_propagates_mtp_metrics(self): + """Worker-reduced MTP metrics survive the TQPolicy aggregation layer.""" + p, _ = _make_tq_policy() + with patch("nemo_rl.models.policy.tq_policy.ray") as mock_ray: + mock_ray.get.return_value = [ + { + "global_loss": 1.0, + "grad_norm": 0.5, + "all_mb_metrics": {"loss": [0.1]}, + "mtp_metrics": { + "mtp_1_loss": 0.25, + "mtp_1_acceptance_rate": 75.0, + "grad_norm": 1.25, + }, + "is_replica_leader": True, + } + ] + out = p.finish_train_step() + + assert out["mtp_metrics"] == { + "mtp_1_loss": 0.25, + "mtp_1_acceptance_rate": 75.0, + "grad_norm": 1.25, + } + def test_abort_consumes_single_data_futures_with_ray_get(self): p, wg = _make_tq_policy() with patch("nemo_rl.models.policy.tq_policy.ray") as mock_ray: