fix(#430): GRPO 训练在 worker 线程建立 thread-local stream

[dahai80]

修复 #430

GRPO 训练作业启动即失败:There is no Stream(gpu, 6) in current thread.

根因

_execute_grpo 经 asyncio.to_thread 在 worker 线程跑训练,而 GRPOTrainer._sample_completions 调用 mlx_lm.generate.generate_step,其内部 with mx.stream(generation_stream) 用的是 mlx_lm.generate 模块级 thread-local generation_stream,worker 线程未绑定,MLX 抛 Stream 错误。SFT/DPO/reward 训练不调 generate_step,故不触发。

改动

_execute_grpo 入口先 mx.new_thread_local_stream(mx.default_device()) 建立本线程默认 stream,原训练/保存逻辑下沉到 _run_grpo。HTTP/数据结构/适配器格式不变。

def _execute_grpo(self, job):
    import mlx.core as mx
    worker_stream = mx.new_thread_local_stream(mx.default_device())
    with mx.stream(worker_stream):
        self._run_grpo(job)

验证

• tests/unit/test_grpo.py + test_grpo_route.py:19 passed, 1 skipped。
• 真实端到端:服务跑起来后创建 GRPO 作业(Qwen3-0.6B-4bit, group_size=2, iters=2, length-reward fallback)
  • 修复前:status=failed,error="There is no Stream(gpu, 6) in current thread."
  • 修复后:status=completed,progress={iter:2, loss:0.0, mean_reward:58.0},adapter(adapters.safetensors + adapter_config.json)已写出。

Closes #430

🤖 Generated with Claude Code