From b7ed678e116402f5e4238f29619d62654a2293b0 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Tue, 23 Jun 2026 09:59:09 +0800 Subject: [PATCH] =?UTF-8?q?[feat]=20=E4=B8=BAverl=5Fv080=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?fixed=20rollout/synthetic=20prefix=E6=B3=A8=E5=85=A5=E7=9A=84mo?= =?UTF-8?q?nkey-patch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增fit_hooks.py：通过patch main_ppo.run_ppo在init_workers后注入固定rollout数据 - 使用模块级函数+functools.partial确保Ray序列化兼容，避免PicklingError - 修复registry.py：同一模块支持多个patch spec（dict改为list存储） - 更新__init__.py：注册新patch为第4个patch（共8个） --- .../verl080_mcore0161_ms0160/__init__.py | 18 ++-- .../verl080_mcore0161_ms0160/fit_hooks.py | 65 +++++++++++++++ .../prefix_sharing/setup/registry.py | 82 ++++++++----------- 3 files changed, 108 insertions(+), 57 deletions(-) create mode 100644 prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/fit_hooks.py diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/__init__.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/__init__.py index 3dcd2762..1ea25189 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/__init__.py +++ b/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 -] \ No newline at end of file +] diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/fit_hooks.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/fit_hooks.py new file mode 100644 index 00000000..79de2da5 --- /dev/null +++ b/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 diff --git a/prefix-sharing/prefix_sharing/setup/registry.py b/prefix-sharing/prefix_sharing/setup/registry.py index cefd65a9..9c6e7bee 100644 --- a/prefix-sharing/prefix_sharing/setup/registry.py +++ b/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,20 +79,18 @@ 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): @@ -111,35 +98,31 @@ def hooked_import(name, globals=None, locals=None, fromlist=(), level=0): 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())}" - ) \ No newline at end of file + f"[PS] Import hook activated for {len(lookup)} modules " + f"({sum(len(v) for v in lookup.values())} specs): {list(lookup.keys())}" + )