Skip to content

[OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shengliangxu merged 60 commits into from
Oct 28, 2026
Merged

[OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
60 commits
Select commit Hold shift + click to select a range
253c7f2
Add the distributed PTQ session and its distribution plan
shengliangxu Oct 10, 2026
0a9333a
Keep synced session attributes the same on every rank
shengliangxu Oct 13, 2026
edf8c54
Let a synced session attribute combine every rank's value
shengliangxu Oct 13, 2026
d597a75
Agree on a one-off value across ranks without keeping it on the session
shengliangxu Oct 14, 2026
0c792bd
Declare a synced session attribute by its type: Synced[T]
shengliangxu Oct 14, 2026
4826253
Set a synced fact through its handle: session.plan.decide(fn, rank)
shengliangxu Oct 14, 2026
ebe52aa
Name the session's collectives after what they do: broadcast_from and…
shengliangxu Oct 14, 2026
683a536
Drop the session object; keep the process group's lifetime and the co…
shengliangxu Oct 14, 2026
354ed35
Record a model's distribution plan on it, and make the plan a value w…
shengliangxu Oct 14, 2026
767931b
Yield this process's DistributedState from distributed_run
shengliangxu Oct 14, 2026
a525796
Own every group a run sets up, and set up FSDP2's one-rank group here
shengliangxu Oct 14, 2026
a2c6e83
Drop setup_fsdp2_process_group for the FSDP2 load to set its group up
shengliangxu Oct 14, 2026
62e5db8
Refuse to record a second, different plan on a model
shengliangxu Oct 15, 2026
7230863
Reject a plan degree below 1
shengliangxu Oct 15, 2026
9b6ce9f
Record a model's plan once, refusing a reload as well as a redistribu…
shengliangxu Oct 15, 2026
2c663b5
Leave a process group the caller set up before the run, under a launc…
shengliangxu Oct 15, 2026
42f7fc8
Exit a rank whose body raised in a distributed run instead of destroy…
shengliangxu Oct 16, 2026
a81a680
Test broadcast_from, all_reduce and the plan guard on two real ranks
shengliangxu Oct 16, 2026
dc3a56c
Trim the distribution plan's and distributed_run's tests
shengliangxu Oct 16, 2026
1646260
Destroy the bounded-timeout group when its block exits
shengliangxu Oct 16, 2026
fe0fefc
Date test_distributed_run 2026, as LICENSE_HEADER does
shengliangxu Oct 22, 2026
8879e72
Refuse a non-CUDA device under a launch, and a rank outside the run
shengliangxu Oct 24, 2026
9292157
Join the collective on a SystemExit too, and bound a barrier named gr…
shengliangxu Oct 24, 2026
e901ce7
Say that a module added after a plan is recorded must read it off its…
shengliangxu Oct 24, 2026
c01124e
Abort a run only on its body's own error, not on one its caller handles
shengliangxu Oct 25, 2026
0a7b1e8
Read a run's rank from its live group; barrier unless its body raised
shengliangxu Oct 25, 2026
efa7e4a
Exit every rank on a collective's SystemExit; fail unpicklable values
shengliangxu Oct 25, 2026
78512d7
Fail every rank of an all_reduce alike when ranks fail differently
shengliangxu Oct 25, 2026
b44032f
Bind the rank's CUDA device under a launch, also to a group the calle…
shengliangxu Oct 27, 2026
11f4330
Refuse the collectives on a launch whose process group is not set up
shengliangxu Oct 27, 2026
ea1385a
Say why the collectives check and unpickle values as they do
shengliangxu Oct 27, 2026
80d1afb
Serialize a collective's value once, where the failure is still sent
shengliangxu Oct 27, 2026
50a45ac
Test that a collective's value the other ranks could not load fails o…
shengliangxu Oct 27, 2026
b176479
Connect a bounded group on every rank before a job can destroy it
shengliangxu Oct 27, 2026
0bbc570
Say what the collectives' load-back check proves, and why it loads fully
shengliangxu Oct 27, 2026
4f2d1bc
Give broadcast_from's sending rank the loaded copy too
shengliangxu Oct 27, 2026
ec53ebe
Say a single process gets fn's value as it is, unserialized
shengliangxu Oct 27, 2026
4841f86
Say how distribution_plan reads a launch before its group, and what t…
shengliangxu Oct 27, 2026
f0be610
Say which runs abort a rank whose body raised
shengliangxu Oct 27, 2026
dbca8fd
Say a job names its group by keyword under bounded_collectives
shengliangxu Oct 27, 2026
ea18be1
Clear WORLD_SIZE in the single-process collectives test
shengliangxu Oct 27, 2026
9b2d2a1
Send a failure whose text cannot be formatted, rather than fail befor…
shengliangxu Oct 27, 2026
5b0ec14
Point a model without a plan at set_distribution_plan only
shengliangxu Oct 27, 2026
e7d59ea
Say inline why _decoded loads with weights_only=False, as SECURITY.md…
shengliangxu Oct 27, 2026
e3cba99
Refuse a plan recorded in a run of another number of ranks
shengliangxu Oct 27, 2026
0aa7306
Keep the collective helpers' docstrings to their contracts
shengliangxu Oct 28, 2026
eb875ea
Format distributed_run's own error report through _text too
shengliangxu Oct 28, 2026
c94339a
Say distributed_run's exit path assumes every rank exits
shengliangxu Oct 28, 2026
df560b7
Restore two contract caveats in the collectives' docstrings
shengliangxu Oct 28, 2026
4f16025
Move the distributed_run tests into test_distributed.py
shengliangxu Oct 28, 2026
2ffa84e
Load the collectives' values onto the CPU
shengliangxu Oct 28, 2026
4c2d6d2
Say in the collectives' docstrings that a tensor comes back on the CPU
shengliangxu Oct 28, 2026
caa8575
Warn when distributed_run's timeout cannot apply to the caller's group
shengliangxu Oct 28, 2026
eba196a
Load the collectives' values with weights_only; yield the device as a…
shengliangxu Oct 28, 2026
d12b0f2
Fold the collectives' encoding into _serialize/_deserialize; trim the…
shengliangxu Oct 28, 2026
cdb2a44
Keep the launch rows telling local from global rank, and is_main
shengliangxu Oct 28, 2026
79b6285
Round-trip a DistributionPlan through the collectives' encoding
shengliangxu Oct 28, 2026
c9f3eb5
Say a single process takes distributed_run's timeout without a warning
shengliangxu Oct 28, 2026
1223b27
Warn about an unused timeout only for a caller's group with peers
shengliangxu Oct 28, 2026
cf4597d
Gate the timeout warning on is_distributed_run and note what loads trust
shengliangxu Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Loading
Loading