[ZCC] Fix offload barrier never waiting for the in-flight D2H

[ForFishes]

Problem

The ZCC offload barrier never waits. ZeroCostCheckpointManager.sync_offload_status() compares
self.current_worker.global_step.value against self.global_step, but:

1. manager.global_step is only ever assigned inside update_zcc_workers(), whose sole caller
   maybe_update_zcc_worker() short-circuits on inner_opt.fused_buffer_version == self.manager.cache_version.
   fused_buffer_version does not change in steady state, so the value is written once per run
   (the step of the first save) and then frozen.
2. The worker simply echoes the value it received back (self.global_step.value = global_step),
   so the manager ends up comparing X against its own X.

The barrier therefore passes on the very first comparison from the second save onwards.

This matters because the ZCC worker is a spawn-ed process that maps the trainer's fused GPU
buffers over CUDA IPC and does the D2H copy on its own CUDA context and stream. There is no
CUDA stream ordering between the trainer's compute stream and the worker's copy stream — this
handshake is the only ordering guarantee. With it broken, nothing stops step N+1 from mutating
GPU memory that the step-N snapshot is still reading, which can silently corrupt saved
parameters / optimizer moments.

Evidence from a 2496-GPU run

log line	count
Waiting current worker offloading done (the branch that actually waits)	0
Current worker offloading done	48, every one with the same frozen manager_step

The run had Offload chunks: 1 (~9 GB per card in a single D2H), and an on_step_begin
callback that quantizes MoE expert weights to FP8 and calls optimizer.clear_param_storage(...),
i.e. a GPU mutation ~3s before on_optimizer_begin where the barrier used to sit.

Changes

1. manager.global_step is refreshed every step it is used. Both branches of
   ZeroCostCheckpointCallback.on_step_end now assign self.manager.global_step = state.global_step
   right after maybe_update_zcc_worker() and before requesting a worker. This is what the
   # set 'on-step-end' comment on the attribute declaration originally intended.
   state.global_step is already incremented before on_step_end, so the value stays >= 1
   and the assert global_step != 0 in get_idle_worker_for_saving is preserved.

2. The synchronization point moves to the beginning of the next step. Even a working barrier
   at on_optimizer_begin is too late, because on_step_begin callbacks may already mutate the
   buffers. New ZeroCostCheckpointManager.maybe_sync_offload_status() is called from
   Trainer._inner_training_loop just before on_step_begin. It is guarded by
   current_pipeline_hook_step != pipeline_hooks_steps: when the offload is still being sliced
   across pipeline hooks (PP > 1, or pipeline_hooks_steps > 1), the remaining chunks are only
   dispatched during this step, so waiting there would deadlock — those topologies keep
   synchronizing at on_optimizer_begin exactly as before. The original
   on_optimizer_begin sync is left in place as a backstop.

Two files, +45 lines, no deletions.

Known remaining limitations (pre-existing, not addressed here)

• When pipeline_hooks_steps > 1, only the first chunk is dispatched at on_step_end; the rest
  come from the next step's on_substep_end / PP hooks, i.e. after any on_step_begin mutation.
  The step-begin guard correctly declines to wait for those topologies.
• The EMA-only branch of on_step_end does not call zcc_pipeline_hook, so its chunk is also
  dispatched by the next step's on_substep_end.
• The wait is for the D2H only (the worker writes global_step.value after both wait_all()
  calls and before process_dump_task), not for the disk write.
• get_fused_param_mappings() captures buffer.param_buffer_ipc_meta once per run, but
  clear_param_storage() / reset_param_storage() free and re-allocate param_storage for the
  cleared colors on every step, so the IPC handles held by the worker stay valid only while the
  caching allocator happens to hand back the same block. Out of scope here.
• Throughput cost: if an on_step_begin callback mutates the fused buffers, the D2H can no
  longer overlap the next step's forward/backward at all. That is inherent to the design, and
  the correct trade against saving corrupt state.

Verification

