[DanielSun11]

Before submitting

• Lint code. If there are lint issues, please format the code first.

# Install and register `pre-commit` in the project folder
pip install pre-commit && pre-commit install

# Process previous code files separately
pre-commit run --file XXXX.py

• Add test cases into tests folder. If there are codecov issues, please add tests cases first.

PR types

Performance optimization

PR changes

Others

Description

Reduce the main-process overhead of a zero cost checkpoint (ZCC) save and its impact on the next step.

• Slim the trainer_state sent to the worker. The PREPARE task used to carry the whole TrainerState; the feeder thread of mp.Queue pickles it while holding the GIL, so large arrays in trial_params (data_status / saved_data_status / last_start_data_status) stalled the training thread. Now data_status is never sent, ranks that do not write trainer_state.bin get none of them, and with a single worker last_start_data_status is sent only when it changes and saved_data_status is sent as "base (once) + accumulated sparse diff" when the data side provides _zcc_saved_base / _zcc_saved_diff. The worker rebuilds the dense arrays (_restore_trainer_state), so the on-disk format is unchanged. Multiple workers keep sending dense copies.
• Defer the save step's D2H until after the step-tail logging (flush_deferred_offload), so it does not compete for PCIe with the small cross-node collectives in logging. Every sync point flushes a pending offload first, so trainers that do not call the flush still work.
• Chunked D2H: new arg zcc_d2h_chunk_mb (default 64, 0 = one copy); each chunk is copied before the next is issued.
• Pinned worker-side CPU buffers for model weights; in-place lerp_ for the EMA of fp32 params; flash (tmpfs) checkpoint files written from a thread pool (persistent writes stay serial).
• TrainerState.save_to_json accepts numpy arrays in trial_params.

Tests: added unit tests for the deferred offload (including the flush-before-sync fallbacks) and for the trainer_state slimming round trip (single worker, multiple workers, ranks without trainer_state, missing cache). End-to-end on 8 nodes (ERNIE MoE, flash every 7 steps): ZCC checkpoints match the original implementation in structure and data status, resuming from a ZCC checkpoint reproduces the first-step loss bit-exactly, old/new checkpoint formats resume in both directions, and the extra time of the save step drops from ~1.9 s to ~0.

[Paddle-Review-Bot]

Paddle-Bot Review Board （review完成)

序号	位置	优先级	规则来源	状态
1	EMA 累计中 pinned 参数缓冲的 lerp_		仓库规则：算子、数值与性能	🚧

_{Powered by Nyanpasu claude with Opus 5 默认推理级别, please check the suggestions carefully.}

[Paddle-Review-Bot]

本 PR 把 init_buffer 中 flatten 的 cpu_buf 改成了 pinned 内存（L599-601），但这里的 ema_buf.lerp_(cpu_buf, ...) 仍然直接以 pinned 张量作为输入。

Paddle 动态图按输入张量的 place 选择 kernel，device_guard("cpu") 改变不了这个结果。默认 FLAGS_pinned_memory_as_cpu_backend=False 时，GPUPINNED 被映射为 GPU backend，所以这里会选中 GPU 的 lerp kernel（L403 的新注释说 zeros_like(pinned) 会分配到 GPU 上，也是同一个原因）。可能的后果：

• ema_buf 是 inplace 输出，会被重新分配到 worker 的显存上；
• TransformData 对 GPUPINNED 输入有 ShareBufferWith 的特殊处理，可能让 inited_buffers[index][1] 本身变成 GPU 缓冲，导致之后的 async_offload_with_offset/state_dict 读写错位置。

开启 zcc_save_ema_coef 且模型里有 fp32 参数（如 MoE gate）时都会走到这条路径。master weights 分支（L428-435）已经用 pin2cpu_zero_copy_fp32 规避了同样的问题，这里建议照同样的方式处理（这条路径上 cpu_buf 由 L407 保证是 fp32，flatten 后也是连续的）：

_, cpu_buf = self.param_fusion_storage_helper.inited_buffers[index]
if cpu_buf.place.is_cuda_pinned_place():
    cpu_buf = pin2cpu_zero_copy_fp32(cpu_buf)
ema_buf.lerp_(cpu_buf, 1 - self.ema_coef)

另外建议在 GPU 上补一个 EMA 场景的验证：累计一次以后，检查 ema_buf.place 和 inited_buffers[index][1].place 都没有变化。