Skip to content

Fix deterministic checkpoint file sharding - #103

Open
koriyoshi2041 wants to merge 1 commit into
MoonshotAI:mainfrom
koriyoshi2041:fix-deterministic-checkpoint-shards
Open

Fix deterministic checkpoint file sharding#103
koriyoshi2041 wants to merge 1 commit into
MoonshotAI:mainfrom
koriyoshi2041:fix-deterministic-checkpoint-shards

Conversation

@koriyoshi2041

Copy link
Copy Markdown

Problem

split_checkpoint_files() partitions the raw result of os.listdir(). Directory enumeration order is not guaranteed, so ranks on different nodes can derive different file-to-rank assignments for the same checkpoint. That can make checkpoint registration overlap some shards and omit others.

Fix

Sort the filtered .safetensors paths before slicing them across ranks.

The regression test presents the same directory contents in forward and reverse enumeration order and verifies that both produce the same complete, non-overlapping assignment.

Tests

uv run pytest -q tests/test_example_update.py
1 passed

uv run ruff check examples/update.py tests/test_example_update.py
All checks passed!

uv run ruff format --check examples/update.py tests/test_example_update.py
2 files already formatted

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.

1 participant