[feat] verl_v080 fixed rollout monkey-patch

prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/__init__.py

@@ -6,6 +6,7 @@
 3. vocab_parallel_log_probs_from_logits  → 自动 logprob restore
 4. no_padding_2_padding                  → PS 物理裁剪后修正序列长度
    (module-level + 所有 from...import 引用)
+5. main_ppo.run_ppo                      → fixed rollout + synthetic prefix injection
 
 所有业务逻辑由 integrations 层处理，本 patch set 只负责 thin wrapper 编排。
 """
@@ -15,14 +16,9 @@
 from .attention import patch_megatron_attention
 from .vocab_logprobs import patch_megatron_vocab
 from .nopadding import patch_no_padding_2_padding
+from .fit_hooks import create_run_ppo_patch
 
-# no_padding_2_padding 被 3 个模块用 from...import 直接引用：
-#   verl.workers.utils.padding           — 原定义模块
-#   verl.workers.utils.losses            — ppo_loss 内调用
-#   verl.trainer.distillation.losses     — distillation 内调用
-#   verl.trainer.ppo.ray_trainer         — trainer 侧调用
-# from...import 创建的是模块级属性，setattr 可以更新。
-# 必须对每个引用模块都 patch，否则该模块的局部引用仍指向原函数。
+# no_padding_2_padding 被 4 个模块用 from...import 直接引用
 _NOPADDING_PATCH_MODULES = [
     "verl.workers.utils.padding",
     "verl.workers.utils.losses",
@@ -53,6 +49,12 @@
         patch_factory=patch_megatron_vocab,
         description="vocab_parallel_log_probs → auto logprob restore (verl 0.8.0)",
     ),
+    PatchSpec(
+        module_name="verl.trainer.main_ppo",
+        target_getter=lambda mod: (mod, "run_ppo"),
+        patch_factory=create_run_ppo_patch,
+        description="run_ppo → fixed rollout + synthetic prefix injection (after init_workers)",
+    ),
 ] + [
     PatchSpec(
         module_name=mod_name,
@@ -61,4 +63,4 @@
         description=f"no_padding_2_padding in {mod_name} → PS trimming-aware",
     )
     for mod_name in _NOPADDING_PATCH_MODULES
-]
+]

prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/fit_hooks.py

@@ -0,0 +1,65 @@
+"""Patch: verl.trainer.main_ppo.run_ppo — fixed rollout + synthetic prefix injection.
+
+Uses default-argument capture so Ray can serialise the patched function.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import os
+from typing import Any
+
+
+def _inject_fixed_data(trainer_self: Any) -> None:
+    """Check env vars and patch generate_sequences on the trainer."""
+    json_path = os.environ.get("USE_FIXED_ROLLOUT", None)
+    if json_path:
+        from prefix_sharing.tools.inject_fixed_rollout import patch_fixed_rollout
+        patch_fixed_rollout(trainer_self, json_path=json_path)
+
+    synthetic_json = os.environ.get("USE_SYNTHETIC_PREFIX", None)
+    if synthetic_json:
+        from prefix_sharing.tools.inject_synthetic_prefix import patch_synthetic_prefix
+        batch_size = trainer_self.config.data.get(
+            "gen_batch_size", trainer_self.config.data.train_batch_size
+        )
+        patch_synthetic_prefix(
+            trainer_self,
+            json_path=synthetic_json,
+            batch_size=batch_size,
+            max_prompt_length=trainer_self.config.data.max_prompt_length,
+            max_response_length=trainer_self.config.data.max_response_length,
+        )
+
+
+def _patched_fit(trainer_self: Any, *, _original_fit: Any) -> Any:
+    """Replacement for RayPPOTrainer.fit — inject fixed data then call original."""
+    _inject_fixed_data(trainer_self)
+    return _original_fit(trainer_self)
+
+
+def create_run_ppo_patch(original_run_ppo: Any) -> Any:
+    """Return a patched run_ppo with original captured as default argument."""
+    import verl.trainer.ppo.ray_trainer as rt
+
+    _original_fit_captured = rt.RayPPOTrainer.fit
+
+    def patched_run_ppo(
+        config: Any,
+        task_runner_class: Any = None,
+        *,
+        _orig=original_run_ppo,
+        _fit=_original_fit_captured,
+        _patched=_patched_fit,
+    ) -> None:
+        import verl.trainer.ppo.ray_trainer as _rt
+
+        # Temporarily override fit
+        _rt.RayPPOTrainer.fit = functools.partial(_patched, _original_fit=_fit)
+        try:
+            return _orig(config, task_runner_class)
+        finally:
+            _rt.RayPPOTrainer.fit = _fit
+
+    return patched_run_ppo

prefix-sharing/prefix_sharing/setup/registry.py

@@ -4,6 +4,7 @@
 - 模块已加载且目标存在 → 立即 patch
 - 模块已加载但目标不存在（模块正在 import 中）→ 加入 pending，稍后重试
 - 模块未加载 → import hook 拦截，加载完成后 patch
