执行摘要

• 一句话：支持 NCCL WORLD 进程组的销毁与重新加载
• 推荐动作：建议深入阅读 slime/utils/reloadable_process_group.py 中的实现，特别是 _destroy_default_nccl_process_group 和 _reload_default_process_group 的设计模式：利用 Gloo 组作为同步和过渡、使用 PrefixStore 保证 store 无冲突、通过 generation 字段保证唯一性。这些是分布式系统中重建通信组时的常用技巧。同时关注环境变量逃生口的设计。

功能与动机

支持训练过程中释放 actor 组后重建通信环境，避免因 WORLD 组损坏需要重启所有进程。PR 2180 引入的 --release-train 功能依赖此基础。此外，slime/backends/megatron_utils/actor.py 的注释明确指出 'Destroying and recreating WORLD invalidates raw dist.group.WORLD references'，需要通过逃生口控制。

实现拆解

1. 新增 slime/utils/reloadable_process_group.py：定义 _DefaultProcessGroupState 数据类捕获 backend、store、rank、world_size 等信息；提供 register_default_process_group 函数供 Actor 初始化时调用，将状态存入进程级字典 default_process_group_states。
2. 实现销毁与重建：_destroy_default_nccl_process_group 通过 Gloo 组同步后调用 dist.destroy_process_group()，然后使用 PrefixStore 创建新命名空间的临时 Gloo WORLD 组；_reload_default_process_group 销毁临时 Gloo 组，恢复原始 backend（如 NCCL）的 WORLD 组，并重新初始化规范 Gloo 组。两个函数均通过 state.generation 递增确保每次使用唯一的 store 前缀。
3. 修改 slime/utils/distributed_utils.py：新增 set_gloo_group 函数允许在 reload 过程中替换全局缓存的 Gloo 组；修改 init_gloo_group 使用 getattr(dist, 'old_new_group', dist.new_group) 绕过猴子补丁，保证规范 Gloo 组不受 ReloadableProcessGroup 包装影响。
4. 集成到 Actor 初始化：在 slime/backends/megatron_utils/actor.py 的 init 方法中，于 monkey_patch_torch_dist() 和 super().init() 之后，调用 register_default_process_group，受 SLIME_DESTROY_WORLD_PROCESS_GROUP 环境变量控制（默认启用）。
5. 修改 slime/ray/rollout.py：仅调整注释缩进，无逻辑变更。
6. 添加单元测试 tests/test_reloadable_process_group_world.py：使用 pytest monkeypatch 模拟分布式调用，验证 register_default_process_group 正确捕获状态（test_register_default_process_group_captures_rendezvous_state）、destroy_process_groups/reload_process_groups 的事件顺序正确（test_world_and_subgroups_follow_destroy_reload_order）、以及未注册 WORLD 时不会误破坏（test_unregistered_world_preserves_subgroup_only_behavior）。
7. 更新 CI 配置：.github/workflows/pr-test.yml 添加新测试到 CPU-only 测试矩阵，并在依赖安装中增加 psutil。

关键文件：

• slime/utils/reloadable_process_group.py（模块 通信组管理；类别 source；类型 core-logic；符号 _DefaultProcessGroupState, register_default_process_group, _uses_nccl, _new_default_process_group）: 新增核心模块，实现了 WORLD 组的注册、销毁、重建以及子组管理。
• tests/test_reloadable_process_group_world.py（模块 测试覆盖；类别 test；类型 test-coverage；符号 test_register_default_process_group_captures_rendezvous_state, test_world_and_subgroups_follow_destroy_reload_order, test_unregistered_world_preserves_subgroup_only_behavior）: 新增单元测试覆盖三种关键场景：正确捕获状态、destroy/reload 事件顺序、未注册时的安全行为。
• slime/utils/distributed_utils.py（模块 分布式工具；类别 source；类型 core-logic；符号 set_gloo_group, init_gloo_group）: 添加 set_gloo_group 函数，修改 init_gloo_group 使用原始 new_group，防止猴子补丁影响规范 Gloo 组。
• slime/backends/megatron_utils/actor.py（模块 Actor；类别 source；类型 dependency-wiring；符号 init）: 在 actor 初始化中注册默认进程组，并通过环境变量提供开关。
• slime/ray/rollout.py（模块 Rollout；类别 source；类型 configuration）: 调整注释缩进，无逻辑变更。
• .github/workflows/pr-test.yml（模块 CI配置；类别 infra；类型 infrastructure）: 添加新测试到 CPU 测试矩阵，并安装 psutil 依赖。
• .github/workflows/pr-test.yml.j2（模块 CI模板；类别 infra；类型 infrastructure）: 模板文件对应调整。

关键符号：_DefaultProcessGroupState, register_default_process_group, _uses_nccl, _new_default_process_group, _destroy_default_nccl_process_group, _reload_default_process_group, set_gloo_group, init_gloo_group (修改), MegatronTrainRayActor.init (修改), destroy_process_groups (现有), reload_process_groups (现有)

关键源码片段

slime/utils/reloadable_process_group.py

新增核心模块，实现了 WORLD 组的注册、销毁、重建以及子组管理。

@dataclass
class _DefaultProcessGroupState:
    '''Captures the rendezvous state of the original WORLD group for later restore.'''
    backend: str
    timeout: timedelta
    store: Any # Keep reference to rank-0 TCPStore
    rank: int
    world_size: int
    generation: int = 0
    nccl_world_destroyed: bool = False
​
​
def register_default_process_group(timeout: timedelta) -> None:
    '''Snapshot the current WORLD group state.
​
    Stores a reference to the underlying rendezvous store to keep it alive
    after ``destroy_process_group()``.
    '''
    if not dist.is_initialized():
        raise RuntimeError('Cannot register WORLD before torch.distributed is initialized')
    pid = os.getpid()
    backend = str(dist.get_backend())
    state = _DefaultProcessGroupState(
        backend=backend,
        timeout=timeout,
        store=_get_default_store(), # prevent store garbage collection
        rank=dist.get_rank(),
        world_size=dist.get_world_size(),
    )
    default_process_group_states[pid] = state
​
​
def _destroy_default_nccl_process_group() -> None:
    '''Destroy NCCL WORLD group and replace with temporary Gloo WORLD.
​
    1. Use canonical Gloo group to barrier all ranks.
    2. Call ``dist.destroy_process_group()`` (destroys all sub-groups).
    3. Replace cached Gloo group with ``None``.
    4. Initialize a new Gloo WORLD via ``_new_default_process_group``.
    5. Re-cache the new default group as the canonical Gloo group.
    '''
    state = default_process_group_states.get(os.getpid())
    if state is None or state.nccl_world_destroyed or not _uses_nccl(state.backend):
        return
​
    # Use the CPU group as an out-of-band synchronization point.
    dist.barrier(group=get_gloo_group())
    dist.destroy_process_group()
    set_gloo_group(None)
​
    _new_default_process_group(state, backend='gloo')
    set_gloo_group(_get_default_group())
    state.nccl_world_destroyed = True

评论区精华

无 review 讨论，PR 由作者自行审查合并。代码注释中体现了关键设计决策：使用 Gloo 组作为过渡（避免 CUDA 操作），通过 PrefixStore 生成唯一 store 前缀防止 stale key，以及通过环境变量提供逃生口。

• 暂无高价值评论线程

风险与影响

• 风险：
  • 若外部代码缓存了 dist.group.WORLD 引用，销毁后引用将失效，可能导致静默错误或崩溃。代码通过 SLIME_DESTROY_WORLD_PROCESS_GROUP=0 提供关闭能力。
  • 修改 init_gloo_group 使用 old_new_group 绕过猴子补丁，如果未来 PyTorch 或 Monkey Patch 逻辑改变，可能出现不一致。
  • 测试仅使用 monkeypatch 模拟复杂分布式调用，未进行实际多进程集成测试，可能遗漏真实环境下的竞态条件或 NCCL 错误处理问题。
  • 新增依赖 psutil 用于测试（测试文件未明显使用，可能是间接依赖），需确认无冲突。
• 影响：
  • 对用户：需在训练启动时调用 register_default_process_group（已在 actor 初始化中自动调用）。如需禁用可通过环境变量关闭。
  • 对系统：新增 slime/utils/reloadable_process_group.py 核心模块，修改 distributed_utils.py 和 actor.py。影响所有使用 MegatronTrainRayActor 的训练流程。
  • 对团队：提供了可复用的 WORLD 组重建模式，便于后续实现动态资源管理（如 release-train）。
  • 影响程度：中等，主要影响需要动态管理通信组的训练场景。
  • 风险标记：核心路径变更（actor.init）, 测试仅依赖 monkeypatch, 环境变量逃生口

关联脉络

• PR #2180 Add --release-train: PR 2180 引入了训练后释放 actor 组的能力，该 PR 提供了 WORLD 组销毁重建的基础机制，两者共同实现动态资源管理。