refactor: FSDP patch lifecycle, naming clarity, and helper relocation

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_step_without_engine_prepare()`
 
 问题：
 
-- `forward_prefix_sharing_fsdp_micro_batch()` 更像测试/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 语义复杂，需要更强测试和更小函数。
@@ -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 清理需要逐项判断，不能批量删除。
 - 保留工具必须补用途说明。

prefix-sharing/prefix_sharing/backends/packed_layout.py

@@ -57,34 +57,35 @@ 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")
         if not 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,
                 )
             )
             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(

prefix-sharing/prefix_sharing/core/config.py

@@ -103,136 +103,134 @@ 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] Current model_type '{model_type}' is not supported in this phase. "
-                f"Phase 1 only supports model_type='text_only_causal_lm' (text-only causal language model). "
-                f"Please use a supported model type or disable 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 does not currently support ulysses_sequence_parallel_size={ulysses_sp_size}. "
-                    "Please disable Ulysses SP or wait for a dedicated adaptation."
+                    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 does not currently support use_fused_kernels=True. "
-                    "Please disable fused kernels or wait for a dedicated adaptation."
+                    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} is invalid. "
-                f"pipeline_model_parallel_size must be >= 1. "
-                f"Please set a valid physical PP size or disable 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} is not supported in this phase. "
-                f"Only physical pipeline parallel is supported; virtual pipeline parallel is not. "
-                f"Please disable virtual PP or disable 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 is not supported in this phase. "
-                "Only physical pipeline parallel is supported; virtual pipeline parallel is not. "
-                "Please disable virtual PP or disable 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} is not supported in this phase. "
-                f"Phase 1 only supports context_parallel_size=1 (no context parallelism). "
-                f"Please set CP size to 1 or disable 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 is not supported in this phase. "
-                "Phase 1 requires rope fusion to be disabled (apply_rope_fusion=False). "
-                "Please update the configuration or disable 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 is not supported in this phase. "
-                "Phase 1 requires fused QKV rope to be disabled (fused_single_qkv_rope=False). "
-                "Please update the configuration or disable prefix sharing."
+                f"fused_single_qkv_rope={fused_qkv_rope} is not supported. "
+                "Supported fused_single_qkv_rope: False"
             )
 
     def validate_for_engine(
         self,
         use_remove_padding: bool = True,
         integrate_mode: str = "verl_megatron_actor",
     ) -> None:
-        """Validate phase-1 constraints for verl engine architecture (verl 0.8.0+).
+        """Validate phase-1 constraints for the verl engine (verl 0.8.0+).
 
-        Unlike validate(), this method reads from engine_config rather than
-        model_config. Used in setup/patches for forward_step patching, where
-        only engine_config (self.engine_config) is available, not the 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 validation
+        # 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"}:
@@ -249,10 +247,10 @@ def validate_for_engine(
         if self.min_group_size < 2:
             raise PrefixSharingConfigError("min_group_size must be >= 2")
 
-        # THD packed layout requires use_remove_padding
+        # THD packed layout requires use_remove_padding.
         if not use_remove_padding:
             raise PrefixSharingConfigError(
                 "[Config Error] Phase 1 THD path requires use_remove_padding=True. "
-                "The BSHD path (use_remove_padding=False) is not yet supported in the current patch set. "
-                "Please enable use_remove_padding or use the BSHD-specific patch set."
+                "BSHD (use_remove_padding=False) is not supported by the current "
+                "patch set. Enable use_remove_padding or use a BSHD-specific patch set."
             )

prefix-sharing/prefix_sharing/integrations/__init__.py

@@ -9,15 +9,15 @@
 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,
-    build_prefix_sharing_micro_batch_fsdp,
-    forward_prefix_sharing_fsdp_micro_batch,
+    prepare_for_prefix_sharing_fsdp,
+    forward_step_without_engine_prepare,
     restore_prefix_sharing_outputs_2d,
 )
 from prefix_sharing.integrations.megatron_runtime import (
@@ -38,7 +38,7 @@
     "prefix_attention",
     "get_megatron_parallel_info",
     "PrefixSharingFSDPAttentionRuntime",
-    "build_prefix_sharing_micro_batch_fsdp",
-    "forward_prefix_sharing_fsdp_micro_batch",
+    "prepare_for_prefix_sharing_fsdp",
+    "forward_step_without_engine_prepare",
     "restore_prefix_sharing_outputs_2d",
 ]

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: