perf: reuse rollout behavior logprobs

Feed sampler-aligned behavior log-probabilities directly into online GRPO instead of allocating, synchronizing, and forwarding a duplicate old-policy model. Keep the old-model path as an offline compatibility fallback and validate supplied rollout tensors before loss computation.
2026-09-02 19:29:53 +08:00
parent e58a728b80
commit 4019ddac31
7 changed files with 115 additions and 44 deletions
@@ -287,13 +287,15 @@ class TrainContextBuilder:
                model=context.model,
                device=get_current_device(),
            )
-        if cfg.strategy in ("grpo", "online_grpo"):
+        if cfg.strategy == "grpo":
            kwargs["old_model"] = create_ref_model(
                cfg.model_fn,
                executor=executor,
                model=context.model,
                device=get_current_device(),
            )
+        elif cfg.strategy == "online_grpo":
+            kwargs["old_model"] = None
        context.strategy = StrategyFactory.create(
            cfg.strategy,
            model=context.model,