[optim] bucket checkpoint save to avoid CPU memory spike during ckpt saving

[yueming-yuan]

Problem

Synchronous torch_dist save stages all write buckets to host memory before any file write starts (FileSystemWriterAsync.preload_tensors runs as one shot in AsyncRequest.execute_sync). The host transient equals the rank's full checkpoint shard.

On GLM-5.2 744B RL training (16x GB300, TP8/PP4/DP2, distributed optimizer), that is ~650 GB/node (fp32 mains + Adam moments + bf16 params) against ~780 GB available — the first save deterministically OOM-killed the job across all nodes within a minute.

Change

Two opt-in args, defaults preserve existing behavior exactly:

• --ckpt-save-staging-buckets N (default 0): sync save stages and writes at most N buckets at a time. The wave loop lives inside write_preloaded_data_multiproc: preload one wave D2H → fork writers → join → free → next. Results are merged across waves and put to the results queue once, so the retrieve_write_results protocol is unchanged. Incompatible with --async-save (asserted): waved D2H must run in the training process, and async staging must complete before training resumes anyway.
• --ckpt-save-thread-count N (default 2): exposes the existing thread_count (buckets/files per rank) so waves have useful granularity.

Host peak becomes waves × bucket_size instead of the full shard: with thread-count 8, staging-buckets 2 on the workload above, ~160 GB/node instead of ~650 GB.

Exactness

Bucket split, file names, and file contents are fixed at planning time (prepare_write_data); this change only reorders when D2H and writes happen. With the same thread_count, checkpoints with staging on/off are byte-identical. The async path is untouched, and staging_max_buckets=0 degenerates to a single wave with no internal preload (current behavior).

Overhead

~1-3% of save wall time (per-wave join stragglers + D2H gaps between waves; writes remain the bottleneck and write concurrency per rank is unchanged). Total D2H bytes and per-bucket fsync behavior are identical to the current path.

Validation

• GLM-5.2 744B, 16x GB300, --save-interval 1 --ckpt-save-thread-count 8 --ckpt-save-staging-buckets 2: 10.4 TB checkpoint written successfully; MemAvailable stable at 585-632 GB/node throughout the save, Dirty ~0 (previous run with unbounded staging was OOM-killed at the same point).
• Wave slicing/result-merging logic unit-checked (global index mapping, single queue put, exception short-circuit, staging=0 equivalence).