From 439b585ff0b6db3de9448d1e9ddc9413aa3c7ea4 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 26 Jun 2026 15:45:45 -0700 Subject: [PATCH] fix(grpo_sync): skip refit for colocated MegatronGeneration In grpo_train_sync, NEED_REFIT was unconditionally True for any non-None policy_generation, so colocated MegatronGeneration ran refit_policy_generation which dispatched to the IPC ZMQ path. MegatronGeneration only overrides update_weights_from_collective, so the colocated IPC ZMQ branch raised NotImplementedError. Mirror the legacy grpo.py:1983-1986 guard so the two trainers behave identically for this case. Signed-off-by: Zhiyu Li --- nemo_rl/algorithms/grpo_sync.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 4ece3f81da..816e7e6d83 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -76,6 +76,7 @@ from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.experience.sync_rollout_actor import SyncRolloutActor from nemo_rl.models.generation.interfaces import GenerationInterface +from nemo_rl.models.generation.megatron import MegatronGeneration from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface from nemo_rl.utils.checkpoint import CheckpointManager from nemo_rl.utils.logger import Logger, print_message_log_samples @@ -404,7 +405,10 @@ def grpo_train_sync( kv_scales_cache = None # Cache reused for computed kv scales - NEED_REFIT = True + NEED_REFIT = not ( + isinstance(policy_generation, MegatronGeneration) + and master_config.policy["generation"]["colocated"]["enabled"] + ) # If policy_generation is None, use the policy as the generation interface (megatron framework backend) if policy_generation is None: policy_generation = policy # type: ignore