`IterablePackingDataset` 使用 fork 模式导致 DeepSpeed ZeRO-3 训练死锁

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

Bug Report: IterablePackingDataset 使用 fork 模式导致 DeepSpeed ZeRO-3 训练死锁

环境信息

• ms-swift 版本: 4.3.1
• Python: 3.10
• 训练框架: DeepSpeed ZeRO-3 + torchrun (8 GPU)
• 模型: Qwen3.5-MoE (多模态)
• 触发条件: --packing true --streaming true --deepspeed zero3

问题描述

当同时启用 packing=true 和 streaming=true 时，IterablePackingDataset.__init__() 中使用默认的 multiprocessing.Process（即 fork 模式）启动 packing worker 进程。在 DeepSpeed ZeRO-3 环境下，父进程已初始化 CUDA context 和 NCCL 通信状态，fork 出的子进程继承了这些状态，导致：

1. Packing worker 在 _in_queue.get() 上阻塞等待数据
2. 主进程继续执行到 get_model_parameter_info() → 触发 ZeRO-3 all_gather → torch.cuda.synchronize() 永久阻塞
3. 主进程卡死 → 永远不会进入训练循环 → 永远不会向 _in_queue 投递数据 → packing worker 也永久阻塞
4. 形成双向死锁，训练完全卡死，无任何错误输出

py-spy 调用栈确认：

100% synchronize (torch/cuda/__init__.py)
  ← all_gather (deepspeed/runtime/zero/partition_parameters.py)
    ← get_model_parameter_info (swift/utils/transformers_utils.py)
      ← run (swift/pipelines/train/sft.py)

该问题具有机器相关性：不同机器的 GPU 拓扑、NCCL 版本、内核版本差异可能导致 fork 后 CUDA/NCCL 状态表现不同，因此同一代码在某些机器上正常、另一些机器上死锁。

根因定位

文件：swift/dataset/packing.py，IterablePackingDataset.__init__() 第 136-142 行：

self._in_queue = mp.Queue()
self._out_queue = mp.Queue()
for _ in range(self.num_proc):
    worker = mp.Process(target=self._processor, daemon=True)  # ← 默认 fork
    worker.start()

mp.Process 和 mp.Queue 使用默认的 start method（Linux 下为 fork），在 CUDA/DeepSpeed 已初始化的进程中 fork 子进程是不安全的。

修复方案

将 IterablePackingDataset 中的 multiprocessing 对象显式切换为 spawn context，避免继承父进程的 CUDA/NCCL 状态：

# swift/dataset/packing.py - IterablePackingDataset.__init__()

# Before:
self._in_queue = mp.Queue()
self._out_queue = mp.Queue()
self.workers = []
self.cyclic = cyclic
for _ in range(self.num_proc):
    worker = mp.Process(target=self._processor, daemon=True)
    worker.start()
    self.workers.append(worker)

# After:
self._mp_ctx = mp.get_context('spawn')
self._in_queue = self._mp_ctx.Queue()
self._out_queue = self._mp_ctx.Queue()
self.workers = []
self.cyclic = cyclic
for _ in range(self.num_proc):
    worker = self._mp_ctx.Process(target=self._processor, daemon=True)
    worker.start()
    self.workers.append(worker)

注意事项

• spawn 模式要求 _processor 方法及其依赖的对象（如 template）可被 pickle 序列化。如果 template 不可序列化，需要将 _processor 改为模块级函数，通过队列传递必要参数而非依赖 self。
• PackingDataset（非 streaming 版本）同样使用了 mp.Process(daemon=True)（第 64 行），建议一并修改以保持一致性。
• 建议在 PackingDataset 和 IterablePackingDataset 中都加上注释说明使用 spawn 的原因，防止后续维护者误改回 fork。

临时 Workaround

在无法升级 swift 的情况下，可通过以下方式绕过：

1. 关闭 packing：--packing false，用 gradient_accumulation_steps 补偿有效 batch size
2. 关闭 streaming：--streaming false，走 PackingDataset 路径（MapDataset，不受此 bug 影响）
3. 手动 patch 服务器上的 packing.py 文件（如上述修复方案）

---

How to Reproduce / 如何复现

如上。

Additional Information / 补充信息

No response

Start in swift/dataset/packing.py at IterablePackingDataset.init(), especially the queue and worker creation around lines 136-142, then compare the multiprocessing setup in PackingDataset around line 64. Reproduce with packing, streaming, DeepSpeed ZeRO-3, and torchrun if the environment is available. Done means both dataset paths use the intended process context without the reported training deadlock, while confirming the objects passed to workers remain serializable.