[gspark-etri]

Background

save_ckpt and save_lora_ckpt wrote checkpoints from every rank to names that depend on the pipeline stage only (#165). With tp > 1, the TP ranks of a stage wrote different shards to one file (one shard lost, reload loads wrong weights). With dp > 1, every replica wrote and pruned the same files (FileNotFoundError race on os.remove). Without ckpt_dir_postfix the directory was checkpoint_None_<stage>, because the or fallback on a non-empty f-string never triggers. save_hf_ckpt already did this right: only DP rank 0 writes, files are named per TP rank.

Changes

All in opt_prime/opt_prime/opti_pri.py, following save_hf_ckpt (rank = (pp*dp + dp_i)*tp + tp_i, so dp_i = (rank // tp) % dp, tp_i = rank % tp):

• __init__ (:465): checkpoint_{postfix}_{stage} only when a postfix is given, otherwise checkpoint_{stage}.
• save_ckpt: returns on DP ranks other than 0; the file gets a _tp{tp_i} suffix when tp > 1, and pruning of old files globs only this TP rank's files.
• load_ckpt: reads the same name (all DP replicas load DP rank 0's file for their TP rank).
• save_lora_ckpt: returns on DP ranks other than 0 (it already named TP shards; load_lora_ckpt is unchanged).
• New opt_prime/tests/test_ckpt_names.py: four ranks of tp=2, dp=2 built without a process group (Optimus_p.__new__, just the attributes the save/load paths read), each with different weights. It checks that exactly one file per TP shard exists, holding DP rank 0's weights, and that DP rank 1 reloads them, for both save_ckpt and save_lora_ckpt. It also checks that PP-only names are unchanged, and the directory name with and without a postfix (runs __init__ with a fake Comm up to the point where it picks the GPU).

Verification

Environment: Python 3.10, torch 2.10.0+cu128, transformers 5.1.0, CPU only. Base: optimus-rm @253be20.

cd opt_prime && TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests/test_ckpt_names.py

• Before the fix: 3 failed, 2 passed.
  • save_ckpt produced only ckpt_step_1_epoch_0.pt, not one file per TP shard.
  • With save_lora_ckpt, the reloaded TP-rank-1 weights were DP rank 1's (the last writer), not DP rank 0's.
  • The directory was checkpoint_None_0.
  • The PP-only name test and the postfix case passed, as expected.
• After the fix: 5 passed.

Full suite (cd opt_prime && TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests):

• Before (253be20): 400 passed, 23 skipped, 135 subtests passed
• After: 405 passed, 23 skipped, 135 subtests passed

Compatibility and limits

• PP-only runs (tp=1, dp=1, which covers all shipped examples) keep the old file names, so existing checkpoints still load. The only change for them is the directory name when no postfix is given: checkpoint_<stage> instead of checkpoint_None_<stage>. To resume from an old run, rename that directory or pass the old name's postfix.
• With dp > 1, reload assumes the checkpoint directory is visible to every DP replica (a shared filesystem, or a single node), as save_hf_ckpt/merge_hf_ckpt.py already assume.
• Old tp > 1 checkpoints were already corrupt (one shard per stage); they are not readable under the new names.
• Not run on a real multi-GPU TP/DP job; the shipped examples are PP-only.

Related Issue

Closes #165

🤖 Generated with Claude Code