ZeroCostCheckpointManager was instantiated via object.__new__ with a fake worker whose
global_step.value write-back is delayed by a configurable number of polls, and module-level
time.sleep monkeypatched to count polls. This exercises the real sync_offload_status,
maybe_sync_offload_status and zcc_pipeline_hook:

• pre-fix frozen step reproduces "barrier returns immediately";
• post-fix the barrier really polls until the worker echoes the current step;
• the guard returns immediately (no deadlock) for pipeline_hooks_steps > 1 and PP > 1;
• multiple workers with monotonically increasing step numbers never let a stale echo through.

[risemeup1111]

序号	位置	优先级	状态
1	IPC 地址漂移		✅
2	ZCC barrier 回归测试		🚧

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[risemeup1111]

地址漂移后不能继续使用旧 IPC handle 保存。这里忽略了 check_ipc_buffer_drift 的返回值，下一行又会在 fused_buffer_version == cache_version 时直接返回；worker 的 ParamFusionStorageHelper 只有通过 update_zcc_workers 收到新的 ipc_meta_mappings 才会刷新。因此 clear_param_storage() / reset_param_storage() 一旦换了地址，本次 PREPARE/OFFLOAD 仍会从已释放内存读取，可能静默损坏 checkpoint；同时上面的 _comm_buffer_ptrs = current 还会让后续相同 stale 地址不再告警。请在发出保存任务前让所有 rank 一致刷新 worker IPC metadata，或 fail fast/跳过本次保存，不能只告警后继续。

[risemeup1111]

当前 head 3b4153d 已移除本 PR 新增的 drift 检测及 warning-only 继续保存路径，增量不再引入该行为；此项按本 PR 范围已解决。

[risemeup1111]

请把这次 barrier 复现加入仓库回归测试。当前 diff 没有测试，现有 tests/ai_edited_test/trainer/test_ai_zero_cost_ckpt.py 也只覆盖枚举、哈希和 optimizer unwrap；PR 描述中的临时 fake-worker 验证不会在 CI 中执行。这个修复依赖 manager.global_step 刷新和 current_pipeline_hook_step guard 的组合，建议至少覆盖：旧 step 必须持续轮询、单 chunk 在 step begin 同步、多 chunk 未发完时不等待、多 worker 复用不会接受 stale echo，以及 on_step_end 两个分支都写入当前 step。

[Paddle-CI-Bot]

PaddleFormers Log Analysis

Run #31097360056 · Attempt 1

日志分析报告

流水线名称	问题标签	修复建议	日志片段
unittest-gpu-ci (job 92593754580)	model-inference / fd_fallback	检查 DeepseekV3 & GLM4-MoE 的 fd_fallback 路径与 test_model_tiny_logits 所用的模型推理逻辑是否受 trainer.py 新增的 maybe_sync_offload_status 调用影响	报错代码
Integration test (H20, multi-card) — Qwen sft (job 92593799954)	distributed-training / NCCL	在 model_utils.py:from_pretrained 的 dist.load_state_dict 前，确保此 PR 新增的 maybe_sync_offload_status() 调用不会在模型加载阶段（barrier 之前）触发不合法的 ZCC 同步，导致多卡 NCCL barrier 错误	报错代码
Integration test (H20, multi-card) — Qwen lora (job 92593799954)	checkpoint-path / API-breakage	qwen3_multicard_lora.yaml 的 model_name_or_path 被设为前一步 SFT 输出的本地路径 /workspace/checkpoints/qwen-sft，但 SFT 步骤已因 NCCL 错误未完成，导致该路径不存在；同时 hf_try_to_load_from_cache 对绝对路径缺乏前置检查（HFValidationError）	报错代码

失败的测试case:

# unittest-gpu-ci (run 31094642421)
tests/transformers/deepseek_v3/test_modeling.py::DeepseekV3IntegrationTest::test_fd_fallback
tests/transformers/deepseek_v3/test_modeling.py::DeepseekV3IntegrationTest::test_model_tiny_logits
tests/transformers/glm4_moe/test_modeling.py::Glm4MoeModelIntegrationTest::test_fd_fallback
tests/transformers/glm4_moe/test_modeling.py::Glm4MoeModelIntegrationTest::test_inference_no_attention

# Integration test H20 multi-card (run 31094642115)
Qwen sft   — OSError: NCCL 'unhandled cuda error' at paddle distributed barrier (process_group_nccl.cc:912)
Qwen lora  — HFValidationError: Repo id must be in the form 'repo_name' or 'namespace/repo_name': '/workspace/checkpoints/qwen-sft'

根本原因分析:

PR #4838 改动了两处：paddleformers/trainer/trainer.py（+11 行）和 paddleformers/trainer/utils/zero_cost_checkpoint.py（+34 行）。

Integration test 失败（直接因果）

• Qwen sft NCCL 崩溃：新增的 maybe_sync_offload_status() 在 _inner_training_loop 的 on_step_begin 之前调用。Qwen sft CI 走的是多卡（8 GPU）PP 模式加载预训练 checkpoint（/workspace/checkpoints/qwen-pt），model_utils.py:2932 处的 dist.load_state_dict 发起 paddle.distributed.barrier。若 ZCC manager 在此时已有一个 stale worker（来自之前的 pt 步骤），maybe_sync_offload_status 触发了错误的同步等待，导致部分 rank 进入 barrier 而另一些 rank 未准备好，NCCL 报 "unhandled cuda error"（rank 4 崩溃，exit code 241）。

• Qwen lora HFValidationError：lora 测试依赖 sft 输出的 /workspace/checkpoints/qwen-sft。sft 因 NCCL 崩溃而未生成该目录，AutoConfig.from_pretrained('/workspace/checkpoints/qwen-sft') 走进 hf_try_to_load_from_cache，huggingface_hub 对绝对路径做 validate_repo_id 校验，抛出 HFValidationError（cascade failure）。

unittest 失败（可能相关）

• DeepseekV3IntegrationTest::test_fd_fallback、test_model_tiny_logits 和 Glm4MoeModelIntegrationTest::test_fd_fallback、test_inference_no_attention 均涉及 MoE 模型推理的 fd（forward dispatch）路径。PR 在 trainer.py 中插入了 maybe_sync_offload_status() 调用，该函数操作 ZCC manager 的共享状态，若测试套件中复用了 trainer 实例或存在 side-effect，可能影响后续 MoE 推理的数值结果或状态机。需要进一步确认这 4 个 test 是否为本 PR 新引入的 regression，还是与 PR 无关的已有 flakiness。

修复建议:

1. Qwen sft NCCL：在 trainer.py 的 maybe_sync_offload_status() 调用前增加守卫条件，确保仅在 self.args.use_zero_cost_checkpoint 且 self.zcc_manager is not None 且 训练循环已启动（即已完成至少一次 on_step_end）时才执行同步；模型加载阶段（from_pretrained 内的 barrier）不能被 ZCC 同步逻辑干扰。

2. Qwen lora HFValidationError：在 paddleformers/utils/download/download.py 的 hf_try_to_load_from_cache 入口处，判断 repo_id 是否为绝对路径（os.path.isabs(repo_id)），若是则直接走本地文件查找，跳过 validate_repo_id；同时修复 CI 脚本依赖链——lora 测试应在 sft 成功后才执行，或将 model_name_or_path 改为不依赖前一步产物的独立路径。

3. unittest MoE fd_fallback 回归：在本地复现 test_fd_fallback / test_inference_no_attention，检查 PR 对 trainer.py 的改动是否修改了任何模块级全局状态（如 Paddle FLAGS、进程组）；若确认为新引入的 regression，隔离 maybe_sync_offload_status() 对 ProcessGroup 的副作用，确保其不影响非 ZCC 代码路径。

---

🔍 准确性记录：请点击评论底部 😊 图标，选择 👍（准确）或 👎（有误），将自动记录到 CI 监控系统

_{🔄 每次 Re-run 后自动更新}