+- 同一模块多个 spec → import hook 批量处理
 - import hook 完成后立即恢复原始 __import__
 """
 
@@ -38,16 +39,7 @@ def register(cls, spec: PatchSpec) -> None:
 
     @classmethod
     def install_all(cls) -> PatchHandle:
-        """应用所有已注册的 patch。
-
-        三种情况：
-        1. 模块已加载且目标可解析 → 立即 patch
-        2. 模块已加载但目标不可解析（模块正在 import 中）→ 加入 pending
-        3. 模块未加载 → 加入 pending，由 import hook 在加载时 patch
-
-        所有 pending 最终统一由 import hook 处理。
-        import hook 在模块加载完成后才尝试解析目标，确保类定义已完成。
-        """
+        """应用所有已注册的 patch。"""
         shared_records: list[PatchRecord] = []
         mgr = LoggedPatchManager(shared_records)
         pending: list[PatchSpec] = []
@@ -64,9 +56,6 @@ def install_all(cls) -> PatchHandle:
                         f"[PS] Immediately patched {spec.description} (module already loaded)"
                     )
                 except (AttributeError, KeyError):
-                    # 模块已加载但目标不存在——
-                    # 可能是模块正在 import 中，类定义尚未完成。
-                    # 加入 pending，等模块完全加载后再 patch。
                     pending.append(spec)
                     print(
                         f"[PS] Target not yet defined in {spec.module_name}, "
@@ -90,56 +79,50 @@ def _activate_import_hook(
     pending_specs: list[PatchSpec],
     shared_records: list[PatchRecord],
 ) -> None:
-    """对未加载或目标尚未定义的模块，临时拦截 __import__。
-
-    模块加载完成后，尝试解析目标并 patch。如果目标仍然不存在
-    （极端情况：模块被 import 但类在延迟定义），记录 warning 并跳过。
-
-    所有 pending 模块处理完毕后立即恢复原始 __import__。
-    """
+    """拦截 __import__，加载完成后批量 patch。"""
     global _original_import
 
     if _original_import is not None:
         print("[PS] Import hook already active, skipping re-activation")
         return
 
-    lookup = {spec.module_name: spec for spec in pending_specs}
+    # 同一模块可能有多条 spec，用 list 保存
+    lookup: dict[str, list[PatchSpec]] = {}
+    for spec in pending_specs:
+        lookup.setdefault(spec.module_name, []).append(spec)
+
     _original_import = builtins.__import__
 
     def hooked_import(name, globals=None, locals=None, fromlist=(), level=0):
         global _original_import
         module = _original_import(name, globals, locals, fromlist, level)
 
         if name in lookup:
-            spec = lookup.pop(name)
-            # __import__ 在 fromlist 为空时返回顶层包而非子模块，
-            # 必须从 sys.modules 取实际加载的模块对象。
+            specs = lookup.pop(name)
             actual_module = sys.modules[name]
 
-            try:
-                target_obj, attr_name = spec.target_getter(actual_module)
-                original = getattr(target_obj, attr_name)
-                patched = spec.patch_factory(original)
-                setattr(target_obj, attr_name, patched)
-                shared_records.append(
-                    PatchRecord(
-                        target=target_obj,
-                        attr_name=attr_name,
-                        original=original,
-                        replacement=patched,
+            for spec in specs:
+                try:
+                    target_obj, attr_name = spec.target_getter(actual_module)
+                    original = getattr(target_obj, attr_name)
+                    patched = spec.patch_factory(original)
+                    setattr(target_obj, attr_name, patched)
+                    shared_records.append(
+                        PatchRecord(
+                            target=target_obj,
+                            attr_name=attr_name,
+                            original=original,
+                            replacement=patched,
+                        )
+                    )
+                    print(
+                        f"[PS] Auto-patched {spec.description} on import of {name}"
+                    )
+                except (AttributeError, KeyError):
+                    print(
+                        f"[PS] Could not resolve target for {spec.description} "
+                        f"after import of {name}; skipping this patch."
                     )
-                )
-                print(
-                    f"[PS] Auto-patched {spec.description} on import of {name}"
-                )
-            except (AttributeError, KeyError):
-                # 模块已加载但目标仍未定义——
-                # 这种情况极少发生，通常是模块结构异常。
-                print(
-                    f"[PS] Could not resolve target for {spec.description} "
-                    f"after import of {name}; skipping this patch. "
-                    f"The patch target may not exist in this module version."
-                )
 
             if not lookup:
                 builtins.__import__ = _original_import
@@ -150,5 +133,6 @@ def hooked_import(name, globals=None, locals=None, fromlist=(), level=0):
 
     builtins.__import__ = hooked_import
     print(
-        f"[PS] Import hook activated for {len(lookup)} modules: {list(lookup.keys())}"
-    )
+        f"[PS] Import hook activated for {len(lookup)} modules "
+        f"({sum(len(v) for v in lookup.values())} specs): {list(lookup.keys())}"
+    )