Skip to content

fix: fence live checkpoint recovery from PG reinit#264

Open
Abhishek21g wants to merge 1 commit into
PrimeIntellect-ai:mainfrom
Abhishek21g:abhishek/live-checkpoint-recovery-fence
Open

fix: fence live checkpoint recovery from PG reinit#264
Abhishek21g wants to merge 1 commit into
PrimeIntellect-ai:mainfrom
Abhishek21g:abhishek/live-checkpoint-recovery-fence

Conversation

@Abhishek21g

Copy link
Copy Markdown

Summary

Live checkpoint P2P transfer used global_pg while surviving ranks could call maybe_reinit_global_pg during inner steps — destroying the process group mid send/recv when a node leaves (#156).

  • Add LiveRecovery cluster-wide fence (begin_recovery / end_recovery / wait_until_recovery_done)
  • Defer maybe_reinit_global_pg while recovery is in progress
  • Wrap send/recv in train.py with fence; all ranks wait before outer-step work
  • Unit tests for fence coordination

Closes #156

Verify

uv run pytest tests/test_live_recovery.py -v
# Full dist repro: 3-node live recovery with node kill mid-transfer (Linux, 2+ GPU)

…#156)

Defer maybe_reinit_global_pg while live recovery is in progress and have
all ranks wait on the recovery fence before inner steps. Prevents surviving
nodes from tearing down global_pg mid send/recv when a peer leaves.
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.

Live checkpoint breaks if node leaves before recovery is complete

1 participant