From b674145f289fe92cb58d9e37207b0895bc3ebee8 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Mon, 10 Aug 2026 17:46:22 +0800 Subject: [PATCH 01/19] refactor: clarify FSDP patch set roles and installation order --- .../setup/patches/verl080_fsdp/__init__.py | 89 +++++++++++++------ .../patches/verl080_fsdp/forward_step.py | 22 ++--- .../test_patch_integrations.py | 18 +++- 3 files changed, 85 insertions(+), 44 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py index 9300f4d2..0855eed9 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py @@ -1,14 +1,12 @@ -"""verl 0.8.0 FSDP patch set. +"""module: prefix_sharing.setup.patches.verl080_fsdp -Patch 目标： -1. FSDPEngineWithLMHead.forward_step → dense FSDP PrefixSharing forward helper -2. 性能验证相关 profiler 注入（不修改 verl 源码） -3. rollout/fixed-data 注入辅助 - -当前 patch set 是 FSDP 开源线的默认入口，可通过兼容矩阵自动选择，也可通过 -``prefix_sharing.setup.install("verl080_fsdp")`` 显式安装。 +verl 0.8.0 FSDP patch set. This is the default entry for the FSDP open-source line. +It can be selected via the compat matrix, or installed explicitly with +``prefix_sharing.setup.install("verl080_fsdp")``. """ +from __future__ import annotations + from prefix_sharing.setup.patch_installer import PatchSpec from .forward_step import ( @@ -26,8 +24,11 @@ PATCH_SET: list[PatchSpec] = [ # ═══════════════════════════════════════════════════════════════ - # 一、fix / feature 注入：PrefixSharing 核心功能与调试辅助 + # PrefixSharing entrypoints # ═══════════════════════════════════════════════════════════════ + + # FSDPEngineWithLMHead.forward_step → verl080_fsdp.patch_fsdp_forward_step + # patch forward_step() for PrefixSharing workflows PatchSpec( module_name="verl.workers.engine.fsdp.transformer_impl", target_getter=lambda mod: ( @@ -35,59 +36,91 @@ "forward_step", ), patch_factory=patch_fsdp_forward_step, - description="FSDPEngineWithLMHead.forward_step → PrefixSharing dense FSDP helper", - eager=True, # verl FSDP engine 仅在 actor 实例化时 lazy-load，必须 eager 触发 + description=("FSDPEngineWithLMHead.forward_step → verl080_fsdp.patch_fsdp_forward_step: " + "patch forward_step() for PrefixSharing workflows"), + # verl FSDP engine is lazy-loaded only when the actor is created; eager is required. + eager=True, ), + + # ═══════════════════════════════════════════════════════════════ + # Debugging functionality, precision and performance + # ═══════════════════════════════════════════════════════════════ + + # RayPPOTrainer.fit → verl080_fsdp.patch_ray_trainer_fit + # patch fit() for capturing rollout results or loading fixed rollout data PatchSpec( module_name="verl.trainer.ppo.ray_trainer", target_getter=lambda mod: (mod.RayPPOTrainer, "fit"), patch_factory=patch_ray_trainer_fit, description=( - "RayPPOTrainer.fit → intercept actor_rollout_wg + async_rollout_manager " - "for PREFIX_SHARING_CAPTURE_ROLLOUT / PREFIX_SHARING_FIXED_ROLLOUT" + "RayPPOTrainer.fit → verl080_fsdp.patch_ray_trainer_fit: " + "patch fit() for capturing rollout results or loading fixed rollout data" ), - eager=True, # ray_trainer is imported by main_ppo at startup; eager ensures patch is in place + # ray_trainer is imported by main_ppo at startup; eager ensures the patch is in place. + eager=True, ), - # ═══════════════════════════════════════════════════════════════ - # 二、性能验证：ProfilerScope 分层 profiling（对应原 verl 源码中的侵入式修改） - # ═══════════════════════════════════════════════════════════════ + + # TrainingWorker.train_mini_batch → step-level ProfilerScope (kind=train) + # patch train_mini_batch() for step-level performance profiling + PatchSpec( + module_name="verl.workers.engine_workers", + target_getter=lambda mod: (mod.TrainingWorker, "train_mini_batch"), + patch_factory=patch_train_mini_batch, + description=("TrainingWorker.train_mini_batch → verl080_fsdp.patch_train_mini_batch: " + "patch train_mini_batch() for step-level performance profiling"), + eager=True, + ), + + # BaseEngine.train_batch is called from train_mini_batch. + # patch train_batch() for optimizer-step performance profiling PatchSpec( module_name="verl.workers.engine.base", target_getter=lambda mod: (mod.BaseEngine, "train_batch"), patch_factory=patch_train_batch, - description="BaseEngine.train_batch → time optimizer_step as PHASE_UPDATE via ProfilerScope", + description=("BaseEngine.train_batch → verl080_fsdp.patch_train_batch: " + "patch train_batch() for optimizer-step performance profiling"), eager=True, ), + + # FSDPEngine.forward_backward_batch → micro-batch ProfilerScope markers + # patch forward_backward_batch() for micro-batch performance profiling PatchSpec( module_name="verl.workers.engine.fsdp.transformer_impl", target_getter=lambda mod: (mod.FSDPEngine, "forward_backward_batch"), patch_factory=patch_forward_backward_batch, - description="FSDPEngine.forward_backward_batch → micro-batch profiling without source edit", + description=("FSDPEngine.forward_backward_batch → verl080_fsdp.patch_forward_backward_batch: " + "patch forward_backward_batch() for micro-batch performance profiling"), eager=True, ), + + # FSDPEngine.forward_backward_batch → dump weight gradients after backward + # patch forward_backward_batch() for precision dumps PatchSpec( module_name="verl.workers.engine.fsdp.transformer_impl", target_getter=lambda mod: (mod.FSDPEngine, "forward_backward_batch"), patch_factory=patch_forward_backward_batch_for_diag_dump, - description="FSDPEngine.forward_backward_batch → dump weight gradients after backward", - eager=True, - ), - PatchSpec( - module_name="verl.workers.engine_workers", - target_getter=lambda mod: (mod.TrainingWorker, "train_mini_batch"), - patch_factory=patch_train_mini_batch, - description="TrainingWorker.train_mini_batch → step-level ProfilerScope (kind=train)", + description=("FSDPEngine.forward_backward_batch → " + "verl080_fsdp.patch_forward_backward_batch_for_diag_dump: " + "patch forward_backward_batch() for weight-gradient diagnostic dumps"), eager=True, ), + + # TrainingWorker.infer_batch → verl080_fsdp.patch_infer_batch + # patch infer_batch() for step-level performance profiling PatchSpec( module_name="verl.workers.engine_workers", target_getter=lambda mod: (mod.TrainingWorker, "infer_batch"), patch_factory=patch_infer_batch, - description="TrainingWorker.infer_batch → step-level ProfilerScope (kind=logp)", + description=("TrainingWorker.infer_batch → verl080_fsdp.patch_infer_batch: " + "patch infer_batch() for step-level performance profiling"), eager=True, ), ] +# ═══════════════════════════════════════════════════════════════ +# PrefixSharing Attention +# ═══════════════════════════════════════════════════════════════ + # Attention patch: directly modify ALL_ATTENTION_FUNCTIONS dict (same approach as verl's PrefixGrouper). # Cannot use PatchSpec because get_interface is not an attribute/key of AttentionInterface. from .attention import install_attention_patch diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 9ce71691..575b91e2 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -1,10 +1,6 @@ -"""Patch: FSDPEngineWithLMHead.forward_step — verl 0.8.0 FSDP 路径。 +"""patch: FSDPEngineWithLMHead.forward_step → verl080_fsdp.patch_fsdp_forward_step -thin wrapper：读取 prefix_sharing_config，优先复用真实 engine 的 -``prepare_model_inputs`` / ``prepare_model_outputs``，在 forward 期间注入 -PrefixSharing runtime，并在输出阶段做 interior / prefix-last restore。 -本 patch 已覆盖 dense 2D 与 verl remove-padding 后的 jagged NestedTensor -形态；Ulysses SP、fused kernels 等未验证形态仍在配置校验阶段显式拒绝。 +forward_step wrapper for PrefixSharing under verl 0.8.0 + FSDP. """ from __future__ import annotations @@ -14,7 +10,7 @@ def patch_fsdp_forward_step(original_forward_step: Any) -> Any: - """创建 FSDPEngineWithLMHead.forward_step 的 patch wrapper。""" + """Build the patched FSDPEngineWithLMHead.forward_step wrapper.""" # Patch _CheckpointFrame.check_recomputed_tensors_match and # _internal_assert to no-op. @@ -74,7 +70,7 @@ def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forwar micro_batch = micro_batch.to(get_device_id()) except Exception: - # 本地单测使用 plain dict / fake engine，不依赖 verl device helper。 + # Local unit tests use plain dict / fake engine without verl device helpers. pass ulysses_sp_size = _read_runtime_value( @@ -212,7 +208,7 @@ def _forward_step_with_engine_prepare( diagnostic_tag, ) - # 获取模型层数以支持 per-layer diagnostic dump + # Read layer count for per-layer diagnostic dumps. _diag_num_layers = int(getattr( getattr(getattr(self, "module", None), "config", None), "num_hidden_layers", 0)) or 0 @@ -359,9 +355,9 @@ def _call_original_like_engine(self: Any, micro_batch: Any, loss_function: Any, import torch from contextlib import nullcontext - # 对齐 verl 原生 forward_step：先把 micro_batch 搬到 device（disable 路径绕过了 - # patched_forward_step 里那段 .to(device)，这里补上，否则 prepare_model_outputs - # 里 logits/temperature device 不一致）。 + # Match native verl forward_step: move micro_batch to device first. + # The disable / no-sharing path skips the .to(device) in patched_forward_step; + # without it, logits/temperature can land on different devices in prepare_model_outputs. if hasattr(micro_batch, "to"): try: from verl.utils.device import get_device_id @@ -370,7 +366,7 @@ def _call_original_like_engine(self: Any, micro_batch: Any, loss_function: Any, pass model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch) - # DIAG_DUMP: ON path dump原始full input_ids（suffix-only dump会缺失prefix tokens） + # DIAG_DUMP: dump original full input_ids (suffix-only dumps miss prefix tokens). import os as _ps_diag_fwd_ids if _ps_diag_fwd_ids.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: _dump_full_input_ids_only(micro_batch, "train") diff --git a/prefix-sharing/tests/integrated_test/test_patch_integrations.py b/prefix-sharing/tests/integrated_test/test_patch_integrations.py index 45ae2bff..23c83b9f 100644 --- a/prefix-sharing/tests/integrated_test/test_patch_integrations.py +++ b/prefix-sharing/tests/integrated_test/test_patch_integrations.py @@ -50,9 +50,21 @@ def test_setup_can_load_explicit_verl080_fsdp_patch_set(): assert "FSDPEngineWithLMHead.forward_step" in patch_set[0].description assert patch_set[1].module_name == "verl.trainer.ppo.ray_trainer" assert "RayPPOTrainer.fit" in patch_set[1].description - # Diag-dump wrapper is installed on top of the profiler wrapper. - assert patch_set[4].module_name == "verl.workers.engine.fsdp.transformer_impl" - assert "dump weight gradients" in patch_set[4].description + forward_backward_specs = [ + spec + for spec in patch_set + if spec.module_name == "verl.workers.engine.fsdp.transformer_impl" + and "forward_backward_batch" in spec.description + ] + # The diagnostic wrapper must wrap the profiler wrapper. + assert [spec.description for spec in forward_backward_specs] == [ + "FSDPEngine.forward_backward_batch → " + "verl080_fsdp.patch_forward_backward_batch: " + "patch forward_backward_batch() for micro-batch performance profiling", + "FSDPEngine.forward_backward_batch → " + "verl080_fsdp.patch_forward_backward_batch_for_diag_dump: " + "patch forward_backward_batch() for weight-gradient diagnostic dumps", + ] From f15e8808f72088b12d4d0a3b5a7f08b99c084266 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Mon, 10 Aug 2026 17:42:25 +0800 Subject: [PATCH 02/19] =?UTF-8?q?fix(fsdp):=20=E6=97=A0=E6=9D=A1=E4=BB=B6?= =?UTF-8?q?=E6=89=A7=E8=A1=8C=20PS=20context=20cleanup=EF=BC=8C=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=20audit=20=E6=97=A5=E5=BF=97=E7=BC=BA=E5=A4=B1?= =?UTF-8?q?=E4=B8=8E=20KV=20store=20=E6=B3=84=E6=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit forward_backward_batch wrapper 中 ctx_cleanup()（含 _log_prefix_sharing_audit 打印 + ctx.store.close()）被误关在 PREFIX_SHARING_DIAG_DUMP 条件块内， 导致正常训练（未开 diag dump）时： - audit 日志（store_count/reuse_hit/matches_expected/sharing_group）完全不输出 - KV store 不 close，多步训练存在 ContextVar/KV 内存累积风险 将 ctx_cleanup 调用移出 DIAG_DUMP 条件块，在 backward 完成后无条件执行 （保留 not forward_only 条件，与原设计一致）。DIAG_DUMP 块只保留 weight grad dump + per-layer grad hook 清理。 回归验证：4090 单卡/多卡 PS ON smoke 测试发现 audit 0 次输出，修复后预期 出现 [PS][audit] summary + 24 层 layer runtime。 Co-Authored-By: Claude Opus 4.8 --- .../setup/patches/verl080_fsdp/forward_step.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 575b91e2..7a70b347 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -580,6 +580,16 @@ def patch_forward_backward_batch_for_diag_dump( def wrapped(self: Any, data: Any, loss_function: Any, forward_only: bool = False) -> Any: result = original_forward_backward_batch(self, data, loss_function, forward_only) + # 无条件清理 PrefixSharing runtime context：打印 audit 日志 + 关闭 KV store。 + # 之前该清理被误关在 PREFIX_SHARING_DIAG_DUMP 条件块内，导致正常训练时 + # audit 日志不输出、KV store 不 close（多步训练存在内存累积风险）。 + if not forward_only: + ctx_cleanup = getattr(self.module, "_ps_ctx_cleanup", None) + if ctx_cleanup is not None: + ctx_cleanup() + delattr(self.module, "_ps_ctx_cleanup") + + # 诊断 dump 专用：dump weight gradients + 清理 per-layer attention grad hooks。 if not forward_only and os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: tag = "train" if self.module.training else "old" print( @@ -591,12 +601,6 @@ def wrapped(self: Any, data: Any, loss_function: Any, forward_only: bool = False dump_weight_grads_verl080(self.module, tag) - # Clean up PrefixSharing context if the root backward hook did not fire. - ctx_cleanup = getattr(self.module, "_ps_ctx_cleanup", None) - if ctx_cleanup is not None: - ctx_cleanup() - delattr(self.module, "_ps_ctx_cleanup") - # Remove per-layer attention gradient hooks. for module in self.module.modules(): try: From e32d5f23ade16e0637fdaaef78cec5cba5115c6d Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Mon, 10 Aug 2026 17:58:19 +0800 Subject: [PATCH 03/19] refactor: hoist FSDP forward_step imports and align config reader names Move hot-path imports to module scope, read PS config via verl_utils, and rename the prefix_grouper helper to read_ps_config_from_prefix_grouper. Co-authored-by: Cursor --- .../prefix_sharing/integrations/verl_utils.py | 4 ++-- .../setup/patches/verl080_fsdp/forward_step.py | 13 ++++++------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/verl_utils.py b/prefix-sharing/prefix_sharing/integrations/verl_utils.py index 1a51e2fc..5cae37b6 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_utils.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_utils.py @@ -28,10 +28,10 @@ def read_ps_config_from_engine_config(engine_config: Any) -> Any | None: if explicit_config is not None: return explicit_config - return _prefix_sharing_config_from_prefix_grouper(engine_config) + return read_ps_config_from_prefix_grouper(engine_config) -def _prefix_sharing_config_from_prefix_grouper(engine_config: Any) -> dict[str, Any] | None: +def read_ps_config_from_prefix_grouper(engine_config: Any) -> dict[str, Any] | None: use_prefix_grouper = _read_actor_value(engine_config, "use_prefix_grouper", False) if not use_prefix_grouper: return None diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 7a70b347..b59f0b7b 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -8,9 +8,14 @@ import os from typing import Any, Callable +import torch.utils.checkpoint as _ckpt + +from prefix_sharing.core.config import PrefixSharingConfig +from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config +from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope + def patch_fsdp_forward_step(original_forward_step: Any) -> Any: - """Build the patched FSDPEngineWithLMHead.forward_step wrapper.""" # Patch _CheckpointFrame.check_recomputed_tensors_match and # _internal_assert to no-op. @@ -18,7 +23,6 @@ def patch_fsdp_forward_step(original_forward_step: Any) -> Any: # computation graph, causing the saved-tensor count mismatch detected by # these methods. The recomputed values are numerically correct — the count # difference is benign. Bypass both checks so ON-path training completes. - import torch.utils.checkpoint as _ckpt # Apply once, globally. if not getattr(patch_fsdp_forward_step, "_cp_patched", False): _ckpt._CheckpointFrame.check_recomputed_tensors_match = lambda self, gid: None # type: ignore[method-assign] @@ -27,10 +31,6 @@ def patch_fsdp_forward_step(original_forward_step: Any) -> Any: patch_fsdp_forward_step._cp_patched = True def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forward_only: bool): - from prefix_sharing.core.config import PrefixSharingConfig - from prefix_sharing.integrations.verl_mcore import read_ps_config_from_engine_config - from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope - raw_config = read_ps_config_from_engine_config(self.engine_config) ps_config = PrefixSharingConfig.from_raw(raw_config) if not ps_config.enable_prefix_sharing: @@ -371,7 +371,6 @@ def _call_original_like_engine(self: Any, micro_batch: Any, loss_function: Any, if _ps_diag_fwd_ids.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: _dump_full_input_ids_only(micro_batch, "train") - autocast_dtype = getattr(self, "_autocast_dtype", torch.float32) autocast_dtype = getattr(self, "_autocast_dtype", torch.float32) device_name = _read_device_name() autocast_ctx = ( From dc7746970b5b8a49d5777d3fdff4dcffc2cd36a5 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Wed, 12 Aug 2026 16:34:33 +0800 Subject: [PATCH 04/19] refactor: rename FSDP plan/trim helper and English-ize config.validate Rename build_prefix_sharing_micro_batch_fsdp to plan_and_trim_microbatch_fsdp with clearer micro_batch/ps_config args, and rewrite PrefixSharingConfig.validate error messages in English while keeping integrate_mode fallback semantics. Co-authored-by: Cursor --- prefix-sharing/prefix_sharing/core/config.py | 127 +++++++++--------- .../prefix_sharing/integrations/__init__.py | 4 +- .../prefix_sharing/integrations/verl_fsdp.py | 48 +++---- .../patches/verl080_fsdp/forward_step.py | 10 +- .../tests/unit_test/test_verl_fsdp_adapter.py | 20 +-- .../test_verl_fsdp_ch4_functional.py | 18 +-- .../unit_test/test_verl_fsdp_ch4_precision.py | 16 +-- 7 files changed, 123 insertions(+), 120 deletions(-) diff --git a/prefix-sharing/prefix_sharing/core/config.py b/prefix-sharing/prefix_sharing/core/config.py index 71d0a186..044bb25d 100644 --- a/prefix-sharing/prefix_sharing/core/config.py +++ b/prefix-sharing/prefix_sharing/core/config.py @@ -103,118 +103,117 @@ def from_raw(cls, raw: Any) -> "PrefixSharingConfig": return cls(**values) def validate(self, model_config: Any | None = None, integrate_mode: str | None = None) -> None: - """Validate phase-1 constraints against a model/config object. + """Validate constraints against model_config. Args: - model_config: Mapping or object with Megatron-like attributes. + model_config: Model attributes. integrate_mode: Optional integration mode name. """ if not self.enable_prefix_sharing: return - if self.detector != "trie": - raise PrefixSharingConfigError("phase 1 supports only detector='trie'") + + # detector guard + supported_detectors = {"trie"} + if self.detector not in supported_detectors: + raise PrefixSharingConfigError( + f"detector='{self.detector}' is not supported. Supported detectors: {supported_detectors}" + ) + + # backend guard supported_backends = {"torch_ref", "flash_atten_gpu", "flash_atten_npu"} if self.backend not in supported_backends: raise PrefixSharingConfigError( - f"backend='{self.backend}' is not supported. " - f"Supported backends: {supported_backends}" + f"backend='{self.backend}' is not supported. Supported backends: {supported_backends}" ) - if self.boundary_strategy != "prefix_last_restore": + + # boundary_strategy guard + supported_boundary_strategies = {"prefix_last_restore"} + if self.boundary_strategy not in supported_boundary_strategies: raise PrefixSharingConfigError( - "phase 1 currently implements only " - "boundary_strategy='prefix_last_restore'; future strategies may include " - "'boundary_token' and 'strict_suffix'" + f"boundary_strategy='{self.boundary_strategy}' is not supported. Supported strategies: {supported_boundary_strategies}" ) + + # min_prefix_len guard if self.min_prefix_len < 1: - raise PrefixSharingConfigError("min_prefix_len must be >= 1") + raise PrefixSharingConfigError(f"min_prefix_len='{self.min_prefix_len}' is not supported. Supported min_prefix_len: >= 1") if self.min_group_size < 2: - raise PrefixSharingConfigError("min_group_size must be >= 2") + raise PrefixSharingConfigError(f"min_group_size='{self.min_group_size}' is not supported. Supported min_group_size: >= 2") - active_mode = integrate_mode or self.integrate_mode + # integrate_mode guard supported_integrate_modes = {"verl_megatron_actor", "verl_fsdp"} + active_mode = integrate_mode or self.integrate_mode if active_mode not in supported_integrate_modes: raise PrefixSharingConfigError( - "phase 1 supports only integrate_mode in " - f"{sorted(supported_integrate_modes)}" + f"integrate_mode='{active_mode}' is not supported. Supported integrate_modes: {supported_integrate_modes}" ) + # model_type guard + supported_model_types = {"text_only_causal_lm"} model_type = _read_config_value(model_config, "model_type", "text_only_causal_lm") - if self.model_type == "text_only_causal_lm" and model_type != "text_only_causal_lm": + if self.model_type in supported_model_types and model_type not in supported_model_types: raise PrefixSharingConfigError( - f"[Config Error] 当前模型类型 '{model_type}' 不支持当前阶段。" - f"Phase 1 仅支持 model_type='text_only_causal_lm' (纯文本因果语言模型)，" - f"请使用支持的模型类型或禁用 prefix sharing。" + f"model_type='{model_type}' is not supported. Supported model_types: {supported_model_types}" ) + + # verl_fsdp guard if active_mode == "verl_fsdp": ulysses_sp_size = _read_config_value(model_config, "ulysses_sequence_parallel_size", 1) use_fused_kernels = _read_config_value(model_config, "use_fused_kernels", False) if int(ulysses_sp_size) != 1: raise PrefixSharingConfigError( - f"[Config Error] verl_fsdp 当前不支持 ulysses_sequence_parallel_size={ulysses_sp_size}。" - "请关闭 Ulysses SP 或等待专门适配。" + f"ulysses_sequence_parallel_size={ulysses_sp_size} is not supported. " + "Supported ulysses_sequence_parallel_size: 1" ) if use_fused_kernels: raise PrefixSharingConfigError( - "[Config Error] verl_fsdp 当前不支持 use_fused_kernels=True。" - "请关闭 fused kernels 或等待专门适配。" + f"use_fused_kernels={use_fused_kernels} is not supported. " + "Supported use_fused_kernels: False" ) return + # parallel strategy guards pp_size = _read_config_value( model_config, "pipeline_model_parallel_size", _read_config_value(model_config, "pipeline_parallel_size", 1), ) - virtual_pp_size = _read_config_value(model_config, "virtual_pipeline_model_parallel_size", None) - num_layers_per_virtual_pipeline_stage = _read_config_value( - model_config, - "num_layers_per_virtual_pipeline_stage", - None, - ) - cp_size = _read_config_value( - model_config, - "context_parallel_size", - self.supported_cp_size, - ) - rope_fusion = _read_config_value(model_config, "apply_rope_fusion", False) - fused_qkv_rope = _read_config_value(model_config, "fused_single_qkv_rope", False) - if int(pp_size) < 1: raise PrefixSharingConfigError( - f"[Config Error] pipeline_model_parallel_size={pp_size} 不合法。" - f"pipeline_model_parallel_size 必须 >= 1，" - f"请修改为合法物理 PP 大小，或禁用 prefix sharing。" + f"pipeline_model_parallel_size={pp_size} is not supported. " + "Supported pipeline_model_parallel_size: >= 1" ) + virtual_pp_size = _read_config_value(model_config, "virtual_pipeline_model_parallel_size", 1) if virtual_pp_size not in (None, 1): raise PrefixSharingConfigError( - f"[Config Error] virtual_pipeline_model_parallel_size={virtual_pp_size} 不支持当前阶段。" - f"当前仅支持物理 pipeline parallel，不支持 virtual pipeline parallel，" - f"请关闭 virtual PP 或禁用 prefix sharing。" + f"virtual_pipeline_model_parallel_size={virtual_pp_size} is not supported. " + "Supported virtual_pipeline_model_parallel_size: 1" ) + num_layers_per_virtual_pipeline_stage = _read_config_value( + model_config, "num_layers_per_virtual_pipeline_stage", None + ) if num_layers_per_virtual_pipeline_stage is not None: raise PrefixSharingConfigError( - "[Config Error] num_layers_per_virtual_pipeline_stage 不支持当前阶段。" - "当前仅支持物理 pipeline parallel，不支持 virtual pipeline parallel，" - "请关闭 virtual PP 或禁用 prefix sharing。" + f"num_layers_per_virtual_pipeline_stage={num_layers_per_virtual_pipeline_stage} " + "is not supported. Supported num_layers_per_virtual_pipeline_stage: None" ) + cp_size = _read_config_value(model_config, "context_parallel_size", self.supported_cp_size) if cp_size != self.supported_cp_size: raise PrefixSharingConfigError( - f"[Config Error] context_parallel_size={cp_size} 不支持当前阶段。" - f"Phase 1 仅支持 context_parallel_size=1 (无上下文并行)，" - f"请修改配置将 CP 大小设为 1，或禁用 prefix sharing。" + f"context_parallel_size={cp_size} is not supported. " + f"Supported context_parallel_size: {self.supported_cp_size}" ) + rope_fusion = _read_config_value(model_config, "apply_rope_fusion", False) if not self.supported_rope_fusion and rope_fusion: raise PrefixSharingConfigError( - "[Config Error] apply_rope_fusion=True 不支持当前阶段。" - "Phase 1 要求关闭 rope fusion (apply_rope_fusion=False)，" - "请修改配置或禁用 prefix sharing。" + f"apply_rope_fusion={rope_fusion} is not supported. " + "Supported apply_rope_fusion: False" ) + fused_qkv_rope = _read_config_value(model_config, "fused_single_qkv_rope", False) if not self.supported_fused_qkv_rope and fused_qkv_rope: raise PrefixSharingConfigError( - "[Config Error] fused_single_qkv_rope=True 不支持当前阶段。" - "Phase 1 要求关闭 fused QKV rope (fused_single_qkv_rope=False)，" - "请修改配置或禁用 prefix sharing。" + f"fused_single_qkv_rope={fused_qkv_rope} is not supported. " + "Supported fused_single_qkv_rope: False" ) def validate_for_engine( @@ -222,16 +221,16 @@ def validate_for_engine( use_remove_padding: bool = True, integrate_mode: str = "verl_megatron_actor", ) -> None: - """Validate phase-1 constraints for verl engine 架构（verl 0.8.0+）。 + """Validate phase-1 constraints for the verl engine (verl 0.8.0+). - 与 validate() 不同，此方法从 engine_config 而非 model_config 读取配置。 - 用于 setup/patches 中的 forward_step patch，此时只有 engine_config - 可用（self.engine_config），而非 Megatron TransformerConfig。 + Unlike ``validate()``, this reads constraints from engine-facing flags + rather than Megatron ``model_config``. Used by setup/patches + ``forward_step`` when only ``self.engine_config`` is available. """ if not self.enable_prefix_sharing: return - # 基础校验 + # Basic field checks. if self.detector != "trie": raise PrefixSharingConfigError("phase 1 supports only detector='trie'") if self.backend not in {"torch_ref", "flash_atten_gpu", "flash_atten_npu"}: @@ -248,10 +247,10 @@ def validate_for_engine( if self.min_group_size < 2: raise PrefixSharingConfigError("min_group_size must be >= 2") - # THD packed layout 需要 use_remove_padding + # THD packed layout requires use_remove_padding. if not use_remove_padding: raise PrefixSharingConfigError( - "[Config Error] Phase 1 THD 路径要求 use_remove_padding=True。" - "BSHD 路径 (use_remove_padding=False) 尚未在当前 patch 中支持，" - "请启用 use_remove_padding 或使用 BSHD 专用 patch set。" + "[Config Error] Phase 1 THD path requires use_remove_padding=True. " + "BSHD (use_remove_padding=False) is not supported by the current " + "patch set. Enable use_remove_padding or use a BSHD-specific patch set." ) diff --git a/prefix-sharing/prefix_sharing/integrations/__init__.py b/prefix-sharing/prefix_sharing/integrations/__init__.py index 2b03ff05..c00326da 100644 --- a/prefix-sharing/prefix_sharing/integrations/__init__.py +++ b/prefix-sharing/prefix_sharing/integrations/__init__.py @@ -16,7 +16,7 @@ ) from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - build_prefix_sharing_micro_batch_fsdp, + plan_and_trim_microbatch_fsdp, forward_prefix_sharing_fsdp_micro_batch, restore_prefix_sharing_outputs_2d, ) @@ -38,7 +38,7 @@ "prefix_attention", "get_megatron_parallel_info", "PrefixSharingFSDPAttentionRuntime", - "build_prefix_sharing_micro_batch_fsdp", + "plan_and_trim_microbatch_fsdp", "forward_prefix_sharing_fsdp_micro_batch", "restore_prefix_sharing_outputs_2d", ] diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index fb5b83aa..758a9aec 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -1,7 +1,9 @@ -"""verl FSDP integration helpers for PrefixSharing. +"""prefix_sharing.integrations.verl_fsdp + +verl FSDP integration helpers for PrefixSharing. The FSDP path follows the same public shape as the Megatron integration: -``build_*`` returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``. +``plan_and_trim_microbatch_fsdp`` returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``. The helpers stay framework-light enough for CPU tests, while the explicit ``verl080_fsdp`` patch set wires them into ``FSDPEngineWithLMHead.forward_step``. """ @@ -114,7 +116,7 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any, def forward_prefix_sharing_fsdp_micro_batch( micro_batch: Any, model: Any, - config: PrefixSharingConfig, + ps_config: PrefixSharingConfig, *, model_config: Any | None = None, backend: Any | None = None, @@ -134,9 +136,9 @@ def forward_prefix_sharing_fsdp_micro_batch( own ``prepare_model_inputs`` / ``prepare_model_outputs`` path. """ - trimmed_micro_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp( + trimmed_micro_batch, runtime_state = plan_and_trim_microbatch_fsdp( micro_batch, - config, + ps_config, model_config=model_config, backend=backend, ) @@ -173,33 +175,33 @@ def forward_prefix_sharing_fsdp_micro_batch( return output -def build_prefix_sharing_micro_batch_fsdp( - batch: Any, - config: PrefixSharingConfig, - *, +def plan_and_trim_microbatch_fsdp( + micro_batch: Any, + ps_config: PrefixSharingConfig, model_config: Any | None = None, backend: Any | None = None, ) -> tuple[Any, PrefixSharingRuntimeState | None]: - """Build a trimmed FSDP micro-batch and PrefixSharing runtime state. + """Plan prefix sharing and trim one FSDP micro-batch. + Returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``. This helper is intentionally framework-light: it accepts dense 2D ``input_ids``/``attention_mask`` or jagged NestedTensor ``input_ids`` from - verl remove-padding, and returns the original batch unchanged when prefix - sharing is disabled or no reusable prefix is detected. + verl remove-padding, and returns the original micro-batch unchanged when + prefix sharing is disabled or no reusable prefix is detected. """ - if not config.enable_prefix_sharing: - return batch, None - config.validate(model_config=model_config, integrate_mode="verl_fsdp") + if not ps_config.enable_prefix_sharing: + return micro_batch, None + ps_config.validate(model_config=model_config, integrate_mode="verl_fsdp") - input_ids = batch["input_ids"] + input_ids = micro_batch["input_ids"] is_nested_input = _is_nested_tensor(input_ids) if is_nested_input: sequences = _extract_seq_from_nested_tensor(input_ids) valid_indices = None attention_mask = None else: - attention_mask = batch["attention_mask"].to(bool) + attention_mask = micro_batch["attention_mask"].to(bool) if input_ids.dim() != 2 or attention_mask.dim() != 2: raise RuntimeError("prefix sharing FSDP path expects 2D or jagged NestedTensor input_ids") if input_ids.shape != attention_mask.shape: @@ -210,19 +212,19 @@ def build_prefix_sharing_micro_batch_fsdp( for row in range(input_ids.shape[0]) ] sequences = _extract_sequences_async(input_ids, valid_indices) - prefix_sharing_plan = PrefixSharingPlanner(config).plan(sequences) + prefix_sharing_plan = PrefixSharingPlanner(ps_config).plan(sequences) if not prefix_sharing_plan.has_sharing: - return batch, None + return micro_batch, None if is_nested_input: - trimmed_micro_batch = _trim_nested_batch(batch, prefix_sharing_plan) + trimmed_micro_batch = _trim_nested_batch(micro_batch, prefix_sharing_plan) kept_position_rows = _collect_kept_position_rows( trimmed_micro_batch, prefix_sharing_plan, is_nested_tensor=True, ) else: - trimmed_micro_batch = _clone_batch(batch) + trimmed_micro_batch = _clone_batch(micro_batch) trimmed_attention_mask = attention_mask.clone() trimmed_attention_mask[:] = False @@ -238,7 +240,7 @@ def build_prefix_sharing_micro_batch_fsdp( trimmed_loss_mask[:] = False for row, indices in enumerate(valid_indices): keep_start, keep_end = prefix_sharing_plan.loss_mask_keep_ranges[row] - trimmed_loss_mask[row, indices[keep_start:keep_end]] = batch["loss_mask"][row, indices[keep_start:keep_end]].to(bool) + trimmed_loss_mask[row, indices[keep_start:keep_end]] = micro_batch["loss_mask"][row, indices[keep_start:keep_end]].to(bool) trimmed_micro_batch["loss_mask"] = trimmed_loss_mask kept_position_rows = _collect_kept_position_rows( trimmed_micro_batch, @@ -253,7 +255,7 @@ def build_prefix_sharing_micro_batch_fsdp( ) runtime_state = PrefixSharingRuntimeState( prefix_sharing_plan=prefix_sharing_plan, - attention_backend=get_backend_instance(config, backend), + attention_backend=get_backend_instance(ps_config, backend), packed_batch_layout=packed_batch_layout, parallel_info=MegatronParallelInfo(), kept_position_ids=trimmed_micro_batch.get("position_ids"), diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index b59f0b7b..5077506b 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -31,8 +31,10 @@ def patch_fsdp_forward_step(original_forward_step: Any) -> Any: patch_fsdp_forward_step._cp_patched = True def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forward_only: bool): - raw_config = read_ps_config_from_engine_config(self.engine_config) - ps_config = PrefixSharingConfig.from_raw(raw_config) + + raw_config = read_ps_config_from_engine_config(self.engine_config) # framework-level config + ps_config = PrefixSharingConfig.from_raw(raw_config) # PrefixSharing-level config + if not ps_config.enable_prefix_sharing: # Memory sampling is managed by the step-level ProfilerScope; this # path records only the model-forward phase. @@ -159,7 +161,7 @@ def _forward_step_with_engine_prepare( from prefix_sharing.integrations.context import create_prefix_sharing_context from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime - from prefix_sharing.integrations.verl_fsdp import build_prefix_sharing_micro_batch_fsdp + from prefix_sharing.integrations.verl_fsdp import plan_and_trim_microbatch_fsdp from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope @@ -173,7 +175,7 @@ def _forward_step_with_engine_prepare( if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: _dump_full_input_ids_only(micro_batch, "train") - trimmed_micro_batch, ps_state = build_prefix_sharing_micro_batch_fsdp( + trimmed_micro_batch, ps_state = plan_and_trim_microbatch_fsdp( micro_batch, ps_config, model_config={ diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py index 1b9a558d..fc79152d 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py @@ -9,7 +9,7 @@ from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - build_prefix_sharing_micro_batch_fsdp, + plan_and_trim_microbatch_fsdp, forward_prefix_sharing_fsdp_micro_batch, restore_prefix_sharing_outputs_2d, ) @@ -132,7 +132,7 @@ def _baseline_attention(query, key, value): return torch.einsum("blmh,bmhd->blhd", probs, value) -def test_build_prefix_sharing_micro_batch_fsdp_returns_trimmed_batch_and_runtime_state(): +def test_plan_and_trim_microbatch_fsdp_returns_trimmed_batch_and_runtime_state(): config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { "input_ids": torch.tensor( @@ -172,7 +172,7 @@ def test_build_prefix_sharing_micro_batch_fsdp_returns_trimmed_batch_and_runtime ), } - trimmed_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed_batch, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) assert runtime_state is not None assert isinstance(runtime_state, PrefixSharingRuntimeState) @@ -193,7 +193,7 @@ def test_build_prefix_sharing_micro_batch_fsdp_returns_trimmed_batch_and_runtime assert torch.equal(trimmed_batch["position_ids"][1, 3:6], torch.tensor([3, 4, 5])) -def test_build_prefix_sharing_micro_batch_fsdp_returns_none_when_no_sharing(): +def test_plan_and_trim_microbatch_fsdp_returns_none_when_no_sharing(): config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { "input_ids": torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long), @@ -201,13 +201,13 @@ def test_build_prefix_sharing_micro_batch_fsdp_returns_none_when_no_sharing(): "position_ids": torch.tensor([[0, 1, 2], [0, 1, 2]], dtype=torch.long), } - returned_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config) + returned_batch, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) assert returned_batch is batch assert runtime_state is None -def test_build_prefix_sharing_micro_batch_fsdp_trims_nested_remove_padding_batch(): +def test_plan_and_trim_microbatch_fsdp_trims_nested_remove_padding_batch(): if not hasattr(torch, "nested"): pytest.skip("torch.nested is unavailable") config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) @@ -235,7 +235,7 @@ def test_build_prefix_sharing_micro_batch_fsdp_trims_nested_remove_padding_batch ), } - trimmed_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed_batch, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) assert runtime_state is not None plan = runtime_state.prefix_sharing_plan @@ -273,7 +273,7 @@ def test_restore_prefix_sharing_outputs_2d_restores_interior_last_logits_entropy dtype=torch.long, ), } - _, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config) + _, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) assert runtime_state is not None vocab = 5 @@ -345,7 +345,7 @@ def test_prefix_sharing_fsdp_attention_runtime_scatter_dense_outputs(): dtype=torch.long, ), } - _, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config) + _, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) assert runtime_state is not None torch.manual_seed(1) @@ -686,7 +686,7 @@ def test_prefix_sharing_fsdp_attention_runtime_supports_packed_single_batch_shap "attention_mask": torch.ones(2, 5, dtype=torch.bool), "position_ids": torch.tensor([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long), } - _, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config) + _, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) assert runtime_state is not None torch.manual_seed(2031) diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py index 7d7710cc..a57e5c73 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py @@ -21,7 +21,7 @@ from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - build_prefix_sharing_micro_batch_fsdp, + plan_and_trim_microbatch_fsdp, restore_prefix_sharing_outputs_2d, ) @@ -52,7 +52,7 @@ def test_scenario1_two_samples_share_arbitrary_prefix(): """A B C D E / A B C X Y -> provider=0, reuser=1, prefix_len=3.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5], [1, 2, 3, 10, 11]]) - trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -76,7 +76,7 @@ def test_scenario2_multiple_reusers_different_prefix_and_suffix_lens(): [1, 2, 20, 21], # reuser, prefix 2, suffix 2 (len 4) ] ) - trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -106,7 +106,7 @@ def test_scenario3_same_provider_serves_multiple_reusers(): [1, 2, 30, 40, 50, 60], # reuser 3 (prefix 2) ] ) - trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -128,7 +128,7 @@ def test_scenario4_chain_reuse_reuser_becomes_provider(): [1, 2, 3, 4, 10, 20, 30], # row2: reuser of row1 (prefix 5) ] ) - trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -146,7 +146,7 @@ def test_scenario5_no_shareable_prefix_returns_none(): """Completely different sequences -> fallback.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4], [5, 6, 7, 8]]) - returned, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + returned, state = plan_and_trim_microbatch_fsdp(batch, config) assert returned is batch assert state is None @@ -156,7 +156,7 @@ def test_scenario5b_single_sample_returns_none(): """Single sample cannot share -> fallback.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5]]) - returned, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + returned, state = plan_and_trim_microbatch_fsdp(batch, config) assert returned is batch assert state is None @@ -205,7 +205,7 @@ def test_scenario7_position_ids_preserve_absolute_positions(): """Reuser suffix position_ids must reflect original absolute positions.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 10, 11, 12, 13]]) - trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -224,7 +224,7 @@ def test_scenario7_position_ids_preserve_absolute_positions(): def test_attention_runtime_records_stats_per_layer(): config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5], [1, 2, 3, 10, 11]]) - _, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + _, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None torch.manual_seed(42) diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py index 07824d69..8bd7c91c 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py @@ -26,7 +26,7 @@ from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - build_prefix_sharing_micro_batch_fsdp, + plan_and_trim_microbatch_fsdp, restore_prefix_sharing_outputs_2d, ) @@ -67,7 +67,7 @@ def test_reuser_suffix_attention_matches_baseline(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan prefix_len = plan.prefix_lens[1] # 3 @@ -143,7 +143,7 @@ def test_gradient_flows_through_provider_prefix_kv(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None H, D = 2, 4 @@ -184,7 +184,7 @@ def test_restore_interior_prefix_matches_provider(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - _, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + _, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is not None # Simulate trimmed output: provider full, reuser has only suffix (prefix zeroed) @@ -238,7 +238,7 @@ def test_restore_prefix_last_recomputed_with_provider_logits_and_reuser_label(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - _, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + _, state = plan_and_trim_microbatch_fsdp(batch, config) plan = state.prefix_sharing_plan prefix_len = plan.prefix_lens[1] # 3 @@ -275,7 +275,7 @@ def test_baseline_matches_when_no_sharing(): [[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.long ), } - returned, state = build_prefix_sharing_micro_batch_fsdp(batch, config) + returned, state = plan_and_trim_microbatch_fsdp(batch, config) assert state is None assert returned is batch # exact same object -> no transformation @@ -308,8 +308,8 @@ def test_padding_positions_do_not_contribute(): } H, D = 2, 4 - _, state_unpadded = build_prefix_sharing_micro_batch_fsdp(batch_unpadded, config) - _, state_padded = build_prefix_sharing_micro_batch_fsdp(batch_padded, config) + _, state_unpadded = plan_and_trim_microbatch_fsdp(batch_unpadded, config) + _, state_padded = plan_and_trim_microbatch_fsdp(batch_padded, config) assert state_unpadded is not None and state_padded is not None # Same plan semantics regardless of padding From 92b10894f2e3cf5b39ac103ef3110b6aff81ae6a Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:20:06 +0800 Subject: [PATCH 05/19] refactor(verl_utils): expose public trim helpers and add dense trim path - Rename private `_trim_*` / `_extract_*` / `_clone_batch` helpers in `verl_utils` to public API without leading underscore. - Add `trim_redundant_prefix_in_dense_tensor` for the FSDP 2D mask-only trim path; return `kept_position_rows` directly from trim helpers so callers do not re-extract position rows. - Update FSDP, Megatron and setup patch callers to use the new helpers and import `is_nested_tensor` from `verl_utils` instead of `verl_mcore`. - Default `PackedBatchLayout.from_kept_position_rows(align_size=1)` so FSDP callers can omit the argument. - Add `test_verl_utils_trim.py` and update boundary test imports. Co-authored-by: Cursor --- .../prefix_sharing/backends/packed_layout.py | 2 +- .../prefix_sharing/integrations/verl_fsdp.py | 74 +++---- .../prefix_sharing/integrations/verl_mcore.py | 47 ++--- .../prefix_sharing/integrations/verl_utils.py | 198 ++++++++++++------ .../patches/verl080_fsdp/forward_step.py | 4 +- .../verl080_mcore0161_ms0160/forward_step.py | 12 +- .../test_verl_integration_boundaries.py | 8 +- .../tests/unit_test/test_verl_utils_trim.py | 143 +++++++++++++ 8 files changed, 341 insertions(+), 147 deletions(-) create mode 100644 prefix-sharing/tests/unit_test/test_verl_utils_trim.py diff --git a/prefix-sharing/prefix_sharing/backends/packed_layout.py b/prefix-sharing/prefix_sharing/backends/packed_layout.py index d30809cc..8d14bc3c 100644 --- a/prefix-sharing/prefix_sharing/backends/packed_layout.py +++ b/prefix-sharing/prefix_sharing/backends/packed_layout.py @@ -57,7 +57,7 @@ def from_kept_position_rows( cls, kept_position_rows: Sequence[Any], *, - align_size: int, + align_size: int = 1, ) -> "PackedBatchLayout": if align_size < 1: raise ValueError("align_size must be >= 1") diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index 758a9aec..4e3f0da8 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -22,12 +22,12 @@ from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.parallel_info import MegatronParallelInfo from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState -from prefix_sharing.integrations.verl_utils import _clone_batch -from prefix_sharing.integrations.verl_utils import _collect_kept_position_rows -from prefix_sharing.integrations.verl_utils import _extract_seq_from_nested_tensor -from prefix_sharing.integrations.verl_utils import _extract_sequences_async -from prefix_sharing.integrations.verl_utils import _is_nested_tensor -from prefix_sharing.integrations.verl_utils import _trim_nested_batch +from prefix_sharing.integrations.verl_utils import clone_batch +from prefix_sharing.integrations.verl_utils import extract_seq_from_dense_tensor +from prefix_sharing.integrations.verl_utils import extract_seq_from_nested_tensor +from prefix_sharing.integrations.verl_utils import is_nested_tensor +from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_nested_tensor +from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_dense_tensor class PrefixSharingFSDPAttentionRuntime: """Standalone FSDP attention runtime for PrefixSharing. @@ -190,69 +190,49 @@ def plan_and_trim_microbatch_fsdp( prefix sharing is disabled or no reusable prefix is detected. """ + # skip if PrefixSharing is disabled if not ps_config.enable_prefix_sharing: return micro_batch, None + + # validate the config ps_config.validate(model_config=model_config, integrate_mode="verl_fsdp") + # extract input_ids from micro-batch and convert it into sequences (Python List) + # for subsequent PrefixSharing planning, which is purely done on CPU input_ids = micro_batch["input_ids"] - is_nested_input = _is_nested_tensor(input_ids) - if is_nested_input: - sequences = _extract_seq_from_nested_tensor(input_ids) + input_ids_is_nested = is_nested_tensor(input_ids) + if input_ids_is_nested: # nested tensor + sequences = extract_seq_from_nested_tensor(input_ids) valid_indices = None attention_mask = None - else: + else: # dense tensor attention_mask = micro_batch["attention_mask"].to(bool) if input_ids.dim() != 2 or attention_mask.dim() != 2: - raise RuntimeError("prefix sharing FSDP path expects 2D or jagged NestedTensor input_ids") + raise RuntimeError("PrefixSharing + FSDP expects 2D or jagged NestedTensor input_ids") if input_ids.shape != attention_mask.shape: raise RuntimeError("input_ids and attention_mask must have the same shape") - valid_indices = [ attention_mask[row].nonzero(as_tuple=False).flatten() for row in range(input_ids.shape[0]) ] - sequences = _extract_sequences_async(input_ids, valid_indices) + sequences = extract_seq_from_dense_tensor(input_ids, valid_indices) + + # plan for PrefixSharing prefix_sharing_plan = PrefixSharingPlanner(ps_config).plan(sequences) if not prefix_sharing_plan.has_sharing: return micro_batch, None - if is_nested_input: - trimmed_micro_batch = _trim_nested_batch(micro_batch, prefix_sharing_plan) - kept_position_rows = _collect_kept_position_rows( - trimmed_micro_batch, - prefix_sharing_plan, - is_nested_tensor=True, + # trim the micro-batch + if input_ids_is_nested: # nested tensor + trimmed_micro_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor( + micro_batch, prefix_sharing_plan ) - else: - trimmed_micro_batch = _clone_batch(micro_batch) - trimmed_attention_mask = attention_mask.clone() - trimmed_attention_mask[:] = False - - for row, indices in enumerate(valid_indices): - keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[row] - kept_indices = indices[keep_start:keep_end] - trimmed_attention_mask[row, kept_indices] = True - - trimmed_micro_batch["attention_mask"] = trimmed_attention_mask - - if "loss_mask" in trimmed_micro_batch: - trimmed_loss_mask = trimmed_micro_batch["loss_mask"].to(bool).clone() - trimmed_loss_mask[:] = False - for row, indices in enumerate(valid_indices): - keep_start, keep_end = prefix_sharing_plan.loss_mask_keep_ranges[row] - trimmed_loss_mask[row, indices[keep_start:keep_end]] = micro_batch["loss_mask"][row, indices[keep_start:keep_end]].to(bool) - trimmed_micro_batch["loss_mask"] = trimmed_loss_mask - kept_position_rows = _collect_kept_position_rows( - trimmed_micro_batch, - prefix_sharing_plan, - is_nested_tensor=False, - valid_indices=valid_indices, + else: # dense tensor + trimmed_micro_batch, kept_position_rows = trim_redundant_prefix_in_dense_tensor( + micro_batch, prefix_sharing_plan, valid_indices ) - packed_batch_layout = PackedBatchLayout.from_kept_position_rows( - kept_position_rows, - align_size=1, - ) + packed_batch_layout = PackedBatchLayout.from_kept_position_rows(kept_position_rows) runtime_state = PrefixSharingRuntimeState( prefix_sharing_plan=prefix_sharing_plan, attention_backend=get_backend_instance(ps_config, backend), diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index becfe915..70da1d5d 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -23,13 +23,12 @@ from prefix_sharing.integrations.context import current_prefix_sharing_context from prefix_sharing.integrations.parallel_info import get_megatron_parallel_info from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState -from prefix_sharing.integrations.verl_utils import _clone_batch -from prefix_sharing.integrations.verl_utils import _collect_kept_position_rows -from prefix_sharing.integrations.verl_utils import _extract_seq_from_nested_tensor -from prefix_sharing.integrations.verl_utils import _extract_sequences_async -from prefix_sharing.integrations.verl_utils import _is_nested_tensor -from prefix_sharing.integrations.verl_utils import _trim_nested_batch -from prefix_sharing.integrations.verl_utils import _trim_plain_batch_thd +from prefix_sharing.integrations.verl_utils import collect_kept_position_rows +from prefix_sharing.integrations.verl_utils import extract_seq_from_dense_tensor +from prefix_sharing.integrations.verl_utils import extract_seq_from_nested_tensor +from prefix_sharing.integrations.verl_utils import is_nested_tensor +from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_nested_tensor +from prefix_sharing.integrations.verl_utils import trim_plain_batch_thd from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config @@ -199,10 +198,10 @@ def restore_via_2d_unfold_verl080( return output log_probs_nested = output.get("log_probs") - if log_probs_nested is None or not _is_nested_tensor(log_probs_nested): + if log_probs_nested is None or not is_nested_tensor(log_probs_nested): return output entropy_nested = output.get("entropy") - has_entropy = entropy_nested is not None and _is_nested_tensor(entropy_nested) + has_entropy = entropy_nested is not None and is_nested_tensor(entropy_nested) original_lengths = plan.original_lengths input_keep_ranges = plan.input_keep_ranges @@ -368,15 +367,15 @@ def build_prefix_sharing_micro_batch_verl080( raise RuntimeError("prefix sharing phase 1 does not support dynamic context parallel") # ── 阶段 3: 从 batch 提取序列 ── - # NestedTensor → 从 offsets/values 提取 - # Plain 2D → 从 attention_mask.nonzero() 提取 - # 同时保留 attention_mask_bool，供阶段 6 的 _collect_kept_position_rows 使用。 + # NestedTensor → 从 offsets/values 提取；kept_position_rows 由 trim 直接返回。 + # Plain 2D → 从 attention_mask.nonzero() 提取，并保留 valid_indices + # 供 trim 后的 collect_kept_position_rows 使用。 input_ids = batch["input_ids"] - is_nested_tensor = _is_nested_tensor(input_ids) + is_nested_input = is_nested_tensor(input_ids) attention_mask_bool_for_layout = None - if is_nested_tensor: - sequences = _extract_seq_from_nested_tensor(input_ids) + if is_nested_input: + sequences = extract_seq_from_nested_tensor(input_ids) else: # plain 2D tensor（需要 attention_mask） attention_mask = batch.get("attention_mask") @@ -388,7 +387,7 @@ def build_prefix_sharing_micro_batch_verl080( attention_mask_bool[row].nonzero(as_tuple=False).flatten() for row in range(input_ids.shape[0]) ] - sequences = _extract_sequences_async(input_ids, valid_indices) + sequences = extract_seq_from_dense_tensor(input_ids, valid_indices) # ── 阶段 4: 前缀共享规划 ── plan = PrefixSharingPlanner(ps_config).plan(sequences) @@ -401,16 +400,14 @@ def build_prefix_sharing_micro_batch_verl080( # 2D path: 只改 attention_mask（Megatron 从 mask 动态重算 packed）， # v080 THD 路径用 preprocess_thd_engine(input_ids) 直接处理数据， # 不看 attention_mask。必须物理裁剪 input_ids/position_ids。 - if is_nested_tensor: - trimmed_batch = _trim_nested_batch(batch, plan) + if is_nested_input: + trimmed_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor(batch, plan) else: - trimmed_batch = _trim_plain_batch_thd(batch, plan, valid_indices) - - # layout 计算：从 trimmed 后的实际 kept position rows 构建 - kept_position_rows = _collect_kept_position_rows( - trimmed_batch, plan, is_nested_tensor, - valid_indices=valid_indices if not is_nested_tensor else None, - ) + trimmed_batch = trim_plain_batch_thd(batch, plan, valid_indices) + kept_position_rows = collect_kept_position_rows( + trimmed_batch, plan, is_nested_input, + valid_indices=valid_indices, + ) # ── 阶段 6: 构建 layout ── parallel_info = get_megatron_parallel_info() diff --git a/prefix-sharing/prefix_sharing/integrations/verl_utils.py b/prefix-sharing/prefix_sharing/integrations/verl_utils.py index 5cae37b6..8f3bc65a 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_utils.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_utils.py @@ -4,6 +4,8 @@ from typing import Any +import torch + from prefix_sharing.core.planner import PrefixSharingPlan @@ -64,7 +66,7 @@ def read_ps_config_from_prefix_grouper(engine_config: Any) -> dict[str, Any] | N return values -def _clone_batch(batch: Any) -> Any: +def clone_batch(batch: Any) -> Any: if hasattr(batch, "clone"): try: return batch.clone() @@ -98,23 +100,32 @@ def _read_actor_value(config: Any, dotted_name: str, default: Any) -> Any: return current -def _trim_nested_batch(batch: Any, plan: PrefixSharingPlan) -> Any: - """Physically trim a NestedTensor batch for verl 0.8 THD paths.""" +def trim_redundant_prefix_in_nested_tensor( + batch: Any, prefix_sharing_plan: PrefixSharingPlan +) -> tuple[Any, list[Any]]: + """Trim redundant prefix tokens in NestedTensor inputs so that their computation + are skipped and not performed. - import torch + Args: + batch: Input batch with NestedTensor ``input_ids``. + prefix_sharing_plan: PrefixSharingPlan with keep ranges. - trimmed_batch = _clone_batch(batch) + Returns: + ``(trimmed_batch, kept_position_rows)``. ``kept_position_rows`` is the + per-row position-id slice used to build the trimmed NestedTensor, so + callers do not need to unpack it again. + """ + trimmed_batch = clone_batch(batch) input_ids = batch["input_ids"] position_ids = batch["position_ids"] + + # trim input_ids + trimmed_batch["input_ids"] = _trim_nested_tensor(input_ids, prefix_sharing_plan) - trimmed_ids_seqs = _slice_nested_sequences(input_ids, plan) - new_input_ids = torch.nested.as_nested_tensor(trimmed_ids_seqs, layout=torch.jagged) - trimmed_batch["input_ids"] = new_input_ids - - if _is_nested_tensor(position_ids): - trimmed_pos_seqs = _slice_nested_sequences(position_ids, plan) - new_position_ids = torch.nested.as_nested_tensor(trimmed_pos_seqs, layout=torch.jagged) + # trim position_ids + if is_nested_tensor(position_ids): + kept_position_rows = _trim_nested_rows(position_ids, prefix_sharing_plan) else: attention_mask = batch.get("attention_mask") if attention_mask is not None: @@ -124,23 +135,74 @@ def _trim_nested_batch(batch: Any, plan: PrefixSharingPlan) -> Any: position_ids.shape[0], position_ids.shape[1], dtype=torch.bool, device=position_ids.device, ) - trimmed_pos_seqs = _slice_2d_position_rows( - position_ids, plan, attention_mask_bool, + kept_position_rows = _trim_2d_tensor( + position_ids, prefix_sharing_plan, attention_mask_bool ) - new_position_ids = torch.nested.as_nested_tensor(trimmed_pos_seqs, layout=torch.jagged) - trimmed_batch["position_ids"] = new_position_ids + trimmed_batch["position_ids"] = torch.nested.as_nested_tensor( + kept_position_rows, layout=torch.jagged + ) + # trim loss_mask loss_mask = batch.get("loss_mask") - if loss_mask is not None and _is_nested_tensor(loss_mask): - trimmed_loss_seqs = _slice_nested_sequences(loss_mask, plan) - trimmed_batch["loss_mask"] = torch.nested.as_nested_tensor( - trimmed_loss_seqs, layout=torch.jagged - ) + if loss_mask is not None and is_nested_tensor(loss_mask): + trimmed_batch["loss_mask"] = _trim_nested_tensor(loss_mask, prefix_sharing_plan) - return trimmed_batch + return trimmed_batch, kept_position_rows + + +def trim_redundant_prefix_in_dense_tensor( + batch: Any, + prefix_sharing_plan: PrefixSharingPlan, + valid_indices: list[Any], +) -> tuple[Any, list[Any]]: + """Trim redundant prefix tokens from a dense 2D batch by masking. + + Unlike ``trim_plain_batch_thd`` which physically removes kept tokens, + this helper keeps ``input_ids`` and ``position_ids`` dense and only + masks the kept positions in ``attention_mask`` and ``loss_mask``. + The position-id rows are trimmed to jagged tensors for packed layout + construction, but the original 2D ``position_ids`` tensor is left + unchanged so downstream dense-path code can still index it. + + Args: + batch: Input batch with 2D ``input_ids``, ``position_ids`` and + ``attention_mask``. + prefix_sharing_plan: Prefix sharing plan with keep ranges. + valid_indices: Pre-computed nonzero indices from ``attention_mask``. + + Returns: + ``(trimmed_batch, kept_position_rows)``. + """ + + trimmed_batch = clone_batch(batch) + attention_mask = batch["attention_mask"].to(bool) + trimmed_attention_mask = attention_mask.clone() + trimmed_attention_mask[:] = False + + for row, indices in enumerate(valid_indices): + keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[row] + kept_indices = indices[keep_start:keep_end] + trimmed_attention_mask[row, kept_indices] = True + trimmed_batch["attention_mask"] = trimmed_attention_mask + + if "loss_mask" in trimmed_batch: + trimmed_loss_mask = trimmed_batch["loss_mask"].to(bool).clone() + trimmed_loss_mask[:] = False + for row, indices in enumerate(valid_indices): + keep_start, keep_end = prefix_sharing_plan.loss_mask_keep_ranges[row] + trimmed_loss_mask[row, indices[keep_start:keep_end]] = ( + batch["loss_mask"][row, indices[keep_start:keep_end]].to(bool) + ) + trimmed_batch["loss_mask"] = trimmed_loss_mask + + kept_position_rows = _trim_2d_tensor( + batch["position_ids"], prefix_sharing_plan, attention_mask + ) + + return trimmed_batch, kept_position_rows -def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: list[Any] | None = None) -> Any: +def trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: list[Any] | None = None) -> Any: """Physically trim a plain 2D tensor batch for verl 0.8 THD paths. Args: @@ -150,8 +212,6 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: li If None, will compute from batch["attention_mask"]. """ - import torch - input_ids = batch["input_ids"] position_ids = batch["position_ids"] @@ -179,7 +239,7 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: li kept_id_rows.append(input_ids[row, kept_indices]) kept_pos_rows.append(position_ids[row, kept_indices]) - trimmed_batch = _clone_batch(batch) + trimmed_batch = clone_batch(batch) trimmed_batch["input_ids"] = torch.nested.as_nested_tensor(kept_id_rows, layout=torch.jagged) trimmed_batch["position_ids"] = torch.nested.as_nested_tensor(kept_pos_rows, layout=torch.jagged) @@ -198,37 +258,55 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: li return trimmed_batch -def _slice_nested_sequences(nested_tensor: Any, plan: PrefixSharingPlan) -> list[Any]: +def _trim_nested_rows(nested_tensor: Any, prefix_sharing_plan: PrefixSharingPlan) -> list[Any]: + """Trim each sequence of a NestedTensor to its keep range in the plan. + + Returns a list of 1D tensors (one per sequence), not packed into NestedTensor. + """ + offsets = nested_tensor.offsets() values = nested_tensor.values() sliced = [] - for i in range(len(plan.input_keep_ranges)): + for i in range(len(prefix_sharing_plan.input_keep_ranges)): seq_values = values[offsets[i]:offsets[i + 1]] - keep_start, keep_end = plan.input_keep_ranges[i] + keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[i] sliced.append(seq_values[keep_start:keep_end]) return sliced -def _slice_2d_position_rows( - position_ids: Any, - plan: PrefixSharingPlan, +def _trim_nested_tensor(nested_tensor: Any, prefix_sharing_plan: PrefixSharingPlan) -> Any: + """Trim each sequence of a NestedTensor and pack into a jagged NestedTensor.""" + + return torch.nested.as_nested_tensor( + _trim_nested_rows(nested_tensor, prefix_sharing_plan), layout=torch.jagged + ) + + +def _trim_2d_tensor( + tensor_2d: Any, + prefix_sharing_plan: PrefixSharingPlan, attention_mask_bool: Any, ) -> list[Any]: + """Trim each row of a 2D dense tensor to its keep range in the plan. + + Returns a list of 1D tensors (one per row), not packed into NestedTensor. + """ + kept_rows = [] - for row in range(position_ids.shape[0]): + for row in range(tensor_2d.shape[0]): indices = attention_mask_bool[row].nonzero(as_tuple=False).flatten() - keep_start, keep_end = plan.input_keep_ranges[row] + keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[row] kept_indices = indices[keep_start:keep_end] - kept_rows.append(position_ids[row, kept_indices]) + kept_rows.append(tensor_2d[row, kept_indices]) return kept_rows -def _collect_kept_position_rows( +def collect_kept_position_rows( trimmed_batch: Any, - plan: PrefixSharingPlan, - is_nested_tensor: bool, + prefix_sharing_plan: PrefixSharingPlan, + is_nested_input: bool, attention_mask_bool: Any | None = None, valid_indices: list[Any] | None = None, ) -> list[Any]: @@ -237,7 +315,7 @@ def _collect_kept_position_rows( Args: trimmed_batch: Batch after trimming. plan: Prefix sharing plan. - is_nested_tensor: Whether position_ids is a NestedTensor. + is_nested_input: Whether position_ids is a NestedTensor. attention_mask_bool: Boolean attention mask (deprecated, use valid_indices). valid_indices: Pre-computed nonzero indices. If provided, skips attention_mask_bool computation. @@ -245,10 +323,10 @@ def _collect_kept_position_rows( position_ids = trimmed_batch["position_ids"] - if is_nested_tensor or _is_nested_tensor(position_ids): + if is_nested_input or is_nested_tensor(position_ids): offsets = position_ids.offsets() values = position_ids.values() - return [values[offsets[i]:offsets[i + 1]] for i in range(len(plan.input_keep_ranges))] + return [values[offsets[i]:offsets[i + 1]] for i in range(len(prefix_sharing_plan.input_keep_ranges))] # 2D tensor — need valid_indices to locate valid column indices if valid_indices is None: @@ -259,19 +337,19 @@ def _collect_kept_position_rows( ) valid_indices = [ attention_mask_bool[i].nonzero(as_tuple=False).flatten() - for i in range(len(plan.input_keep_ranges)) + for i in range(len(prefix_sharing_plan.input_keep_ranges)) ] rows = [] - for i in range(len(plan.input_keep_ranges)): + for i in range(len(prefix_sharing_plan.input_keep_ranges)): indices = valid_indices[i] - keep_start, keep_end = plan.input_keep_ranges[i] + keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[i] kept_indices = indices[keep_start:keep_end] rows.append(position_ids[i, kept_indices]) return rows -def _is_nested_tensor(tensor: Any) -> bool: +def is_nested_tensor(tensor: Any) -> bool: return ( hasattr(tensor, "offsets") and callable(tensor.offsets) @@ -283,7 +361,7 @@ def _is_nested_tensor(tensor: Any) -> bool: def _copy_tensors_to_cpu_lists(tensors: list[Any]) -> list[list[int]]: """Copy device tensors to Python int lists with a single sync point. - Fast path (CUDA / NPU): stage all slices into pinned host buffers with + Fast path (GPU / NPU): stage all slices into pinned host buffers with non-blocking copies, then synchronize the source device once — avoiding one pipeline stall per tensor. Backends without pinned-memory support (older torch_npu, CPU, …) fall back to plain per-tensor copies: correct @@ -293,8 +371,6 @@ def _copy_tensors_to_cpu_lists(tensors: list[Any]) -> list[list[int]]: works on any accelerator backend without hardcoding ``torch.cuda``. """ - import torch - if not tensors: return [] @@ -322,31 +398,29 @@ def _copy_tensors_to_cpu_lists(tensors: list[Any]) -> list[list[int]]: return [buf.tolist() for buf in pinned_buffers] -def _extract_seq_from_nested_tensor(nested_tensor: Any) -> list[list[int]]: - """从 NestedTensor (jagged layout) 中提取每个序列的 token ID 列表。 - - 批量异步 device→CPU 拷贝（CUDA/NPU），单次同步，避免 pipeline stall。 - """ +def extract_seq_from_nested_tensor(nested_tensor: Any) -> list[list[int]]: + """Extract token sequences from a NestedTensor (jagged layout).""" offsets = nested_tensor.offsets() values = nested_tensor.values() - slices = [ + tensor_slices = [ values[offsets[i]:offsets[i + 1]].detach() for i in range(offsets.numel() - 1) ] - return _copy_tensors_to_cpu_lists(slices) + return _copy_tensors_to_cpu_lists(tensor_slices) -def _extract_sequences_async( +def extract_seq_from_dense_tensor( input_ids: Any, valid_indices: list[Any], ) -> list[list[int]]: - """Extract token sequences from a device tensor with batched async copy. + """Extract token sequences from a dense 2D device tensor with batched async copy. - Instead of calling .cpu().tolist() per row (which syncs each time), - all device→CPU copies are queued non-blocking, then synchronized once - (CUDA/NPU fast path; other backends fall back to plain copies). + Uses a batched async copy: instead of calling .cpu().tolist() per row + (which syncs each time), all device→CPU copies are queued non-blocking, + then synchronized once (CUDA/NPU fast path; other backends fall back to + plain copies). Args: input_ids: 2D device tensor of shape [batch_size, seq_len]. @@ -356,8 +430,8 @@ def _extract_sequences_async( List of token ID lists (CPU Python ints). """ - tensors = [ + tensor_slices = [ input_ids[row, indices].detach() for row, indices in enumerate(valid_indices) ] - return _copy_tensors_to_cpu_lists(tensors) + return _copy_tensors_to_cpu_lists(tensor_slices) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 5077506b..4125e842 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -464,7 +464,7 @@ def _restore_engine_model_output(model_output: dict[str, Any]) -> dict[str, Any] entropy_from_logits = None from prefix_sharing.integrations.verl_mcore import restore_via_2d_unfold_verl080 - from prefix_sharing.integrations.verl_mcore import _is_nested_tensor + from prefix_sharing.integrations.verl_utils import is_nested_tensor from prefix_sharing.integrations.verl_fsdp import restore_prefix_sharing_outputs_2d restored = restore_via_2d_unfold_verl080( @@ -473,7 +473,7 @@ def _restore_engine_model_output(model_output: dict[str, Any]) -> dict[str, Any] entropy_from_logits, ) log_probs = restored.get("log_probs") - if log_probs is not None and not _is_nested_tensor(log_probs): + if log_probs is not None and not is_nested_tensor(log_probs): return restore_prefix_sharing_outputs_2d(restored, logprobs_from_logits) return restored diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py index 2d4da6aa..51b2e476 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py @@ -130,9 +130,9 @@ def patched_forward_step( build_attention_mask_2d, build_label_mask_2d, nested_offsets_to_cu, ) - from prefix_sharing.integrations.verl_mcore import _is_nested_tensor + from prefix_sharing.integrations.verl_utils import is_nested_tensor _ids_nested = batch_for_forward["input_ids"] - if _is_nested_tensor(_ids_nested): + if is_nested_tensor(_ids_nested): if ps_state is not None: _plan = ps_state.prefix_sharing_plan _prefix_lens = list(_plan.prefix_lens) @@ -160,7 +160,7 @@ def patched_forward_step( # loss_mask 行 sum，与坐标系无关，据此推 prompt_len 最稳（2D/NestedTensor 均适用）。 _lm = original_batch.get("loss_mask") if _lm is not None: - if _is_nested_tensor(_lm): + if is_nested_tensor(_lm): _lm_off = _lm.offsets() _lm_val = _lm.values() _response_lens = [ @@ -228,7 +228,7 @@ def patched_forward_step( from prefix_sharing.tools.diagnostic_dump import ( nested_to_2d_full, dump_logprobs_2d_verl080, dump_entropy_2d_verl080, ) - from prefix_sharing.integrations.verl_mcore import _is_nested_tensor + from prefix_sharing.integrations.verl_utils import is_nested_tensor # 对齐 v070: tag = "old" if forward_only else "train"。 # forward_step 拿不到 forward_only，用 model.training 等价区分 # （eval_mode→training=False→"old" 对应 old_logp 阶段； @@ -237,7 +237,7 @@ def patched_forward_step( _tag = "train" if model.training else "old" _out_dict, _ = output _lp = _out_dict.get("log_probs") - if _is_nested_tensor(_lp): + if is_nested_tensor(_lp): if ps_state is not None: _ol = list(ps_state.prefix_sharing_plan.original_lengths) else: @@ -245,7 +245,7 @@ def patched_forward_step( _Lmax = max(_ol) if _ol else 0 dump_logprobs_2d_verl080(nested_to_2d_full(_lp, _ol, _Lmax), _tag) _ent = _out_dict.get("entropy") - if _is_nested_tensor(_ent): + if is_nested_tensor(_ent): dump_entropy_2d_verl080(nested_to_2d_full(_ent, _ol, _Lmax), _tag) # ##### [PS-diag] dump 2D logprobs/entropy end ##### _ps_forward_step_probe("after_original_forward_step") diff --git a/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py b/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py index f9ddcba3..5a00a292 100644 --- a/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py +++ b/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py @@ -22,10 +22,10 @@ def test_verl_fsdp_does_not_import_shared_helpers_from_mcore(): source = inspect.getsource(verl_fsdp) forbidden_imports = ( "from prefix_sharing.integrations.verl_mcore import PrefixSharingRuntimeState", - "from prefix_sharing.integrations.verl_mcore import _collect_kept_position_rows", - "from prefix_sharing.integrations.verl_mcore import _extract_seq_from_nested_tensor", - "from prefix_sharing.integrations.verl_mcore import _is_nested_tensor", - "from prefix_sharing.integrations.verl_mcore import _trim_nested_batch", + "from prefix_sharing.integrations.verl_mcore import collect_kept_position_rows", + "from prefix_sharing.integrations.verl_mcore import extract_seq_from_nested_tensor", + "from prefix_sharing.integrations.verl_mcore import is_nested_tensor", + "from prefix_sharing.integrations.verl_mcore import trim_redundant_prefix_in_nested_tensor", ) for forbidden_import in forbidden_imports: assert forbidden_import not in source diff --git a/prefix-sharing/tests/unit_test/test_verl_utils_trim.py b/prefix-sharing/tests/unit_test/test_verl_utils_trim.py new file mode 100644 index 00000000..875b23c8 --- /dev/null +++ b/prefix-sharing/tests/unit_test/test_verl_utils_trim.py @@ -0,0 +1,143 @@ +"""Tests for nested-tensor trim helpers in verl_utils.""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +from prefix_sharing.core.config import PrefixSharingConfig +from prefix_sharing.core.planner import PrefixSharingPlanner +from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_dense_tensor +from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_nested_tensor + + +def _prefix_sharing_plan(): + planner = PrefixSharingPlanner(PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3)) + return planner.plan([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21, 22]]) + + +def _nested(rows): + return torch.nested.as_nested_tensor( + [torch.tensor(row, dtype=torch.long) for row in rows], + layout=torch.jagged, + ) + + +def _unpacked_rows(nested_tensor): + offsets = nested_tensor.offsets() + values = nested_tensor.values() + return [ + values[offsets[i] : offsets[i + 1]].tolist() + for i in range(offsets.numel() - 1) + ] + + +def test_trim_redundant_prefix_returns_kept_position_rows_from_nested_slice(): + if not hasattr(torch, "nested"): + pytest.skip("torch.nested is unavailable") + + plan = _prefix_sharing_plan() + batch = { + "input_ids": _nested([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21, 22]]), + "position_ids": _nested([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4, 5]]), + "loss_mask": _nested([[1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 0]]), + } + + trimmed_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor(batch, plan) + + assert [row.tolist() for row in kept_position_rows] == [[0, 1, 2, 3, 4], [3, 4, 5]] + assert _unpacked_rows(trimmed_batch["position_ids"]) == [ + row.tolist() for row in kept_position_rows + ] + assert _unpacked_rows(trimmed_batch["input_ids"]) == [[1, 2, 3, 10, 11], [20, 21, 22]] + assert _unpacked_rows(trimmed_batch["loss_mask"]) == [[1, 1, 1, 1, 0], [1, 1, 0]] + + +def test_trim_redundant_prefix_returns_kept_position_rows_from_2d_slice(): + if not hasattr(torch, "nested"): + pytest.skip("torch.nested is unavailable") + + plan = _prefix_sharing_plan() + batch = { + "input_ids": _nested([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21, 22]]), + "position_ids": torch.tensor( + [ + [0, 1, 2, 3, 4, 0], + [0, 1, 2, 3, 4, 5], + ], + dtype=torch.long, + ), + "attention_mask": torch.tensor( + [ + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 1], + ], + dtype=torch.bool, + ), + } + + trimmed_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor(batch, plan) + + assert [row.tolist() for row in kept_position_rows] == [[0, 1, 2, 3, 4], [3, 4, 5]] + assert _unpacked_rows(trimmed_batch["position_ids"]) == [ + row.tolist() for row in kept_position_rows + ] + + +def test_trim_redundant_prefix_in_dense_tensor_crops_masks_and_returns_position_rows(): + if not hasattr(torch, "nested"): + pytest.skip("torch.nested is unavailable") + + plan = _prefix_sharing_plan() + batch = { + "input_ids": torch.tensor( + [ + [1, 2, 3, 10, 11, 0], + [1, 2, 3, 20, 21, 22], + ], + dtype=torch.long, + ), + "position_ids": torch.tensor( + [ + [0, 1, 2, 3, 4, 0], + [0, 1, 2, 3, 4, 5], + ], + dtype=torch.long, + ), + "attention_mask": torch.tensor( + [ + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 1], + ], + dtype=torch.bool, + ), + "loss_mask": torch.tensor( + [ + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 1], + ], + dtype=torch.bool, + ), + } + valid_indices = [ + batch["attention_mask"][row].nonzero(as_tuple=False).flatten() + for row in range(batch["input_ids"].shape[0]) + ] + + trimmed_batch, kept_position_rows = trim_redundant_prefix_in_dense_tensor( + batch, plan, valid_indices + ) + + assert [row.tolist() for row in kept_position_rows] == [[0, 1, 2, 3, 4], [3, 4, 5]] + assert trimmed_batch["position_ids"] is batch["position_ids"] + assert trimmed_batch["position_ids"][1, 3:6].tolist() == [3, 4, 5] + assert trimmed_batch["attention_mask"].tolist() == [ + [1, 1, 1, 1, 1, 0], + [0, 0, 0, 1, 1, 1], + ] + assert trimmed_batch["loss_mask"].tolist() == [ + [1, 1, 1, 1, 1, 0], + [0, 0, 0, 1, 1, 1], + ] + assert trimmed_batch["input_ids"] is batch["input_ids"] From cd550c231b0c5e96952cff0c6c4782a7afab7fd8 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:26:36 +0800 Subject: [PATCH 06/19] refactor(packed_layout): use local device variable in from_kept_position_rows Replace repeated `first.device` lookups with a single `device` local variable to improve readability. Co-authored-by: Cursor --- prefix-sharing/prefix_sharing/backends/packed_layout.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/prefix-sharing/prefix_sharing/backends/packed_layout.py b/prefix-sharing/prefix_sharing/backends/packed_layout.py index 8d14bc3c..5eebf900 100644 --- a/prefix-sharing/prefix_sharing/backends/packed_layout.py +++ b/prefix-sharing/prefix_sharing/backends/packed_layout.py @@ -65,18 +65,19 @@ def from_kept_position_rows( return cls.from_valid_lengths([]) first = kept_position_rows[0] + device = first.device valid_lengths = [int(row.shape[0]) for row in kept_position_rows] padded_lengths = [_pad_to_multiple(length, align_size) for length in valid_lengths] packed_position_rows = [] valid_mask_rows = [] for row, valid_length, padded_length in zip(kept_position_rows, valid_lengths, padded_lengths): - row = row.to(first.device) + row = row.to(device) pad_length = padded_length - valid_length valid_mask_rows.append( torch.cat( [ - torch.ones(valid_length, dtype=torch.bool, device=first.device), - torch.zeros(pad_length, dtype=torch.bool, device=first.device), + torch.ones(valid_length, dtype=torch.bool, device=device), + torch.zeros(pad_length, dtype=torch.bool, device=device), ], dim=0, ) @@ -84,7 +85,7 @@ def from_kept_position_rows( if pad_length == 0: packed_position_rows.append(row) continue - padding = torch.zeros(pad_length, dtype=row.dtype, device=first.device) + padding = torch.zeros(pad_length, dtype=row.dtype, device=device) packed_position_rows.append(torch.cat([row, padding], dim=0)) return cls( From 9dc2f8c15be43184556db411e068b8b99a0285b3 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:34:15 +0800 Subject: [PATCH 07/19] refactor(verl_fsdp): rename plan_and_trim_microbatch_fsdp to prepare_for_prefix_sharing_fsdp - Rename `plan_and_trim_microbatch_fsdp` to `prepare_for_prefix_sharing_fsdp` to better reflect its responsibility (plan, trim, build layout, and construct runtime state). - Update all callers: verl_fsdp patch, integrations __init__, and FSDP unit tests. - Add STEP comments to clarify the prepare/execute phases in `verl_fsdp.py`. Co-authored-by: Cursor --- .../prefix_sharing/integrations/__init__.py | 4 ++-- .../prefix_sharing/integrations/verl_fsdp.py | 13 +++++++++--- .../patches/verl080_fsdp/forward_step.py | 4 ++-- .../tests/unit_test/test_verl_fsdp_adapter.py | 20 +++++++++---------- .../test_verl_fsdp_ch4_functional.py | 18 ++++++++--------- .../unit_test/test_verl_fsdp_ch4_precision.py | 16 +++++++-------- 6 files changed, 41 insertions(+), 34 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/__init__.py b/prefix-sharing/prefix_sharing/integrations/__init__.py index c00326da..e0fbdb27 100644 --- a/prefix-sharing/prefix_sharing/integrations/__init__.py +++ b/prefix-sharing/prefix_sharing/integrations/__init__.py @@ -16,7 +16,7 @@ ) from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - plan_and_trim_microbatch_fsdp, + prepare_for_prefix_sharing_fsdp, forward_prefix_sharing_fsdp_micro_batch, restore_prefix_sharing_outputs_2d, ) @@ -38,7 +38,7 @@ "prefix_attention", "get_megatron_parallel_info", "PrefixSharingFSDPAttentionRuntime", - "plan_and_trim_microbatch_fsdp", + "prepare_for_prefix_sharing_fsdp", "forward_prefix_sharing_fsdp_micro_batch", "restore_prefix_sharing_outputs_2d", ] diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index 4e3f0da8..c822860a 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -3,7 +3,7 @@ verl FSDP integration helpers for PrefixSharing. The FSDP path follows the same public shape as the Megatron integration: -``plan_and_trim_microbatch_fsdp`` returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``. +``prepare_for_prefix_sharing_fsdp`` returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``. The helpers stay framework-light enough for CPU tests, while the explicit ``verl080_fsdp`` patch set wires them into ``FSDPEngineWithLMHead.forward_step``. """ @@ -136,7 +136,10 @@ def forward_prefix_sharing_fsdp_micro_batch( own ``prepare_model_inputs`` / ``prepare_model_outputs`` path. """ - trimmed_micro_batch, runtime_state = plan_and_trim_microbatch_fsdp( + ######################################################### + # STEP 1: prepare for prefix sharing + ######################################################### + trimmed_micro_batch, runtime_state = prepare_for_prefix_sharing_fsdp( micro_batch, ps_config, model_config=model_config, @@ -146,6 +149,9 @@ def forward_prefix_sharing_fsdp_micro_batch( autocast = autocast_context if autocast_context is not None else nullcontext() with context as ctx, autocast: + ######################################################### + # STEP 2: call the model with PrefixSharing r + ######################################################### model_output = _call_fsdp_model( model, trimmed_micro_batch, @@ -175,7 +181,7 @@ def forward_prefix_sharing_fsdp_micro_batch( return output -def plan_and_trim_microbatch_fsdp( +def prepare_for_prefix_sharing_fsdp( micro_batch: Any, ps_config: PrefixSharingConfig, model_config: Any | None = None, @@ -232,6 +238,7 @@ def plan_and_trim_microbatch_fsdp( micro_batch, prefix_sharing_plan, valid_indices ) + # build layout and runtime state for the trimmed micro-batch packed_batch_layout = PackedBatchLayout.from_kept_position_rows(kept_position_rows) runtime_state = PrefixSharingRuntimeState( prefix_sharing_plan=prefix_sharing_plan, diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 4125e842..d984df0d 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -161,7 +161,7 @@ def _forward_step_with_engine_prepare( from prefix_sharing.integrations.context import create_prefix_sharing_context from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime - from prefix_sharing.integrations.verl_fsdp import plan_and_trim_microbatch_fsdp + from prefix_sharing.integrations.verl_fsdp import prepare_for_prefix_sharing_fsdp from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope @@ -175,7 +175,7 @@ def _forward_step_with_engine_prepare( if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: _dump_full_input_ids_only(micro_batch, "train") - trimmed_micro_batch, ps_state = plan_and_trim_microbatch_fsdp( + trimmed_micro_batch, ps_state = prepare_for_prefix_sharing_fsdp( micro_batch, ps_config, model_config={ diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py index fc79152d..68856be3 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py @@ -9,7 +9,7 @@ from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - plan_and_trim_microbatch_fsdp, + prepare_for_prefix_sharing_fsdp, forward_prefix_sharing_fsdp_micro_batch, restore_prefix_sharing_outputs_2d, ) @@ -132,7 +132,7 @@ def _baseline_attention(query, key, value): return torch.einsum("blmh,bmhd->blhd", probs, value) -def test_plan_and_trim_microbatch_fsdp_returns_trimmed_batch_and_runtime_state(): +def test_prepare_for_prefix_sharing_fsdp_returns_trimmed_batch_and_runtime_state(): config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { "input_ids": torch.tensor( @@ -172,7 +172,7 @@ def test_plan_and_trim_microbatch_fsdp_returns_trimmed_batch_and_runtime_state() ), } - trimmed_batch, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed_batch, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config) assert runtime_state is not None assert isinstance(runtime_state, PrefixSharingRuntimeState) @@ -193,7 +193,7 @@ def test_plan_and_trim_microbatch_fsdp_returns_trimmed_batch_and_runtime_state() assert torch.equal(trimmed_batch["position_ids"][1, 3:6], torch.tensor([3, 4, 5])) -def test_plan_and_trim_microbatch_fsdp_returns_none_when_no_sharing(): +def test_prepare_for_prefix_sharing_fsdp_returns_none_when_no_sharing(): config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { "input_ids": torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long), @@ -201,13 +201,13 @@ def test_plan_and_trim_microbatch_fsdp_returns_none_when_no_sharing(): "position_ids": torch.tensor([[0, 1, 2], [0, 1, 2]], dtype=torch.long), } - returned_batch, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) + returned_batch, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config) assert returned_batch is batch assert runtime_state is None -def test_plan_and_trim_microbatch_fsdp_trims_nested_remove_padding_batch(): +def test_prepare_for_prefix_sharing_fsdp_trims_nested_remove_padding_batch(): if not hasattr(torch, "nested"): pytest.skip("torch.nested is unavailable") config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) @@ -235,7 +235,7 @@ def test_plan_and_trim_microbatch_fsdp_trims_nested_remove_padding_batch(): ), } - trimmed_batch, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed_batch, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config) assert runtime_state is not None plan = runtime_state.prefix_sharing_plan @@ -273,7 +273,7 @@ def test_restore_prefix_sharing_outputs_2d_restores_interior_last_logits_entropy dtype=torch.long, ), } - _, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) + _, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config) assert runtime_state is not None vocab = 5 @@ -345,7 +345,7 @@ def test_prefix_sharing_fsdp_attention_runtime_scatter_dense_outputs(): dtype=torch.long, ), } - _, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) + _, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config) assert runtime_state is not None torch.manual_seed(1) @@ -686,7 +686,7 @@ def test_prefix_sharing_fsdp_attention_runtime_supports_packed_single_batch_shap "attention_mask": torch.ones(2, 5, dtype=torch.bool), "position_ids": torch.tensor([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long), } - _, runtime_state = plan_and_trim_microbatch_fsdp(batch, config) + _, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config) assert runtime_state is not None torch.manual_seed(2031) diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py index a57e5c73..6dabd727 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py @@ -21,7 +21,7 @@ from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - plan_and_trim_microbatch_fsdp, + prepare_for_prefix_sharing_fsdp, restore_prefix_sharing_outputs_2d, ) @@ -52,7 +52,7 @@ def test_scenario1_two_samples_share_arbitrary_prefix(): """A B C D E / A B C X Y -> provider=0, reuser=1, prefix_len=3.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5], [1, 2, 3, 10, 11]]) - trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -76,7 +76,7 @@ def test_scenario2_multiple_reusers_different_prefix_and_suffix_lens(): [1, 2, 20, 21], # reuser, prefix 2, suffix 2 (len 4) ] ) - trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -106,7 +106,7 @@ def test_scenario3_same_provider_serves_multiple_reusers(): [1, 2, 30, 40, 50, 60], # reuser 3 (prefix 2) ] ) - trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -128,7 +128,7 @@ def test_scenario4_chain_reuse_reuser_becomes_provider(): [1, 2, 3, 4, 10, 20, 30], # row2: reuser of row1 (prefix 5) ] ) - trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -146,7 +146,7 @@ def test_scenario5_no_shareable_prefix_returns_none(): """Completely different sequences -> fallback.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4], [5, 6, 7, 8]]) - returned, state = plan_and_trim_microbatch_fsdp(batch, config) + returned, state = prepare_for_prefix_sharing_fsdp(batch, config) assert returned is batch assert state is None @@ -156,7 +156,7 @@ def test_scenario5b_single_sample_returns_none(): """Single sample cannot share -> fallback.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5]]) - returned, state = plan_and_trim_microbatch_fsdp(batch, config) + returned, state = prepare_for_prefix_sharing_fsdp(batch, config) assert returned is batch assert state is None @@ -205,7 +205,7 @@ def test_scenario7_position_ids_preserve_absolute_positions(): """Reuser suffix position_ids must reflect original absolute positions.""" config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 10, 11, 12, 13]]) - trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan @@ -224,7 +224,7 @@ def test_scenario7_position_ids_preserve_absolute_positions(): def test_attention_runtime_records_stats_per_layer(): config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2) batch = _make_batch([[1, 2, 3, 4, 5], [1, 2, 3, 10, 11]]) - _, state = plan_and_trim_microbatch_fsdp(batch, config) + _, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None torch.manual_seed(42) diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py index 8bd7c91c..e6a72729 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py @@ -26,7 +26,7 @@ from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, - plan_and_trim_microbatch_fsdp, + prepare_for_prefix_sharing_fsdp, restore_prefix_sharing_outputs_2d, ) @@ -67,7 +67,7 @@ def test_reuser_suffix_attention_matches_baseline(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None plan = state.prefix_sharing_plan prefix_len = plan.prefix_lens[1] # 3 @@ -143,7 +143,7 @@ def test_gradient_flows_through_provider_prefix_kv(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - trimmed, state = plan_and_trim_microbatch_fsdp(batch, config) + trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None H, D = 2, 4 @@ -184,7 +184,7 @@ def test_restore_interior_prefix_matches_provider(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - _, state = plan_and_trim_microbatch_fsdp(batch, config) + _, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is not None # Simulate trimmed output: provider full, reuser has only suffix (prefix zeroed) @@ -238,7 +238,7 @@ def test_restore_prefix_last_recomputed_with_provider_logits_and_reuser_label(): [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long ), } - _, state = plan_and_trim_microbatch_fsdp(batch, config) + _, state = prepare_for_prefix_sharing_fsdp(batch, config) plan = state.prefix_sharing_plan prefix_len = plan.prefix_lens[1] # 3 @@ -275,7 +275,7 @@ def test_baseline_matches_when_no_sharing(): [[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.long ), } - returned, state = plan_and_trim_microbatch_fsdp(batch, config) + returned, state = prepare_for_prefix_sharing_fsdp(batch, config) assert state is None assert returned is batch # exact same object -> no transformation @@ -308,8 +308,8 @@ def test_padding_positions_do_not_contribute(): } H, D = 2, 4 - _, state_unpadded = plan_and_trim_microbatch_fsdp(batch_unpadded, config) - _, state_padded = plan_and_trim_microbatch_fsdp(batch_padded, config) + _, state_unpadded = prepare_for_prefix_sharing_fsdp(batch_unpadded, config) + _, state_padded = prepare_for_prefix_sharing_fsdp(batch_padded, config) assert state_unpadded is not None and state_padded is not None # Same plan semantics regardless of padding From d71fad266f200f30e54e39eb54be15f4a364f4ea Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:34:41 +0800 Subject: [PATCH 08/19] docs(verl_fsdp): add STEP comments to clarify prepare/execute phases Add explicit STEP comments to mark the prepare and execute phases in `forward_prefix_sharing_fsdp_micro_batch` and the internal steps inside `prepare_for_prefix_sharing_fsdp`. Co-authored-by: Cursor --- .../prefix_sharing/integrations/verl_fsdp.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index c822860a..fbc279dc 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -150,7 +150,7 @@ def forward_prefix_sharing_fsdp_micro_batch( with context as ctx, autocast: ######################################################### - # STEP 2: call the model with PrefixSharing r + # STEP 2: call the model with PrefixSharing runtime context ######################################################### model_output = _call_fsdp_model( model, @@ -223,12 +223,16 @@ def prepare_for_prefix_sharing_fsdp( ] sequences = extract_seq_from_dense_tensor(input_ids, valid_indices) - # plan for PrefixSharing + ######################################################### + # STEP 1: plan for PrefixSharing + ######################################################### prefix_sharing_plan = PrefixSharingPlanner(ps_config).plan(sequences) if not prefix_sharing_plan.has_sharing: return micro_batch, None - # trim the micro-batch + ######################################################### + # STEP 2: trim the micro-batch + ######################################################### if input_ids_is_nested: # nested tensor trimmed_micro_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor( micro_batch, prefix_sharing_plan @@ -238,8 +242,14 @@ def prepare_for_prefix_sharing_fsdp( micro_batch, prefix_sharing_plan, valid_indices ) - # build layout and runtime state for the trimmed micro-batch + ######################################################### + # STEP 3: build batch layout for the micro-batch + ######################################################### packed_batch_layout = PackedBatchLayout.from_kept_position_rows(kept_position_rows) + + ######################################################### + # STEP 4: build runtime state for the micro-batch + ######################################################### runtime_state = PrefixSharingRuntimeState( prefix_sharing_plan=prefix_sharing_plan, attention_backend=get_backend_instance(ps_config, backend), @@ -247,6 +257,7 @@ def prepare_for_prefix_sharing_fsdp( parallel_info=MegatronParallelInfo(), kept_position_ids=trimmed_micro_batch.get("position_ids"), ) + return trimmed_micro_batch, runtime_state From 722d290fd28e35d45f7b187ab3ebaae2975cc9e3 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:46:16 +0800 Subject: [PATCH 09/19] refactor(verl_fsdp): rename forward_prefix_sharing_fsdp_micro_batch Rename `forward_prefix_sharing_fsdp_micro_batch` to `forward_prefix_sharing_micro_batch_fsdp` for naming consistency with `prepare_for_prefix_sharing_fsdp`. Update all callers, imports, and docs. Also clarify the STEP comments and variable names inside the forward helper. Co-authored-by: Cursor --- docs/developer-docs/impr-refactor.md | 4 +- .../prefix_sharing/integrations/__init__.py | 4 +- .../prefix_sharing/integrations/verl_fsdp.py | 44 +++++++++---------- .../patches/verl080_fsdp/forward_step.py | 4 +- .../tests/unit_test/test_verl_fsdp_adapter.py | 10 ++--- 5 files changed, 33 insertions(+), 33 deletions(-) diff --git a/docs/developer-docs/impr-refactor.md b/docs/developer-docs/impr-refactor.md index 619f3bd7..28b8799b 100644 --- a/docs/developer-docs/impr-refactor.md +++ b/docs/developer-docs/impr-refactor.md @@ -270,11 +270,11 @@ integrations/verl_fsdp.py # FSDP 专属流程 - `build_prefix_sharing_micro_batch_fsdp()` - `PrefixSharingFSDPAttentionRuntime` - `restore_prefix_sharing_outputs_2d()` -- `forward_prefix_sharing_fsdp_micro_batch()` +- `forward_prefix_sharing_micro_batch_fsdp()` 问题： -- `forward_prefix_sharing_fsdp_micro_batch()` 更像测试/fake engine helper，不一定是合入 verl 的主路径，应避免让 reviewer 误以为这是生产接入方式。 +- `forward_prefix_sharing_micro_batch_fsdp()` 更像测试/fake engine helper，不一定是合入 verl 的主路径，应避免让 reviewer 误以为这是生产接入方式。 - `PrefixSharingFSDPAttentionRuntime.forward()` 直接忽略 `attn_func/attention_mask/kwargs`，对 HF attention 接口兼容性说明不足。 - dense `[B,L,H,D]` 与 packed `[1,T,H,D]` 两种路径混在同一个 runtime，缺少清晰的 input contract。 - FSDP restore 的 logits/log_probs/entropy/attention_output copy 语义复杂，需要更强测试和更小函数。 diff --git a/prefix-sharing/prefix_sharing/integrations/__init__.py b/prefix-sharing/prefix_sharing/integrations/__init__.py index e0fbdb27..0b6c6477 100644 --- a/prefix-sharing/prefix_sharing/integrations/__init__.py +++ b/prefix-sharing/prefix_sharing/integrations/__init__.py @@ -17,7 +17,7 @@ from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, prepare_for_prefix_sharing_fsdp, - forward_prefix_sharing_fsdp_micro_batch, + forward_prefix_sharing_micro_batch_fsdp, restore_prefix_sharing_outputs_2d, ) from prefix_sharing.integrations.megatron_runtime import ( @@ -39,6 +39,6 @@ "get_megatron_parallel_info", "PrefixSharingFSDPAttentionRuntime", "prepare_for_prefix_sharing_fsdp", - "forward_prefix_sharing_fsdp_micro_batch", + "forward_prefix_sharing_micro_batch_fsdp", "restore_prefix_sharing_outputs_2d", ] diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index fbc279dc..23ed71fb 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -113,7 +113,7 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any, return dense_output -def forward_prefix_sharing_fsdp_micro_batch( +def forward_prefix_sharing_micro_batch_fsdp( micro_batch: Any, model: Any, ps_config: PrefixSharingConfig, @@ -126,40 +126,40 @@ def forward_prefix_sharing_fsdp_micro_batch( entropy_fn: Any | None = None, autocast_context: Any | None = None, ) -> dict[str, Any]: - """Run one dense verl/FSDP-style micro-batch with PrefixSharing. - - This is the executable helper used by fake/local FSDP tests and by engines - that do not expose prepare hooks: - prepare the micro-batch, open the runtime context, run the model with a - PrefixSharing attention runtime, compute token-level outputs, then restore - reuser prefix columns. Real verl FSDP patching should prefer the engine's - own ``prepare_model_inputs`` / ``prepare_model_outputs`` path. - """ + """Run one forward pass of a verl/FSDP-style micro-batch with PrefixSharing.""" ######################################################### - # STEP 1: prepare for prefix sharing + # STEP 1: pre-processing inputs for PrefixSharing ######################################################### - trimmed_micro_batch, runtime_state = prepare_for_prefix_sharing_fsdp( + micro_batch_modified, prefix_sharing_runtime_state = prepare_for_prefix_sharing_fsdp( micro_batch, ps_config, model_config=model_config, backend=backend, ) - context = prefix_sharing_runtime_context(runtime_state) if runtime_state is not None else nullcontext(None) - autocast = autocast_context if autocast_context is not None else nullcontext() + ######################################################### + # STEP 2: create PrefixSharing runtime context + ######################################################### + context = prefix_sharing_runtime_context(prefix_sharing_runtime_state) if prefix_sharing_runtime_state is not None else nullcontext(None) + autocast = autocast_context if autocast_context is not None else nullcontext() with context as ctx, autocast: + ######################################################### - # STEP 2: call the model with PrefixSharing runtime context + # STEP 3: call the model with PrefixSharing+FSDP attention runtime ######################################################### model_output = _call_fsdp_model( model, - trimmed_micro_batch, + micro_batch_modified, prefix_sharing_runtime=PrefixSharingFSDPAttentionRuntime( num_layers=model.config.num_hidden_layers if hasattr(model, "config") else 0, ), - enable_prefix_sharing=runtime_state is not None, + enable_prefix_sharing=prefix_sharing_runtime_state is not None, ) + + ######################################################### + # STEP 4: post-processing outputs for PrefixSharing + ######################################################### logits = _extract_logits(model_output) / float(temperature) output = { "model_output": model_output, @@ -231,7 +231,7 @@ def prepare_for_prefix_sharing_fsdp( return micro_batch, None ######################################################### - # STEP 2: trim the micro-batch + # STEP 2: trim redundant prefix of reusers in the micro-batch ######################################################### if input_ids_is_nested: # nested tensor trimmed_micro_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor( @@ -243,14 +243,14 @@ def prepare_for_prefix_sharing_fsdp( ) ######################################################### - # STEP 3: build batch layout for the micro-batch + # STEP 3: build batch layout for the trimmed micro-batch ######################################################### packed_batch_layout = PackedBatchLayout.from_kept_position_rows(kept_position_rows) ######################################################### - # STEP 4: build runtime state for the micro-batch + # STEP 4: build runtime state for the trimmed micro-batch ######################################################### - runtime_state = PrefixSharingRuntimeState( + prefix_sharing_runtime_state = PrefixSharingRuntimeState( prefix_sharing_plan=prefix_sharing_plan, attention_backend=get_backend_instance(ps_config, backend), packed_batch_layout=packed_batch_layout, @@ -258,7 +258,7 @@ def prepare_for_prefix_sharing_fsdp( kept_position_ids=trimmed_micro_batch.get("position_ids"), ) - return trimmed_micro_batch, runtime_state + return trimmed_micro_batch, prefix_sharing_runtime_state def restore_prefix_sharing_outputs_2d( diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index d984df0d..1e15a8c0 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -101,13 +101,13 @@ def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forwar self, micro_batch, loss_function, forward_only, ps_config, ) - from prefix_sharing.integrations.verl_fsdp import forward_prefix_sharing_fsdp_micro_batch + from prefix_sharing.integrations.verl_fsdp import forward_prefix_sharing_micro_batch_fsdp calculate_entropy = bool( _read_runtime_value(self.engine_config, micro_batch, "calculate_entropy", default=False) ) temperature = _read_temperature(micro_batch) - output = forward_prefix_sharing_fsdp_micro_batch( + output = forward_prefix_sharing_micro_batch_fsdp( micro_batch, self.module, ps_config, diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py index 68856be3..ab54f8c0 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py @@ -10,7 +10,7 @@ from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, prepare_for_prefix_sharing_fsdp, - forward_prefix_sharing_fsdp_micro_batch, + forward_prefix_sharing_micro_batch_fsdp, restore_prefix_sharing_outputs_2d, ) from prefix_sharing.integrations.verl_mcore import PrefixSharingRuntimeState @@ -368,7 +368,7 @@ def test_prefix_sharing_fsdp_attention_runtime_scatter_dense_outputs(): assert not torch.allclose(dense_output[1, 3:6], torch.zeros_like(dense_output[1, 3:6])) -def test_forward_prefix_sharing_fsdp_micro_batch_matches_tiny_hf_model_baseline(): +def test_forward_prefix_sharing_micro_batch_fsdp_matches_tiny_hf_model_baseline(): torch.manual_seed(2026) config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { @@ -403,7 +403,7 @@ def test_forward_prefix_sharing_fsdp_micro_batch_matches_tiny_hf_model_baseline( baseline_log_probs = _mock_log_probs_fn(baseline_logits, labels) baseline_entropy = _entropy_from_logits(baseline_logits) - prefix_output = forward_prefix_sharing_fsdp_micro_batch( + prefix_output = forward_prefix_sharing_micro_batch_fsdp( batch, model, config, @@ -418,7 +418,7 @@ def test_forward_prefix_sharing_fsdp_micro_batch_matches_tiny_hf_model_baseline( assert torch.allclose(prefix_output["attention_output"], baseline.attention_output, atol=1e-5) -def test_forward_prefix_sharing_fsdp_micro_batch_keeps_provider_prefix_grad_path(): +def test_forward_prefix_sharing_micro_batch_fsdp_keeps_provider_prefix_grad_path(): torch.manual_seed(2027) config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { @@ -443,7 +443,7 @@ def test_forward_prefix_sharing_fsdp_micro_batch_keeps_provider_prefix_grad_path batch["labels"] = labels model = _TinyHFStyleModel(vocab_size=32) - output = forward_prefix_sharing_fsdp_micro_batch( + output = forward_prefix_sharing_micro_batch_fsdp( batch, model, config, From 7eee5d522ded3bed22e894c49c9604aecf03a39b Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:50:08 +0800 Subject: [PATCH 10/19] refactor(verl_fsdp): rename attention runtime parameter in _call_fsdp_model Rename the `prefix_sharing_runtime` keyword parameter of `_call_fsdp_model` to `attention_runtime` to reflect that the helper is agnostic to the specific attention runtime implementation. The model input key remains `prefix_sharing_runtime` since that is the contract expected by the model forward / attention patch. Co-authored-by: Cursor --- prefix-sharing/prefix_sharing/integrations/verl_fsdp.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index 23ed71fb..40398f10 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -151,7 +151,7 @@ def forward_prefix_sharing_micro_batch_fsdp( model_output = _call_fsdp_model( model, micro_batch_modified, - prefix_sharing_runtime=PrefixSharingFSDPAttentionRuntime( + attention_runtime=PrefixSharingFSDPAttentionRuntime( num_layers=model.config.num_hidden_layers if hasattr(model, "config") else 0, ), enable_prefix_sharing=prefix_sharing_runtime_state is not None, @@ -426,7 +426,7 @@ def _call_fsdp_model( model: Any, micro_batch: Any, *, - prefix_sharing_runtime: PrefixSharingFSDPAttentionRuntime, + attention_runtime: PrefixSharingFSDPAttentionRuntime, enable_prefix_sharing: bool, ) -> Any: model_inputs = { @@ -436,7 +436,7 @@ def _call_fsdp_model( } model_inputs["use_cache"] = False if enable_prefix_sharing: - model_inputs["prefix_sharing_runtime"] = prefix_sharing_runtime + model_inputs["prefix_sharing_runtime"] = attention_runtime try: return model(**model_inputs) except TypeError: From f2d21e9618c72a020799f8cf2b0a9e7f3d5bba76 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:58:19 +0800 Subject: [PATCH 11/19] refactor(fsdp_patch): hoist imports in _forward_step_with_engine_prepare Move function-local imports in `_forward_step_with_engine_prepare` to the module top level for consistency and to avoid repeated import overhead. Other functions in the same file are left unchanged as requested. Co-authored-by: Cursor --- .../setup/patches/verl080_fsdp/forward_step.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 1e15a8c0..2eab6b9c 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -6,11 +6,16 @@ from __future__ import annotations import os +from contextlib import nullcontext from typing import Any, Callable +import torch import torch.utils.checkpoint as _ckpt from prefix_sharing.core.config import PrefixSharingConfig +from prefix_sharing.integrations.context import create_prefix_sharing_context +from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime +from prefix_sharing.integrations.verl_fsdp import prepare_for_prefix_sharing_fsdp from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope @@ -156,15 +161,6 @@ def _forward_step_with_engine_prepare( forward_only: bool, ps_config: Any, ) -> Any: - import torch - from contextlib import nullcontext - - from prefix_sharing.integrations.context import create_prefix_sharing_context - from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime - from prefix_sharing.integrations.verl_fsdp import prepare_for_prefix_sharing_fsdp - - from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope - profiler = ProfilerScope.current() if profiler is not None: profiler.start_phase(PerfProfiler.PHASE_PLAN) From 296ce1eb7ace37ae67d4199318e96e88cd02440f Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:58:57 +0800 Subject: [PATCH 12/19] refactor(fsdp_patch): align variable naming with verl_fsdp helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In `_forward_step_with_engine_prepare`, rename: - `trimmed_micro_batch` → `micro_batch_modified` - `ps_state` → `prefix_sharing_runtime_state` These names match the standalone helper in `verl_fsdp.py` and make the relationship between the two paths clearer. Co-authored-by: Cursor --- .../setup/patches/verl080_fsdp/forward_step.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 2eab6b9c..a267464c 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -171,7 +171,7 @@ def _forward_step_with_engine_prepare( if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: _dump_full_input_ids_only(micro_batch, "train") - trimmed_micro_batch, ps_state = prepare_for_prefix_sharing_fsdp( + micro_batch_modified, prefix_sharing_runtime_state = prepare_for_prefix_sharing_fsdp( micro_batch, ps_config, model_config={ @@ -193,8 +193,8 @@ def _forward_step_with_engine_prepare( if profiler is not None: profiler.stop_phase(PerfProfiler.PHASE_PLAN) # Detect, plan, and trim on CPU. - if ps_state is None: - return _call_original_like_engine(self, trimmed_micro_batch, loss_function, forward_only) + if prefix_sharing_runtime_state is None: + return _call_original_like_engine(self, micro_batch_modified, loss_function, forward_only) if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: from prefix_sharing.tools.diagnostic_dump import dump_fsdp_on_metadata_verl080 @@ -202,7 +202,7 @@ def _forward_step_with_engine_prepare( diagnostic_tag = "train" if self.module.training else "old" dump_fsdp_on_metadata_verl080( micro_batch, - ps_state.prefix_sharing_plan, + prefix_sharing_runtime_state.prefix_sharing_plan, diagnostic_tag, ) @@ -211,7 +211,7 @@ def _forward_step_with_engine_prepare( getattr(getattr(self, "module", None), "config", None), "num_hidden_layers", 0)) or 0 - model_inputs, output_args = self.prepare_model_inputs(micro_batch=trimmed_micro_batch) + model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch_modified) model_inputs["prefix_sharing_runtime"] = PrefixSharingFSDPAttentionRuntime() model_inputs["prefix_sharing_runtime"].num_layers = _diag_num_layers autocast_dtype = getattr(self, "_autocast_dtype", torch.float32) @@ -222,7 +222,7 @@ def _forward_step_with_engine_prepare( else torch.autocast(device_type=device_name, dtype=autocast_dtype) ) # ── Create PS context with manual lifecycle (survives backward for AC) ── - ctx, ctx_cleanup = create_prefix_sharing_context(ps_state) + ctx, ctx_cleanup = create_prefix_sharing_context(prefix_sharing_runtime_state) # Set _ps_ctx on every attention module so the attention patch reads # the context from the module itself rather than ContextVar (compatible @@ -261,7 +261,7 @@ def _forward_step_with_engine_prepare( model_output = self.prepare_model_outputs( output=raw_output, output_args=output_args, - micro_batch=trimmed_micro_batch, + micro_batch=micro_batch_modified, logits_processor_func=loss_function, ) @@ -276,7 +276,7 @@ def _forward_step_with_engine_prepare( dump_fsdp_model_output_2d_verl080( model_output, - list(ps_state.prefix_sharing_plan.original_lengths), + list(prefix_sharing_runtime_state.prefix_sharing_plan.original_lengths), diagnostic_tag, ) @@ -529,7 +529,7 @@ def _read_temperature(micro_batch: Any) -> float: def _dump_full_input_ids_only(micro_batch: Any, tag: str) -> None: """Dump the original (full) input_ids before prefix sharing trimming. - The ON path dumps ``input_ids_train.pt`` from the ``trimmed_micro_batch``, + The ON path dumps ``input_ids_train.pt`` from the ``micro_batch_modified``, which has shared prefix tokens removed. This helper saves the **original** ``micro_batch`` input_ids so that ``cmp_diag_verl080`` can compare the full input against the OFF baseline, rather than reporting 186+ differing From 7a4e0a299c71ecdb6d0626d496d8a9880fbe05ed Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 16:59:57 +0800 Subject: [PATCH 13/19] style(fsdp_patch): collapse _read_runtime_value calls to one line Collapse the two multi-line `_read_runtime_value` calls in the `model_config` dict to single-line form for readability. Co-authored-by: Cursor --- .../setup/patches/verl080_fsdp/forward_step.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index a267464c..40a797c9 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -171,23 +171,16 @@ def _forward_step_with_engine_prepare( if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: _dump_full_input_ids_only(micro_batch, "train") + ######################################################### + # STEP 1: pre-processing inputs for PrefixSharing + ######################################################### micro_batch_modified, prefix_sharing_runtime_state = prepare_for_prefix_sharing_fsdp( micro_batch, ps_config, model_config={ "model_type": "text_only_causal_lm", - "ulysses_sequence_parallel_size": _read_runtime_value( - self.engine_config, - micro_batch, - "ulysses_sequence_parallel_size", - default=1, - ), - "use_fused_kernels": _read_runtime_value( - self.engine_config, - micro_batch, - "use_fused_kernels", - default=False, - ), + "ulysses_sequence_parallel_size": _read_runtime_value(self.engine_config, micro_batch, "ulysses_sequence_parallel_size", default=1), + "use_fused_kernels": _read_runtime_value(self.engine_config, micro_batch, "use_fused_kernels", default=False), }, ) if profiler is not None: From 180452d362d87e056fa6128b65155242d39e46de Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 17:13:33 +0800 Subject: [PATCH 14/19] style(fsdp_patch): inline one-time device_name variable Remove the `device_name` intermediate variable in `_forward_step_with_engine_prepare` and call `_read_device_name()` directly inside `torch.autocast()`. Co-authored-by: Cursor --- .../setup/patches/verl080_fsdp/forward_step.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 40a797c9..1e0ed590 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -185,7 +185,7 @@ def _forward_step_with_engine_prepare( ) if profiler is not None: profiler.stop_phase(PerfProfiler.PHASE_PLAN) # Detect, plan, and trim on CPU. - + if prefix_sharing_runtime_state is None: return _call_original_like_engine(self, micro_batch_modified, loss_function, forward_only) @@ -208,11 +208,10 @@ def _forward_step_with_engine_prepare( model_inputs["prefix_sharing_runtime"] = PrefixSharingFSDPAttentionRuntime() model_inputs["prefix_sharing_runtime"].num_layers = _diag_num_layers autocast_dtype = getattr(self, "_autocast_dtype", torch.float32) - device_name = _read_device_name() autocast_ctx = ( nullcontext() if autocast_dtype == torch.float32 - else torch.autocast(device_type=device_name, dtype=autocast_dtype) + else torch.autocast(device_type=_read_device_name(), dtype=autocast_dtype) ) # ── Create PS context with manual lifecycle (survives backward for AC) ── ctx, ctx_cleanup = create_prefix_sharing_context(prefix_sharing_runtime_state) From 5d3c65466986607909e1262e3979b512d17e764e Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 17:20:25 +0800 Subject: [PATCH 15/19] refactor(verl_fsdp): rename forward helper to forward_step_without_engine_prepare Rename `forward_prefix_sharing_micro_batch_fsdp` to `forward_step_without_engine_prepare` to mirror the production path `_forward_step_with_engine_prepare` and make the dispatch logic clearer. Update all imports, callers, tests, and docs. Co-authored-by: Cursor --- docs/developer-docs/impr-refactor.md | 4 ++-- prefix-sharing/prefix_sharing/integrations/__init__.py | 4 ++-- .../prefix_sharing/integrations/verl_fsdp.py | 9 ++++----- .../setup/patches/verl080_fsdp/forward_step.py | 4 ++-- .../tests/unit_test/test_verl_fsdp_adapter.py | 10 +++++----- 5 files changed, 15 insertions(+), 16 deletions(-) diff --git a/docs/developer-docs/impr-refactor.md b/docs/developer-docs/impr-refactor.md index 28b8799b..b88f46b9 100644 --- a/docs/developer-docs/impr-refactor.md +++ b/docs/developer-docs/impr-refactor.md @@ -270,11 +270,11 @@ integrations/verl_fsdp.py # FSDP 专属流程 - `build_prefix_sharing_micro_batch_fsdp()` - `PrefixSharingFSDPAttentionRuntime` - `restore_prefix_sharing_outputs_2d()` -- `forward_prefix_sharing_micro_batch_fsdp()` +- `forward_step_without_engine_prepare()` 问题： -- `forward_prefix_sharing_micro_batch_fsdp()` 更像测试/fake engine helper，不一定是合入 verl 的主路径，应避免让 reviewer 误以为这是生产接入方式。 +- `forward_step_without_engine_prepare()` 更像测试/fake engine helper，不一定是合入 verl 的主路径，应避免让 reviewer 误以为这是生产接入方式。 - `PrefixSharingFSDPAttentionRuntime.forward()` 直接忽略 `attn_func/attention_mask/kwargs`，对 HF attention 接口兼容性说明不足。 - dense `[B,L,H,D]` 与 packed `[1,T,H,D]` 两种路径混在同一个 runtime，缺少清晰的 input contract。 - FSDP restore 的 logits/log_probs/entropy/attention_output copy 语义复杂，需要更强测试和更小函数。 diff --git a/prefix-sharing/prefix_sharing/integrations/__init__.py b/prefix-sharing/prefix_sharing/integrations/__init__.py index 0b6c6477..7961771b 100644 --- a/prefix-sharing/prefix_sharing/integrations/__init__.py +++ b/prefix-sharing/prefix_sharing/integrations/__init__.py @@ -17,7 +17,7 @@ from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, prepare_for_prefix_sharing_fsdp, - forward_prefix_sharing_micro_batch_fsdp, + forward_step_without_engine_prepare, restore_prefix_sharing_outputs_2d, ) from prefix_sharing.integrations.megatron_runtime import ( @@ -39,6 +39,6 @@ "get_megatron_parallel_info", "PrefixSharingFSDPAttentionRuntime", "prepare_for_prefix_sharing_fsdp", - "forward_prefix_sharing_micro_batch_fsdp", + "forward_step_without_engine_prepare", "restore_prefix_sharing_outputs_2d", ] diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index 40398f10..d1825c9b 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -113,7 +113,7 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any, return dense_output -def forward_prefix_sharing_micro_batch_fsdp( +def forward_step_without_engine_prepare( micro_batch: Any, model: Any, ps_config: PrefixSharingConfig, @@ -160,11 +160,10 @@ def forward_prefix_sharing_micro_batch_fsdp( ######################################################### # STEP 4: post-processing outputs for PrefixSharing ######################################################### + output = {} logits = _extract_logits(model_output) / float(temperature) - output = { - "model_output": model_output, - "logits": logits.clone(), - } + output["model_output"] = model_output + output["logits"] = logits.clone() labels = _labels_for_log_probs(micro_batch) if labels is not None: output["log_probs"] = _compute_log_probs(logits, labels, log_probs_fn) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index 1e0ed590..ddd6871a 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -106,13 +106,13 @@ def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forwar self, micro_batch, loss_function, forward_only, ps_config, ) - from prefix_sharing.integrations.verl_fsdp import forward_prefix_sharing_micro_batch_fsdp + from prefix_sharing.integrations.verl_fsdp import forward_step_without_engine_prepare calculate_entropy = bool( _read_runtime_value(self.engine_config, micro_batch, "calculate_entropy", default=False) ) temperature = _read_temperature(micro_batch) - output = forward_prefix_sharing_micro_batch_fsdp( + output = forward_step_without_engine_prepare( micro_batch, self.module, ps_config, diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py index ab54f8c0..4c694c40 100644 --- a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py +++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py @@ -10,7 +10,7 @@ from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, prepare_for_prefix_sharing_fsdp, - forward_prefix_sharing_micro_batch_fsdp, + forward_step_without_engine_prepare, restore_prefix_sharing_outputs_2d, ) from prefix_sharing.integrations.verl_mcore import PrefixSharingRuntimeState @@ -368,7 +368,7 @@ def test_prefix_sharing_fsdp_attention_runtime_scatter_dense_outputs(): assert not torch.allclose(dense_output[1, 3:6], torch.zeros_like(dense_output[1, 3:6])) -def test_forward_prefix_sharing_micro_batch_fsdp_matches_tiny_hf_model_baseline(): +def test_forward_step_without_engine_prepare_matches_tiny_hf_model_baseline(): torch.manual_seed(2026) config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { @@ -403,7 +403,7 @@ def test_forward_prefix_sharing_micro_batch_fsdp_matches_tiny_hf_model_baseline( baseline_log_probs = _mock_log_probs_fn(baseline_logits, labels) baseline_entropy = _entropy_from_logits(baseline_logits) - prefix_output = forward_prefix_sharing_micro_batch_fsdp( + prefix_output = forward_step_without_engine_prepare( batch, model, config, @@ -418,7 +418,7 @@ def test_forward_prefix_sharing_micro_batch_fsdp_matches_tiny_hf_model_baseline( assert torch.allclose(prefix_output["attention_output"], baseline.attention_output, atol=1e-5) -def test_forward_prefix_sharing_micro_batch_fsdp_keeps_provider_prefix_grad_path(): +def test_forward_step_without_engine_prepare_keeps_provider_prefix_grad_path(): torch.manual_seed(2027) config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3) batch = { @@ -443,7 +443,7 @@ def test_forward_prefix_sharing_micro_batch_fsdp_keeps_provider_prefix_grad_path batch["labels"] = labels model = _TinyHFStyleModel(vocab_size=32) - output = forward_prefix_sharing_micro_batch_fsdp( + output = forward_step_without_engine_prepare( batch, model, config, From eb18596b0aad64cda93843a67d83bfcf7b71f4bc Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 17:22:06 +0800 Subject: [PATCH 16/19] refactor(fsdp_patch): pass model_config to _forward_step_with_engine_prepare Construct `model_config` once in `patched_forward_step`, reuse it for `ps_config.validate`, and pass it as a parameter to `_forward_step_with_engine_prepare` and `forward_step_without_engine_prepare`. This eliminates duplicate `_read_runtime_value` calls and aligns the two dispatch paths with the same helper API. Co-authored-by: Cursor --- .../patches/verl080_fsdp/forward_step.py | 38 +++++-------------- 1 file changed, 10 insertions(+), 28 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index ddd6871a..d9aee013 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -80,30 +80,19 @@ def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forwar # Local unit tests use plain dict / fake engine without verl device helpers. pass - ulysses_sp_size = _read_runtime_value( - self.engine_config, - micro_batch, - "ulysses_sequence_parallel_size", - default=1, - ) - use_fused_kernels = _read_runtime_value( - self.engine_config, - micro_batch, - "use_fused_kernels", - default=False, - ) + model_config = { + "model_type": "text_only_causal_lm", + "ulysses_sequence_parallel_size": _read_runtime_value(self.engine_config, micro_batch, "ulysses_sequence_parallel_size", default=1), + "use_fused_kernels": _read_runtime_value(self.engine_config, micro_batch, "use_fused_kernels", default=False), + } ps_config.validate( - model_config={ - "model_type": "text_only_causal_lm", - "ulysses_sequence_parallel_size": ulysses_sp_size, - "use_fused_kernels": use_fused_kernels, - }, + model_config=model_config, integrate_mode="verl_fsdp", ) if hasattr(self, "prepare_model_inputs") and hasattr(self, "prepare_model_outputs"): return _forward_step_with_engine_prepare( - self, micro_batch, loss_function, forward_only, ps_config, + self, micro_batch, loss_function, forward_only, ps_config, model_config, ) from prefix_sharing.integrations.verl_fsdp import forward_step_without_engine_prepare @@ -116,11 +105,7 @@ def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forwar micro_batch, self.module, ps_config, - model_config={ - "model_type": "text_only_causal_lm", - "ulysses_sequence_parallel_size": ulysses_sp_size, - "use_fused_kernels": use_fused_kernels, - }, + model_config=model_config, temperature=temperature, calculate_entropy=calculate_entropy, entropy_fn=getattr(self, "compute_entropy_from_logits", None), @@ -160,6 +145,7 @@ def _forward_step_with_engine_prepare( loss_function: Any, forward_only: bool, ps_config: Any, + model_config: Any, ) -> Any: profiler = ProfilerScope.current() if profiler is not None: @@ -177,11 +163,7 @@ def _forward_step_with_engine_prepare( micro_batch_modified, prefix_sharing_runtime_state = prepare_for_prefix_sharing_fsdp( micro_batch, ps_config, - model_config={ - "model_type": "text_only_causal_lm", - "ulysses_sequence_parallel_size": _read_runtime_value(self.engine_config, micro_batch, "ulysses_sequence_parallel_size", default=1), - "use_fused_kernels": _read_runtime_value(self.engine_config, micro_batch, "use_fused_kernels", default=False), - }, + model_config=model_config, ) if profiler is not None: profiler.stop_phase(PerfProfiler.PHASE_PLAN) # Detect, plan, and trim on CPU. From 304fa93d24ab0e8ee72543b72606348c77ff4acb Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 18:46:31 +0800 Subject: [PATCH 17/19] refactor(verl_utils): relocate restore helpers from verl_mcore Move restore_reuser_prefix_columns_2d, restore_via_2d_unfold_verl080 and their helpers from verl_mcore to verl_utils so that verl_mcore only contains Megatron-specific micro-batch building logic. Update all import sites and tests accordingly. Co-authored-by: Cursor --- .../prefix_sharing/integrations/__init__.py | 6 +- .../prefix_sharing/integrations/verl_mcore.py | 298 ------------------ .../prefix_sharing/integrations/verl_utils.py | 290 +++++++++++++++++ .../verl080_mcore0161_ms0160/forward_step.py | 2 +- .../test_verl_megatron_runtime_helpers.py | 2 +- .../test_verl080_restore_e2e.py | 6 +- .../unit_test/test_restore_unfold_verl080.py | 4 +- 7 files changed, 299 insertions(+), 309 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/__init__.py b/prefix-sharing/prefix_sharing/integrations/__init__.py index 7961771b..4ab50dc1 100644 --- a/prefix-sharing/prefix_sharing/integrations/__init__.py +++ b/prefix-sharing/prefix_sharing/integrations/__init__.py @@ -9,11 +9,11 @@ from prefix_sharing.integrations.parallel_info import MegatronParallelInfo, get_megatron_parallel_info from prefix_sharing.backends.packed_layout import PackedBatchLayout from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState -from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config -from prefix_sharing.integrations.verl_mcore import ( - build_prefix_sharing_micro_batch_verl080, +from prefix_sharing.integrations.verl_utils import ( + read_ps_config_from_engine_config, restore_reuser_prefix_columns_2d, ) +from prefix_sharing.integrations.verl_mcore import build_prefix_sharing_micro_batch_verl080 from prefix_sharing.integrations.verl_fsdp import ( PrefixSharingFSDPAttentionRuntime, prepare_for_prefix_sharing_fsdp, diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index 70da1d5d..666359e3 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -20,7 +20,6 @@ from prefix_sharing.core.config import PrefixSharingConfig from prefix_sharing.core.planner import PrefixSharingPlan from prefix_sharing.core.planner import PrefixSharingPlanner -from prefix_sharing.integrations.context import current_prefix_sharing_context from prefix_sharing.integrations.parallel_info import get_megatron_parallel_info from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState from prefix_sharing.integrations.verl_utils import collect_kept_position_rows @@ -32,303 +31,6 @@ from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config -def restore_reuser_prefix_columns_2d( - output: dict[str, Any], - vocab_parallel_log_probs_fn: Any, - vocab_parallel_entropy_fn: Any = None, -) -> dict[str, Any]: - """Restore reuser prefix columns in 2D space — build_kv-style slice + concat. - - Mirrors :meth:`TorchReferenceBackend.build_kv` - (``torch.cat([provider_kv[:prefix_len], own_suffix])``): instead of writing - each prefix token one scalar at a time, the whole prefix interval is sliced - off the **direct provider's already-restored 2D row** and only the single - prefix-last logprob is recomputed. - - Per ``reuser_idx`` with direct provider ``provider_idx = provider_index[reuser_idx]`` - and ``P = prefix_lens[reuser_idx]`` (columns are identity-mapped in the - unfolded 2D tensor, so ``target_2d_pos`` == column): - - - **interior ``[0, P-2]``**: ``log_probs[reuser_idx, 0:P-1] = log_probs[provider_idx, 0:P-1]`` - (bulk copy). Identical across the shared prefix (same logits + labels), - and ``provider_idx`` was restored earlier in the batch-order loop, so its - row already holds correct values — no per-position provider resolution - needed. - - **prefix-last ``P-1``**: recompute ``log_probs[reuser_idx, P-1]`` from the - saved provider logits + the reuser's own first-suffix label (differs from - the provider's). When the reuser has no suffix (``suffix_len == 0``) the - planner emits no prefix-last spec; that column is masked downstream, so - the provider's value is copied as a safe placeholder. - - **entropy ``[0, P-1]``**: ``entropy[reuser_idx, 0:P] = entropy[provider_idx, 0:P]`` - (whole prefix copied, including prefix-last — entropy is label-independent). - - Rows are visited in ``range(B)`` order so a provider is always restored - before any reuser that reads it (the same online-detector invariant - ``build_kv`` relies on). - - Args: - output: Output dict with ``log_probs`` [B, L] and optionally - ``entropy`` [B, L] in 2D space (unfolded from the trimmed - NestedTensor by :func:`restore_via_2d_unfold_verl080`). - vocab_parallel_log_probs_fn: ``logits [1, V//tp]``, ``label [1]`` → - scalar; used only for the prefix-last recompute. - vocab_parallel_entropy_fn: Retained for call-site compatibility; - unused (entropy is copied, never recomputed). - - Returns: - ``output`` with ``log_probs`` and ``entropy`` mutated in-place. - """ - - ctx = current_prefix_sharing_context() - if ctx is None: - return output - plan = ctx.prefix_sharing_plan - # Guard on reuser presence (not on prefix_last_restore_indices): a batch - # whose reusers all have suffix_len == 0 emits no prefix-last spec but still - # needs its interior prefix columns restored. - if not plan.has_sharing: - return output - - import torch - - log_probs = output.get("log_probs") - if log_probs is None: - return output - entropy = output.get("entropy") - - provider_index = plan.provider_index - prefix_lens = plan.prefix_lens - - # reuser row → its prefix-last restore spec (one per reuser-with-suffix; - # interior positions have no spec — they are bulk-sliced below). - prefix_last_spec_by_reuser = { - spec.reuse_idx_in_batch: spec for spec in ctx.prefix_last_restore_indices - } - - restored_reusers = 0 - # Row 0 is always a provider (nothing precedes it to reuse), so start at 1. - # A reuser's provider always has a smaller batch index (online-detector - # invariant), so it is already restored when we reach reuser_idx. - for reuser_idx in range(1, len(prefix_lens)): - prefix_len = prefix_lens[reuser_idx] - if provider_index[reuser_idx] == reuser_idx or prefix_len <= 0: - continue # provider / non-reuser: row already complete - provider_idx = provider_index[reuser_idx] - - # interior [0, prefix_len-2]: bulk-copy from the provider's restored row. - if prefix_len - 1 > 0: - log_probs[reuser_idx, 0:prefix_len - 1] = log_probs[provider_idx, 0:prefix_len - 1] - - # prefix-last (position prefix_len-1): recompute with the reuser's label. - prefix_last_spec = prefix_last_spec_by_reuser.get(reuser_idx) - if prefix_last_spec is not None: - saved_logits_key = (reuser_idx, prefix_last_spec.target_2d_pos) - saved_provider_logits = ctx.prefix_last_logits_saved[saved_logits_key] # [1, V//tp] - reuser_label = torch.tensor( - [prefix_last_spec.label_value], dtype=torch.long, device=log_probs.device, - ) # [1] - log_probs[reuser_idx, prefix_len - 1] = vocab_parallel_log_probs_fn( - saved_provider_logits, reuser_label, - ).reshape(()) - else: - # suffix_len == 0: no prefix-last spec; column is masked downstream. - log_probs[reuser_idx, prefix_len - 1] = log_probs[provider_idx, prefix_len - 1] - - # entropy [0, prefix_len-1]: whole prefix copied (label-independent). - if entropy is not None: - entropy[reuser_idx, 0:prefix_len] = entropy[provider_idx, 0:prefix_len] - - restored_reusers += 1 - - if ctx.stats is not None: - ctx.stats.record_restore(restored_reusers) - return output - - -# ═══════════════════════════════════════════════════════════════ -# v080 restore 包装：NestedTensor → 2D left-pad → 复用 2D restore → 压回 -# ═══════════════════════════════════════════════════════════════ - - -def restore_via_2d_unfold_verl080( - output: dict, - vocab_parallel_log_probs_fn: Any, - vocab_parallel_entropy_fn: Any = None, -) -> dict: - """v080 restore 包装：NestedTensor → 2D left-pad → 复用 restore_reuser_prefix_columns_2d → 压回。 - - v080 物理裁剪后 reuser NestedTensor 行只含 suffix 区段，prefix 区段（含 - prefix-last）被物理删除。本函数在 forward_step 出口（context 仍激活、provider - prefix-last logits 已存于 ``ctx.prefix_last_logits_saved``）完成重组： - - 1. 展开裁剪后 NestedTensor 各行为完整 2D ``[B, L_max]``（reuser prefix 区段 - left-pad 0，尾部 right-pad 0 到 L_max） - 2. 复用 :func:`restore_reuser_prefix_columns_2d`：interior 整段从直接 - provider 的已恢复 2D 行 bulk 切片复制，prefix-last 用存的 logits + - ``index.label_value`` 重算 - 3. 按各 ``original_lengths`` 切片压回 NestedTensor (jagged) - - 列映射为 identity：left-pad 后 valid-content 的 0-based 偏移即 2D 列号， - ``target_2d_pos`` 直接当列索引用，无需 ``valid_indices`` / 列映射表。 - - Must be called inside ``prefix_sharing_runtime_context`` (reads - ``current_prefix_sharing_context``), after the vocab_logprobs patch has saved - provider prefix-last logits into ``ctx.prefix_last_logits_saved``. - - Args: - output: forward_step 返回的 output_dict，含 ``"log_probs"`` NestedTensor - （裁剪后 jagged），可选 ``"entropy"`` NestedTensor。**不含** tuple 外层 - （tuple 解包由调用方负责）。 - vocab_parallel_log_probs_fn: 用于 prefix-last logp 重算。 - vocab_parallel_entropy_fn: 可选，当前未直接使用（entropy 走复制路径—— - interior 和 prefix-last 都从 provider 复制，不重算）。 - - Returns: - ``output``（``log_probs``/``entropy`` 被替换为重组后的 NestedTensor）。 - """ - - ctx = current_prefix_sharing_context() - if ctx is None: - return output - plan = ctx.prefix_sharing_plan - # Guard on reuser presence, not on prefix_last_restore_indices: a batch - # whose reusers all have suffix_len == 0 has no prefix-last spec but still - # needs interior prefix columns restored. - if not plan.has_sharing: - return output - - log_probs_nested = output.get("log_probs") - if log_probs_nested is None or not is_nested_tensor(log_probs_nested): - return output - entropy_nested = output.get("entropy") - has_entropy = entropy_nested is not None and is_nested_tensor(entropy_nested) - - original_lengths = plan.original_lengths - input_keep_ranges = plan.input_keep_ranges - B = len(original_lengths) - if B == 0: - return output - L_max = max(original_lengths) - - # --- Step 1: 展开裁剪后 NestedTensor → 完整 2D [B, L_max] --- - log_probs_2d, entropy_2d = _unfold_trimmed_nested_to_2d( - log_probs_nested, - entropy_nested if has_entropy else None, - original_lengths, - input_keep_ranges, - L_max, - B, - ) - - # --- Step 2: 复用 restore_reuser_prefix_columns_2d --- - # build_kv 式区间拼接：interior 整段从直接 provider 的已恢复 2D 行切片， - # prefix-last 用 index.label_value + saved logits 重算。identity 列映射 - # （target_2d_pos 即 2D 列号，无 left padding）。 - output_2d: dict[str, Any] = {"log_probs": log_probs_2d} - if entropy_2d is not None: - output_2d["entropy"] = entropy_2d - output_2d = restore_reuser_prefix_columns_2d( - output_2d, - vocab_parallel_log_probs_fn, - vocab_parallel_entropy_fn, - ) - - # --- Step 3: 按各 original_lengths 压回 NestedTensor --- - output["log_probs"] = _fold_2d_to_nested(output_2d["log_probs"], original_lengths) - if entropy_2d is not None: - output["entropy"] = _fold_2d_to_nested(output_2d["entropy"], original_lengths) - - num_prefix_last = len(ctx.prefix_last_restore_indices) - print( - f"[PS][restore_verl080] unfolded B={B} L_max={L_max}, " - f"restored reusers={num_prefix_last} (prefix-last entries; " - f"interior handled by bulk slice)", - flush=True, - ) - return output - - -def _unfold_trimmed_nested_to_2d( - log_probs_nested: Any, - entropy_nested: Any, - original_lengths: list[int], - input_keep_ranges: list, - L_max: int, - B: int, -) -> tuple[Any, Any | None]: - """展开裁剪后 NestedTensor → 完整 2D [B, L_max]（reuser prefix left-pad 0）。 - - 裁剪后各行： - - provider (keep_start=0): 完整 [prefix | suffix]，长度 = original_lengths[i] - - reuser (keep_start=prefix_len>0): 仅 [suffix]，长度 = original_lengths[i]-prefix_len - - 展开后每行恢复成 [prefix_zeros | suffix]，再 right-pad 0 到 L_max。 - left-pad 的 zeros 不在 autograd 图里，但 restore 会覆盖 prefix 区段（interior - 复制 provider、prefix-last 重算），最终值在图里。right-pad 尾部在压回时丢弃。 - """ - import torch - - log_probs_offsets = log_probs_nested.offsets() - log_probs_values = log_probs_nested.values() - if entropy_nested is not None: - entropy_offsets = entropy_nested.offsets() - entropy_values = entropy_nested.values() - - log_probs_rows: list[Any] = [] - entropy_rows: list[Any] | None = [] if entropy_nested is not None else None - - for seq_idx in range(B): - orig_len = original_lengths[seq_idx] - prefix_len = input_keep_ranges[seq_idx][0] - - log_probs_suffix = log_probs_values[log_probs_offsets[seq_idx]:log_probs_offsets[seq_idx + 1]] - log_probs_rows.append(_build_padded_row(log_probs_suffix, prefix_len, orig_len, L_max)) - - if entropy_nested is not None: - entropy_suffix = entropy_values[entropy_offsets[seq_idx]:entropy_offsets[seq_idx + 1]] - entropy_rows.append(_build_padded_row(entropy_suffix, prefix_len, orig_len, L_max)) - - log_probs_2d = torch.stack(log_probs_rows, dim=0) - entropy_2d = torch.stack(entropy_rows, dim=0) if entropy_rows else None - return log_probs_2d, entropy_2d - - -def _build_padded_row( - suffix_data: Any, prefix_len: int, orig_len: int, L_max: int, -) -> Any: - """构造一行完整 2D ``[prefix_zeros | suffix]`` right-pad 0 到 L_max。""" - import torch - - device = suffix_data.device - dtype = suffix_data.dtype - tail_shape = tuple(suffix_data.shape[1:]) - pieces: list[Any] = [] - if prefix_len > 0: - pieces.append(torch.zeros((prefix_len,) + tail_shape, dtype=dtype, device=device)) - pieces.append(suffix_data) - row = torch.cat(pieces, dim=0) # [orig_len, ...] - if orig_len < L_max: - pad = torch.zeros((L_max - orig_len,) + tail_shape, dtype=dtype, device=device) - row = torch.cat([row, pad], dim=0) - return row - - -def _fold_2d_to_nested(tensor_2d: Any, original_lengths: list[int]) -> Any: - """完整 2D [B, L_max] → NestedTensor (jagged)，按各 original_lengths 切片。""" - import torch - - rows = [tensor_2d[seq_idx, :original_lengths[seq_idx]] for seq_idx in range(len(original_lengths))] - values = torch.cat(rows, dim=0) if rows else tensor_2d.reshape(0, *tensor_2d.shape[2:]) - offsets = torch.tensor( - [0] + [sum(original_lengths[: idx + 1]) for idx in range(len(original_lengths))], - dtype=torch.long, - device=tensor_2d.device, - ) - if hasattr(torch.nested, "nested_tensor_from_jagged"): - return torch.nested.nested_tensor_from_jagged(values, offsets) - return torch.nested.as_nested_tensor(rows, layout=torch.jagged) - - def build_prefix_sharing_micro_batch_verl080( engine_self: Any, batch: Any, diff --git a/prefix-sharing/prefix_sharing/integrations/verl_utils.py b/prefix-sharing/prefix_sharing/integrations/verl_utils.py index 8f3bc65a..f387ed07 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_utils.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_utils.py @@ -7,6 +7,7 @@ import torch from prefix_sharing.core.planner import PrefixSharingPlan +from prefix_sharing.integrations.context import current_prefix_sharing_context def read_ps_config_from_engine_config(engine_config: Any) -> Any | None: @@ -435,3 +436,292 @@ def extract_seq_from_dense_tensor( for row, indices in enumerate(valid_indices) ] return _copy_tensors_to_cpu_lists(tensor_slices) + + +def restore_reuser_prefix_columns_2d( + output: dict[str, Any], + vocab_parallel_log_probs_fn: Any, + vocab_parallel_entropy_fn: Any = None, +) -> dict[str, Any]: + """Restore reuser prefix columns in 2D space — build_kv-style slice + concat. + + Mirrors :meth:`TorchReferenceBackend.build_kv` + (``torch.cat([provider_kv[:prefix_len], own_suffix])``): instead of writing + each prefix token one scalar at a time, the whole prefix interval is sliced + off the **direct provider's already-restored 2D row** and only the single + prefix-last logprob is recomputed. + + Per ``reuser_idx`` with direct provider ``provider_idx = provider_index[reuser_idx]`` + and ``P = prefix_lens[reuser_idx]`` (columns are identity-mapped in the + unfolded 2D tensor, so ``target_2d_pos`` == column): + + - **interior ``[0, P-2]``**: ``log_probs[reuser_idx, 0:P-1] = log_probs[provider_idx, 0:P-1]`` + (bulk copy). Identical across the shared prefix (same logits + labels), + and ``provider_idx`` was restored earlier in the batch-order loop, so its + row already holds correct values — no per-position provider resolution + needed. + - **prefix-last ``P-1``**: recompute ``log_probs[reuser_idx, P-1]`` from the + saved provider logits + the reuser's own first-suffix label (differs from + the provider's). When the reuser has no suffix (``suffix_len == 0``) the + planner emits no prefix-last spec; that column is masked downstream, so + the provider's value is copied as a safe placeholder. + - **entropy ``[0, P-1]``**: ``entropy[reuser_idx, 0:P] = entropy[provider_idx, 0:P]`` + (whole prefix copied, including prefix-last — entropy is label-independent). + + Rows are visited in ``range(B)`` order so a provider is always restored + before any reuser that reads it (the same online-detector invariant + ``build_kv`` relies on). + + Args: + output: Output dict with ``log_probs`` [B, L] and optionally + ``entropy`` [B, L] in 2D space (unfolded from the trimmed + NestedTensor by :func:`restore_via_2d_unfold_verl080`). + vocab_parallel_log_probs_fn: ``logits [1, V//tp]``, ``label [1]`` → + scalar; used only for the prefix-last recompute. + vocab_parallel_entropy_fn: Retained for call-site compatibility; + unused (entropy is copied, never recomputed). + + Returns: + ``output`` with ``log_probs`` and ``entropy`` mutated in-place. + """ + + ctx = current_prefix_sharing_context() + if ctx is None: + return output + plan = ctx.prefix_sharing_plan + # Guard on reuser presence (not on prefix_last_restore_indices): a batch + # whose reusers all have suffix_len == 0 emits no prefix-last spec but still + # needs its interior prefix columns restored. + if not plan.has_sharing: + return output + + log_probs = output.get("log_probs") + if log_probs is None: + return output + entropy = output.get("entropy") + + provider_index = plan.provider_index + prefix_lens = plan.prefix_lens + + # reuser row → its prefix-last restore spec (one per reuser-with-suffix; + # interior positions have no spec — they are bulk-sliced below). + prefix_last_spec_by_reuser = { + spec.reuse_idx_in_batch: spec for spec in ctx.prefix_last_restore_indices + } + + restored_reusers = 0 + # Row 0 is always a provider (nothing precedes it to reuse), so start at 1. + # A reuser's provider always has a smaller batch index (online-detector + # invariant), so it is already restored when we reach reuser_idx. + for reuser_idx in range(1, len(prefix_lens)): + prefix_len = prefix_lens[reuser_idx] + if provider_index[reuser_idx] == reuser_idx or prefix_len <= 0: + continue # provider / non-reuser: row already complete + provider_idx = provider_index[reuser_idx] + + # interior [0, prefix_len-2]: bulk-copy from the provider's restored row. + if prefix_len - 1 > 0: + log_probs[reuser_idx, 0:prefix_len - 1] = log_probs[provider_idx, 0:prefix_len - 1] + + # prefix-last (position prefix_len-1): recompute with the reuser's label. + prefix_last_spec = prefix_last_spec_by_reuser.get(reuser_idx) + if prefix_last_spec is not None: + saved_logits_key = (reuser_idx, prefix_last_spec.target_2d_pos) + saved_provider_logits = ctx.prefix_last_logits_saved[saved_logits_key] # [1, V//tp] + reuser_label = torch.tensor( + [prefix_last_spec.label_value], dtype=torch.long, device=log_probs.device, + ) # [1] + log_probs[reuser_idx, prefix_len - 1] = vocab_parallel_log_probs_fn( + saved_provider_logits, reuser_label, + ).reshape(()) + else: + # suffix_len == 0: no prefix-last spec; column is masked downstream. + log_probs[reuser_idx, prefix_len - 1] = log_probs[provider_idx, prefix_len - 1] + + # entropy [0, prefix_len-1]: whole prefix copied (label-independent). + if entropy is not None: + entropy[reuser_idx, 0:prefix_len] = entropy[provider_idx, 0:prefix_len] + + restored_reusers += 1 + + if ctx.stats is not None: + ctx.stats.record_restore(restored_reusers) + return output + + +# ═══════════════════════════════════════════════════════════════ +# v080 restore 包装：NestedTensor → 2D left-pad → 复用 2D restore → 压回 +# ═══════════════════════════════════════════════════════════════ + + +def restore_via_2d_unfold_verl080( + output: dict, + vocab_parallel_log_probs_fn: Any, + vocab_parallel_entropy_fn: Any = None, +) -> dict: + """v080 restore 包装：NestedTensor → 2D left-pad → 复用 restore_reuser_prefix_columns_2d → 压回。 + + v080 物理裁剪后 reuser NestedTensor 行只含 suffix 区段，prefix 区段（含 + prefix-last）被物理删除。本函数在 forward_step 出口（context 仍激活、provider + prefix-last logits 已存于 ``ctx.prefix_last_logits_saved``）完成重组： + + 1. 展开裁剪后 NestedTensor 各行为完整 2D ``[B, L_max]``（reuser prefix 区段 + left-pad 0，尾部 right-pad 0 到 L_max） + 2. 复用 :func:`restore_reuser_prefix_columns_2d`：interior 整段从直接 + provider 的已恢复 2D 行 bulk 切片复制，prefix-last 用存的 logits + + ``index.label_value`` 重算 + 3. 按各 ``original_lengths`` 切片压回 NestedTensor (jagged) + + 列映射为 identity：left-pad 后 valid-content 的 0-based 偏移即 2D 列号， + ``target_2d_pos`` 直接当列索引用，无需 ``valid_indices`` / 列映射表。 + + Must be called inside ``prefix_sharing_runtime_context`` (reads + ``current_prefix_sharing_context``), after the vocab_logprobs patch has saved + provider prefix-last logits into ``ctx.prefix_last_logits_saved``. + + Args: + output: forward_step 返回的 output_dict，含 ``"log_probs"`` NestedTensor + （裁剪后 jagged），可选 ``"entropy"`` NestedTensor。**不含** tuple 外层 + （tuple 解包由调用方负责）。 + vocab_parallel_log_probs_fn: 用于 prefix-last logp 重算。 + vocab_parallel_entropy_fn: 可选，当前未直接使用（entropy 走复制路径—— + interior 和 prefix-last 都从 provider 复制，不重算）。 + + Returns: + ``output``（``log_probs``/``entropy`` 被替换为重组后的 NestedTensor）。 + """ + + ctx = current_prefix_sharing_context() + if ctx is None: + return output + plan = ctx.prefix_sharing_plan + # Guard on reuser presence, not on prefix_last_restore_indices: a batch + # whose reusers all have suffix_len == 0 has no prefix-last spec but still + # needs interior prefix columns restored. + if not plan.has_sharing: + return output + + log_probs_nested = output.get("log_probs") + if log_probs_nested is None or not is_nested_tensor(log_probs_nested): + return output + entropy_nested = output.get("entropy") + has_entropy = entropy_nested is not None and is_nested_tensor(entropy_nested) + + original_lengths = plan.original_lengths + input_keep_ranges = plan.input_keep_ranges + B = len(original_lengths) + if B == 0: + return output + L_max = max(original_lengths) + + # --- Step 1: 展开裁剪后 NestedTensor → 完整 2D [B, L_max] --- + log_probs_2d, entropy_2d = _unfold_trimmed_nested_to_2d( + log_probs_nested, + entropy_nested if has_entropy else None, + original_lengths, + input_keep_ranges, + L_max, + B, + ) + + # --- Step 2: 复用 restore_reuser_prefix_columns_2d --- + # build_kv 式区间拼接：interior 整段从直接 provider 的已恢复 2D 行切片， + # prefix-last 用 index.label_value + saved logits 重算。identity 列映射 + # （target_2d_pos 即 2D 列号，无 left padding）。 + output_2d: dict[str, Any] = {"log_probs": log_probs_2d} + if entropy_2d is not None: + output_2d["entropy"] = entropy_2d + output_2d = restore_reuser_prefix_columns_2d( + output_2d, + vocab_parallel_log_probs_fn, + vocab_parallel_entropy_fn, + ) + + # --- Step 3: 按各 original_lengths 压回 NestedTensor --- + output["log_probs"] = _fold_2d_to_nested(output_2d["log_probs"], original_lengths) + if entropy_2d is not None: + output["entropy"] = _fold_2d_to_nested(output_2d["entropy"], original_lengths) + + num_prefix_last = len(ctx.prefix_last_restore_indices) + print( + f"[PS][restore_verl080] unfolded B={B} L_max={L_max}, " + f"restored reusers={num_prefix_last} (prefix-last entries; " + f"interior handled by bulk slice)", + flush=True, + ) + return output + + +def _unfold_trimmed_nested_to_2d( + log_probs_nested: Any, + entropy_nested: Any, + original_lengths: list[int], + input_keep_ranges: list, + L_max: int, + B: int, +) -> tuple[Any, Any | None]: + """展开裁剪后 NestedTensor → 完整 2D [B, L_max]（reuser prefix left-pad 0）。 + + 裁剪后各行： + - provider (keep_start=0): 完整 [prefix | suffix]，长度 = original_lengths[i] + - reuser (keep_start=prefix_len>0): 仅 [suffix]，长度 = original_lengths[i]-prefix_len + + 展开后每行恢复成 [prefix_zeros | suffix]，再 right-pad 0 到 L_max。 + left-pad 的 zeros 不在 autograd 图里，但 restore 会覆盖 prefix 区段（interior + 复制 provider、prefix-last 重算），最终值在图里。right-pad 尾部在压回时丢弃。 + """ + log_probs_offsets = log_probs_nested.offsets() + log_probs_values = log_probs_nested.values() + if entropy_nested is not None: + entropy_offsets = entropy_nested.offsets() + entropy_values = entropy_nested.values() + + log_probs_rows: list[Any] = [] + entropy_rows: list[Any] | None = [] if entropy_nested is not None else None + + for seq_idx in range(B): + orig_len = original_lengths[seq_idx] + prefix_len = input_keep_ranges[seq_idx][0] + + log_probs_suffix = log_probs_values[log_probs_offsets[seq_idx]:log_probs_offsets[seq_idx + 1]] + log_probs_rows.append(_build_padded_row(log_probs_suffix, prefix_len, orig_len, L_max)) + + if entropy_nested is not None: + entropy_suffix = entropy_values[entropy_offsets[seq_idx]:entropy_offsets[seq_idx + 1]] + entropy_rows.append(_build_padded_row(entropy_suffix, prefix_len, orig_len, L_max)) + + log_probs_2d = torch.stack(log_probs_rows, dim=0) + entropy_2d = torch.stack(entropy_rows, dim=0) if entropy_rows else None + return log_probs_2d, entropy_2d + + +def _build_padded_row( + suffix_data: Any, prefix_len: int, orig_len: int, L_max: int, +) -> Any: + """构造一行完整 2D ``[prefix_zeros | suffix]`` right-pad 0 到 L_max。""" + device = suffix_data.device + dtype = suffix_data.dtype + tail_shape = tuple(suffix_data.shape[1:]) + pieces: list[Any] = [] + if prefix_len > 0: + pieces.append(torch.zeros((prefix_len,) + tail_shape, dtype=dtype, device=device)) + pieces.append(suffix_data) + row = torch.cat(pieces, dim=0) # [orig_len, ...] + if orig_len < L_max: + pad = torch.zeros((L_max - orig_len,) + tail_shape, dtype=dtype, device=device) + row = torch.cat([row, pad], dim=0) + return row + + +def _fold_2d_to_nested(tensor_2d: Any, original_lengths: list[int]) -> Any: + """完整 2D [B, L_max] → NestedTensor (jagged)，按各 original_lengths 切片。""" + rows = [tensor_2d[seq_idx, :original_lengths[seq_idx]] for seq_idx in range(len(original_lengths))] + values = torch.cat(rows, dim=0) if rows else tensor_2d.reshape(0, *tensor_2d.shape[2:]) + offsets = torch.tensor( + [0] + [sum(original_lengths[: idx + 1]) for idx in range(len(original_lengths))], + dtype=torch.long, + device=tensor_2d.device, + ) + if hasattr(torch.nested, "nested_tensor_from_jagged"): + return torch.nested.nested_tensor_from_jagged(values, offsets) + return torch.nested.as_nested_tensor(rows, layout=torch.jagged) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py index 51b2e476..7f9ed666 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py @@ -202,7 +202,7 @@ def patched_forward_step( # 解包处理 output_dict 再重包。restore_via_2d_unfold_verl080 内部 # 会检查 context / restore_indices，无 restore 需求时 early return。 if ps_state is not None: - from prefix_sharing.integrations.verl_mcore import restore_via_2d_unfold_verl080 + from prefix_sharing.integrations.verl_utils import restore_via_2d_unfold_verl080 from prefix_sharing.integrations.context import current_prefix_sharing_context from verl.utils.megatron.tensor_parallel import ( vocab_parallel_entropy, diff --git a/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py b/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py index e6285ead..b736002c 100644 --- a/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py +++ b/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py @@ -18,7 +18,7 @@ from prefix_sharing.integrations.megatron_runtime import prefix_attention from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState from prefix_sharing.integrations.parallel_info import MegatronParallelInfo -from prefix_sharing.integrations.verl_mcore import restore_reuser_prefix_columns_2d +from prefix_sharing.integrations.verl_utils import restore_reuser_prefix_columns_2d def _make_state() -> tuple[PrefixSharingRuntimeState, list]: diff --git a/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py b/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py index 8a0e2ed2..a01138ce 100644 --- a/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py +++ b/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py @@ -23,10 +23,8 @@ megatron = pytest.importorskip("megatron") torch = pytest.importorskip("torch") -from prefix_sharing.integrations.verl_mcore import ( - build_prefix_sharing_micro_batch_verl080, - restore_via_2d_unfold_verl080, -) +from prefix_sharing.integrations.verl_mcore import build_prefix_sharing_micro_batch_verl080 +from prefix_sharing.integrations.verl_utils import restore_via_2d_unfold_verl080 # ═══════════════════════════════════════ diff --git a/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py b/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py index 38e2dd09..07647f2c 100644 --- a/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py +++ b/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py @@ -22,8 +22,8 @@ from prefix_sharing.core.planner import PrefixSharingPlanner from prefix_sharing.integrations.context import prefix_sharing_runtime_context from prefix_sharing.integrations.parallel_info import MegatronParallelInfo -from prefix_sharing.integrations.verl_mcore import ( - PrefixSharingRuntimeState, +from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState +from prefix_sharing.integrations.verl_utils import ( _fold_2d_to_nested, _unfold_trimmed_nested_to_2d, restore_via_2d_unfold_verl080, From f55adc72be80df9cb5e9161a6096dfb370c7643d Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Sun, 16 Aug 2026 18:46:50 +0800 Subject: [PATCH 18/19] refactor(fsdp): clarify context lifecycle and attention path Rename the private context attribute from `_ps_ctx` to `_prefix_sharing_context` and the cleanup callback to `_cleanup_prefix_sharing_context` so the lifecycle is self-documenting. Add STEP comments to the FSDP attention patch and forward_step to make the prepare/execute/restore phases explicit. Reorganize imports and simplify local variable names while preserving the original behavior. Co-authored-by: Cursor --- .../prefix_sharing/integrations/context.py | 6 +- .../prefix_sharing/integrations/verl_fsdp.py | 17 ++++ .../setup/patches/verl080_fsdp/attention.py | 44 ++++++----- .../patches/verl080_fsdp/forward_step.py | 78 +++++++++++-------- 4 files changed, 92 insertions(+), 53 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/context.py b/prefix-sharing/prefix_sharing/integrations/context.py index b28498dc..b25f7937 100644 --- a/prefix-sharing/prefix_sharing/integrations/context.py +++ b/prefix-sharing/prefix_sharing/integrations/context.py @@ -151,7 +151,7 @@ def create_prefix_sharing_context( the ContextVar is left set until the caller invokes the returned cleanup function. This is required for activation‑checkpointing compatibility: AC recompute runs inside ``backward()`` and reads the PS context from - ``module._ps_ctx`` (set independently by the caller), while the store must + ``module._prefix_sharing_context`` (set independently by the caller), while the store must still contain the per‑layer KV populated during the first forward. Returns: @@ -164,12 +164,12 @@ def create_prefix_sharing_context( ctx = PrefixSharingRuntimeContext(prefix_sharing_runtime_state, store) ctxvar_token = _current_context.set(ctx) - def cleanup() -> None: + def cleanup_context() -> None: _current_context.reset(ctxvar_token) _log_prefix_sharing_audit(ctx) ctx.store.close() - return ctx, cleanup + return ctx, cleanup_context def _log_prefix_sharing_audit(ctx: PrefixSharingRuntimeContext) -> None: diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py index d1825c9b..77fdc385 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py @@ -54,6 +54,11 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any, from prefix_sharing.tools.perf_profiler import PerfProfiler profiler = PerfProfiler.current() # [PS-perf] end —————————————————————————————————————— + + ######################################################### + # handle THD format input + ######################################################### + if query.shape[0] == 1 and key.shape[0] == 1 and value.shape[0] == 1: packed_query = query.squeeze(0) packed_key = key.squeeze(0) @@ -70,6 +75,10 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any, if query.shape[:2] != key.shape[:2] or query.shape[:2] != value.shape[:2]: raise RuntimeError("query, key, and value must share dense batch/sequence dimensions") + ######################################################### + # handle BSHD format input + ######################################################### + plan = ctx.prefix_sharing_plan # [PS-perf] start — attn.pack ——————————————————————— _per_layer_ok = profiler is not None and getattr(profiler, "per_layer_enabled", False) @@ -374,6 +383,10 @@ def _run_packed_attention_runtime( if _per_layer_ok: profiler.start_phase(f"attn.kv.l{layer_id}") # [PS-perf] end ———————————————————————————————————————— + + ######################################################### + # STEP 1: build key and value tensors + ######################################################### expanded_key, expanded_value = ctx.attention_backend.build_kv( packed_key, packed_value, @@ -398,6 +411,10 @@ def _run_packed_attention_runtime( dump_expanded_kv_on(layer_number, expanded_key, expanded_value, num_layers) # ##### [PS-diag] end ##### + ######################################################### + # STEP 2: compute attention + ######################################################### + if profiler is not None: profiler.start_phase(PerfProfiler.PHASE_ATTN_COMPUTE) # [PS-perf] start — per-layer compute —————————————————— diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py index 8b99edf8..1cf1a5d7 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py @@ -18,6 +18,8 @@ import os from typing import Any +from prefix_sharing.integrations.context import _current_context +from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime from prefix_sharing.tools.perf_profiler import PerfProfiler _SUPPORTED_ATTENTIONS = { @@ -86,14 +88,19 @@ def create_attention_wrapper(original_fn: Any) -> Any: def patched_attention(module: Any, query: Any, key: Any, value: Any, attention_mask: Any, *args: Any, **kwargs: Any) -> Any: + + ######################################################### + # STEP 1: get prefix sharing context + ######################################################### + # 优先 module 属性（AC recompute 兼容），回退 ContextVar（Megatron 等路径） - ctx = getattr(module, '_ps_ctx', None) - if ctx is None: + prefix_sharing_context = getattr(module, '_prefix_sharing_context', None) + if prefix_sharing_context is None: from prefix_sharing.integrations.context import current_prefix_sharing_context - ctx = current_prefix_sharing_context() + prefix_sharing_context = current_prefix_sharing_context() # ── OFF path: no prefix sharing context → transparent passthrough ── - if ctx is None: + if prefix_sharing_context is None: # [PS-perf] start — OFF attention timing (cross-layer + per-layer) — profiler = PerfProfiler.current() _per_layer_ok = profiler is not None and getattr(profiler, "per_layer_enabled", False) @@ -119,22 +126,23 @@ def patched_attention(module: Any, query: Any, key: Any, value: Any, # ##### [PS-diag] end ##### return result - # ── ON path: route through PrefixSharing attention runtime ── - from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime - from prefix_sharing.integrations.context import _current_context + ######################################################### + # STEP 2: call PrefixSharing attention runtime + ######################################################### + # ── ON path: route through PrefixSharing attention runtime ── layer_id = int(getattr(module, "layer_idx", 0) or 0) _num_layers = _resolve_num_layers(module) - runtime = PrefixSharingFSDPAttentionRuntime(layer_id=layer_id, num_layers=_num_layers) + attention_runtime = PrefixSharingFSDPAttentionRuntime(layer_id=layer_id, num_layers=_num_layers) - # HF attention interface expects [B, H, L, D]; runtime works in [B, L, H, D] - query_ld = query.transpose(1, 2) - key_ld = key.transpose(1, 2) - value_ld = value.transpose(1, 2) + # HF attention interface expects [B, H, L, D]; attention_runtime works in [B, L, H, D] + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) - # runtime.forward() 内部读 ContextVar；AC recompute 时 ContextVar - # 可能已过期，但 ctx 来自 module._ps_ctx 仍然有效。临时注入 ContextVar。 - _ctxvar_token = _current_context.set(ctx) + # attention_runtime.forward() 内部读 ContextVar；AC recompute 时 ContextVar + # 可能已过期，但 prefix_sharing_context 来自 module._prefix_sharing_context 仍然有效。临时注入 ContextVar。 + _ctxvar_token = _current_context.set(prefix_sharing_context) # [PS-perf] start — ON attention timing (attn.on = pack+kv+comp+unpack) — profiler = PerfProfiler.current() _per_layer_ok = profiler is not None and getattr(profiler, "per_layer_enabled", False) @@ -143,7 +151,7 @@ def patched_attention(module: Any, query: Any, key: Any, value: Any, if _per_layer_ok: profiler.start_phase(f"attn.on.l{layer_id}") try: - output_ld = runtime.forward(None, query_ld, key_ld, value_ld) + output = attention_runtime.forward(None, query, key, value) finally: if _per_layer_ok: _on_elapsed = profiler.stop_phase(f"attn.on.l{layer_id}") @@ -155,8 +163,8 @@ def patched_attention(module: Any, query: Any, key: Any, value: Any, # ##### [PS-diag] ON attn output dump（context 激活 = PS 路径） ##### if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: - _dump_attn_output(output_ld, module) - return output_ld, None + _dump_attn_output(output, module) + return output, None return patched_attention diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py index d9aee013..e3c45820 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py @@ -16,6 +16,9 @@ from prefix_sharing.integrations.context import create_prefix_sharing_context from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime from prefix_sharing.integrations.verl_fsdp import prepare_for_prefix_sharing_fsdp +from prefix_sharing.integrations.verl_fsdp import restore_prefix_sharing_outputs_2d +from prefix_sharing.integrations.verl_utils import restore_via_2d_unfold_verl080 +from prefix_sharing.integrations.verl_utils import is_nested_tensor from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope @@ -45,11 +48,7 @@ def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forwar # path records only the model-forward phase. profiler = ProfilerScope.current() if profiler is not None: - forward_phase = ( - PerfProfiler.PHASE_FORWARD_OLD - if forward_only - else PerfProfiler.PHASE_FORWARD - ) + forward_phase = PerfProfiler.PHASE_FORWARD_OLD if forward_only else PerfProfiler.PHASE_FORWARD profiler.start_phase(forward_phase) # Without diagnostics this path must delegate directly to verl. @@ -147,6 +146,11 @@ def _forward_step_with_engine_prepare( ps_config: Any, model_config: Any, ) -> Any: + + ######################################################### + # STEP 1: pre-processing inputs for PrefixSharing + ######################################################### + profiler = ProfilerScope.current() if profiler is not None: profiler.start_phase(PerfProfiler.PHASE_PLAN) @@ -157,14 +161,12 @@ def _forward_step_with_engine_prepare( if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: _dump_full_input_ids_only(micro_batch, "train") - ######################################################### - # STEP 1: pre-processing inputs for PrefixSharing - ######################################################### micro_batch_modified, prefix_sharing_runtime_state = prepare_for_prefix_sharing_fsdp( micro_batch, ps_config, model_config=model_config, ) + if profiler is not None: profiler.stop_phase(PerfProfiler.PHASE_PLAN) # Detect, plan, and trim on CPU. @@ -181,6 +183,10 @@ def _forward_step_with_engine_prepare( diagnostic_tag, ) + ######################################################### + # STEP 2: pre-processing inputs for the model + ######################################################### + # Read layer count for per-layer diagnostic dumps. _diag_num_layers = int(getattr( getattr(getattr(self, "module", None), "config", None), @@ -195,16 +201,20 @@ def _forward_step_with_engine_prepare( if autocast_dtype == torch.float32 else torch.autocast(device_type=_read_device_name(), dtype=autocast_dtype) ) - # ── Create PS context with manual lifecycle (survives backward for AC) ── - ctx, ctx_cleanup = create_prefix_sharing_context(prefix_sharing_runtime_state) - # Set _ps_ctx on every attention module so the attention patch reads + ######################################################### + # STEP 3: create PrefixSharing runtime context + ######################################################### + + prefix_sharing_context, cleanup_prefix_sharing_context = create_prefix_sharing_context(prefix_sharing_runtime_state) + + # Set _prefix_sharing_context on every attention module so the attention patch reads # the context from the module itself rather than ContextVar (compatible # with activation-checkpointing recompute, which bypasses the context # manager that set the ContextVar). for attention_module in self.module.modules(): if hasattr(attention_module, "layer_idx") and hasattr(attention_module, "q_proj"): - attention_module._ps_ctx = ctx + attention_module._prefix_sharing_context = prefix_sharing_context # Register diagnostic gradient hooks when the diagnostic dump is enabled. _register_grad_dump_hooks(self.module, forward_only) @@ -212,15 +222,15 @@ def _forward_step_with_engine_prepare( # Attach cleanup callback so the forward_backward_batch wrapper can release # PrefixSharing state after backward. The root full-backward hook is # unreliable here because ``self.module`` returns a CausalLMOutput dataclass. - self.module._ps_ctx_cleanup = ctx_cleanup + self.module._cleanup_prefix_sharing_context = cleanup_prefix_sharing_context + + ######################################################### + # STEP 4: call the model with PrefixSharing+FSDP attention runtime + ######################################################### with autocast_ctx: if profiler is not None: - forward_phase = ( - PerfProfiler.PHASE_FORWARD_OLD - if forward_only - else PerfProfiler.PHASE_FORWARD - ) + forward_phase = PerfProfiler.PHASE_FORWARD_OLD if forward_only else PerfProfiler.PHASE_FORWARD profiler.start_phase(forward_phase) raw_output = self.module(**model_inputs, use_cache=False) if profiler is not None: @@ -231,6 +241,10 @@ def _forward_step_with_engine_prepare( dump_raw_logits_verl080(raw_output, dp_aware=_get_dp_size() > 1) + ######################################################### + # STEP 5: post-processing outputs for the model + ######################################################### + _save_prefix_last_logits_from_raw_output(raw_output) model_output = self.prepare_model_outputs( output=raw_output, @@ -239,6 +253,10 @@ def _forward_step_with_engine_prepare( logits_processor_func=loss_function, ) + ######################################################### + # STEP 6: post-processing outputs for PrefixSharing + ######################################################### + if profiler is not None: profiler.start_phase(PerfProfiler.PHASE_RESTORE) model_output = _restore_engine_model_output(model_output) @@ -254,6 +272,10 @@ def _forward_step_with_engine_prepare( diagnostic_tag, ) + ######################################################### + # STEP 7: compute loss & metrics + ######################################################### + if loss_function is not None: if profiler is not None: profiler.start_phase(PerfProfiler.PHASE_LOSS) @@ -357,11 +379,7 @@ def _call_original_like_engine(self: Any, micro_batch: Any, loss_function: Any, profiler = ProfilerScope.current() if profiler is not None: - forward_phase = ( - profiler.PHASE_FORWARD_OLD - if forward_only - else profiler.PHASE_FORWARD - ) + forward_phase = profiler.PHASE_FORWARD_OLD if forward_only else profiler.PHASE_FORWARD profiler.start_phase(forward_phase) with autocast_ctx: @@ -433,10 +451,6 @@ def _restore_engine_model_output(model_output: dict[str, Any]) -> dict[str, Any] except Exception: entropy_from_logits = None - from prefix_sharing.integrations.verl_mcore import restore_via_2d_unfold_verl080 - from prefix_sharing.integrations.verl_utils import is_nested_tensor - from prefix_sharing.integrations.verl_fsdp import restore_prefix_sharing_outputs_2d - restored = restore_via_2d_unfold_verl080( model_output, logprobs_from_logits, @@ -555,10 +569,10 @@ def wrapped(self: Any, data: Any, loss_function: Any, forward_only: bool = False # 之前该清理被误关在 PREFIX_SHARING_DIAG_DUMP 条件块内，导致正常训练时 # audit 日志不输出、KV store 不 close（多步训练存在内存累积风险）。 if not forward_only: - ctx_cleanup = getattr(self.module, "_ps_ctx_cleanup", None) - if ctx_cleanup is not None: - ctx_cleanup() - delattr(self.module, "_ps_ctx_cleanup") + cleanup_prefix_sharing_context = getattr(self.module, "_cleanup_prefix_sharing_context", None) + if cleanup_prefix_sharing_context is not None: + cleanup_prefix_sharing_context() + delattr(self.module, "_cleanup_prefix_sharing_context") # 诊断 dump 专用：dump weight gradients + 清理 per-layer attention grad hooks。 if not forward_only and os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: @@ -575,7 +589,7 @@ def wrapped(self: Any, data: Any, loss_function: Any, forward_only: bool = False # Remove per-layer attention gradient hooks. for module in self.module.modules(): try: - del module._ps_ctx + del module._prefix_sharing_context except AttributeError: pass From 2ce771acdd883e06359709984c12bd68f4c6c8d2 Mon Sep 17 00:00:00 2001 From: Jackie2049 Date: Mon, 17 Aug 2026 18:46:06 +0800 Subject: [PATCH 19/19] docs(impr-refactor): add FSDP activation-checkpointing runtime context lifecycle note Co-authored-by: Cursor --- docs/developer-docs/impr-refactor.md | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/docs/developer-docs/impr-refactor.md b/docs/developer-docs/impr-refactor.md index b88f46b9..10d31146 100644 --- a/docs/developer-docs/impr-refactor.md +++ b/docs/developer-docs/impr-refactor.md @@ -2280,25 +2280,32 @@ PYTHONPATH=prefix-sharing pytest -q \ - 需要固定数据和随机种子，形成可复现的精度对齐脚本。 - optional GPU/NPU 环境测试不能作为每个本地 PR 的硬门槛，但 release 前必须复跑。 -### 6.4 Patch 与 import hook +### 6.4 FSDP activation-checkpointing runtime context 生命周期 + +- `origin/open-source` 已存在 FSDP PrefixSharing runtime context 生命周期缺口：`forward_only=True` 没有 backward 触发清理；普通训练在未开启诊断 dump 时也可能遗留 module 上的 prefix-sharing context。该问题起源于 activation-checkpointing 兼容改造，不是当前第三波重构引入。 +- 当前第三波重构以 `origin/open-source` 为行为基线闭环，不在本 PR 内修复此既有缺口；但所有新增或修改的测试不得掩盖它，也不得将其误归因为重构回归。 +- 后续应以独立修复 PR 处理：统一 ContextVar、store 与 attention module binding 的所有权和幂等 cleanup；覆盖普通训练、`forward_only`、activation-checkpointing recompute、异常路径及连续 micro-batch。 +- 在该独立修复完成前，FSDP 精度验证应固定 replay fixture，并明确记录所覆盖的执行模式；不得把单 micro-batch 的 ON/OFF 对齐外推为上述生命周期场景全部已验证。 + +### 6.5 Patch 与 import hook - import hook 是否长期保留，需要等社区对 monkey patch 方式的反馈后再定。 - 如果正式合入 verl，显式调用路径可能替代外部包 import auto patch。 - 当前阶段必须保留 import auto patch，因为脚本化训练仍依赖 `VERL_USE_EXTERNAL_MODULES=prefix_sharing`。 -### 6.5 HybridAttention / Gated DeltaNet +### 6.6 HybridAttention / Gated DeltaNet - 当前重构清理 Qwen3.5/GDN 专门化代码，不代表永久放弃 HybridAttention。 - 后续需等待训练引擎侧真实接口稳定，再重新设计 activation/cache_param store。 - 未来重新引入时，应以实际 mixer 类型命名，避免把 DeltaNet 泛化成所有 linear attention。 -### 6.6 Megatron / MCore / NPU +### 6.7 Megatron / MCore / NPU - MCore path 保留 advanced/internal 定位。 - NPU/MindSpeed/Megatron-Bridge 后续可继续支持，但不能阻塞 FSDP-first 开源主线。 - 若后续重新提高 Megatron 优先级，需要单独补兼容矩阵、真实环境测试和文档。 -### 6.7 tools 与 diagnostics +### 6.8 tools 与 diagnostics - tools 清理需要逐项判断，不能批量删除。 - 保留工具必须补用途说明。