Skip to content

[fix] IPC barrier - #2114

Open
PengchengShi00 wants to merge 2 commits into
InternLM:mainfrom
PengchengShi00:fix_ipc
Open

PengchengShi00 wants to merge 2 commits into
InternLM:mainfrom
PengchengShi00:fix_ipc

Conversation

@PengchengShi00

Copy link
Copy Markdown
Collaborator

Summary

  1. 修复 LMDeploy IPC 持久化 buffer 写覆盖

LMDeploy IPC 按 dtype 缓存一块 GPU buffer,后续 bucket 是 copy 进同一块,对端读的也是这块。

  • rollout TP=1:每个 rank 都是 head,写完就 POST,等 HTTP 200 再写下一批。训练写和推理读完全串行,不会因抢跑覆写。
  • rollout TP>1:只有 head POST。非 head 只 gather_object,以前中间 batch 的 cpu_group barrier 还绑在 train_enable_ep 上。训练 没开 EP 时,非 head 不必等 head 的 200,就可以把下一批 copy 进同一块 buffer。若下一次 unshard 的 FSDP 组 不含该 engine 的 head,就不会被集合通信卡住,推理还在读第 N 批时内容已变成 N+1。

代码改动:LMDeployIPCBackendAdapter.after_update_per_batch 从 finished or (train_enable_ep and rollout_tp > 1) 改为 finished or rollout_tp > 1。

即 只要 TP>1,每个中间 batch 都 barrier,非 head 等到 head HTTP 回来再写下一拍。顺带去掉每个 submodule group 后的 empty_cache(),正确性改靠 barrier,不再依赖 allocator 副作用。

  1. 其他引擎 / 其他传输为什么不容易踩同一坑

这是 LMDeploy 持久化 buffer 特有的「同一块还活着的 IPC 页被写下一批」。

SGLang / vLLM IPC:没有 pin 住的 buffer。每个 bucket 会 cat 出新的 flattened / reduce_tensor,发完引用就可以没了,不会主动 copy 进「上一拍那块还在被对端 map 的 tensor」。因此没有这次修的那条竞态。

NCCL / checkpoint-engine:不复用 CUDA IPC handle 指向的同一块 sender buffer,接收端自己的内存,没有这条写覆盖。

  1. 单测

共卡(IPC)和分离部署(NCCL)都加了 generate 正确性检查,覆盖 SGLang 和 LMDeploy:

同一引擎 greedy generate 两次(确认引擎自身可复现)
权重更新(无训练 step)
再 generate,和第一次比 response / response_ids / sampled-token logprobs(atol/rtol=1e-5)
LMDeploy /generate 写死 do_sample=True,greedy 用 top_k=1 且 temperature=1.0(temperature=0 会除零)。SGLang greedy 仍用 temperature=0。

checkpoint 检查(SGLang WeightChecker)是互补用例,依赖私有 patch,默认 skip。

@PengchengShi00

Copy link
Copy Markdown
Collaborator Author

@claude review

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant