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

[Jackie2049]

Overview

This PR completes the second FSDP-focused refactor wave after the setup lifecycle consolidation. It makes the FSDP patch set easier to inspect, restores runtime cleanup after every training backward pass, removes avoidable work from the FSDP forward-step hot path, and relocates shared restore helpers so that framework-specific modules stay focused.

Design

1. Manifest as the integration contract: the FSDP patch-set manifest groups PrefixSharing entrypoints separately from diagnostic and performance instrumentation, and records the wrapper order that composes on the same verl target.
2. Cleanup follows runtime lifetime: PrefixSharing context cleanup is performed after every non-forward-only backward pass, independent of whether diagnostic dumping is enabled.
3. Keep the hot path explicit: stable imports and configuration readers are resolved at module scope; the forward wrapper only constructs the PrefixSharing configuration from framework-level input.
4. Keep framework-specific modules focused: restore helpers that are shared between FSDP and Megatron paths live in verl_utils, while Megatron-specific micro-batch building stays in verl_mcore.

Key Changes

• Clarify the FSDP patch manifest: reorganized PATCH_SET comments and descriptions, documented each wrapper purpose, and added an integration assertion for the profiler then diagnostic-wrapper order.
• Fix context cleanup outside diagnostic mode: moved the PrefixSharing context cleanup out of the diagnostic-only block after backward. This restores audit logging and closes the KV store during ordinary training, preventing runtime-state accumulation across steps.
• Simplify FSDP forward-step dependencies: hoisted stable imports, used the shared verl configuration reader, and renamed the PrefixGrouper bridge helper for clearer ownership.
• Relocate restore helpers from verl_mcore to verl_utils: moved restore_reuser_prefix_columns_2d, restore_via_2d_unfold_verl080 and their helpers so that verl_mcore only contains Megatron-specific micro-batch building logic. Updated all import sites (FSDP patch, Megatron patch, tests) accordingly.
• Clarify FSDP context lifecycle and attention path: renamed the private context attribute from _ps_ctx to _prefix_sharing_context and the cleanup callback to _cleanup_prefix_sharing_context, added explicit STEP comments in the FSDP attention patch and forward_step to mark the prepare/execute/restore phases, and simplified local variable names while preserving behavior.

Test Results

Local CPU regression

PYTHONPATH=. PYTHONPYCACHEPREFIX=/private/tmp/prefix-sharing-pycache python3 -m pytest -q tests/unit_test tests/integrated_test tests/system_test -k 'not test_auto_activation_ and not test_fixed_rollout_skips_validation_and_replays_training_output'

Result: 290 passed, 31 skipped, 4 deselected.

4090 verl 0.8.0 FSDP environment

PYTHONPATH=. PYTHONPYCACHEPREFIX=/tmp/prefix-sharing-fsdp-pycache python -m pytest -q tests/integrated_test/test_patch_integrations.py tests/unit_test/test_version_detector.py tests/unit_test/test_verl080_migration.py

Result: 49 passed, 9 warnings.

Static checks

ruff check --select F,E9 prefix_sharing/setup/patches/verl080_fsdp tests/integrated_test/test_patch_integrations.py

Result: passed. git diff --check also passed.

Remaining Items and Potential Impact

• The FSDP attention registry patch still lives outside the PatchHandle rollback ledger because it patches a registry entry rather than a normal module attribute. It should be moved into a reversible installer primitive in a later refactor.
• The temporary builtins import hook remains necessary for lazy-loaded training modules. Coverage for multiple pending patches targeting one module, timeout recovery, and repeated install/disable sequences should continue to be maintained.
• No end-to-end multi-step device training was added specifically for this PR. The cleanup fix was covered by setup-level regression tests; a future device smoke run should confirm audit output and KV-store release across consecutive training steps.