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

Open
shengliangxu wants to merge 48 commits into
base: main
Choose a base branch
from
+1,056 −3
Open

[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
48 commits
Select commit Hold shift + click to select a range
132eef1
Add the distributed PTQ session and its distribution plan
shengliangxu Oct 10, 2026
df59de9
Keep synced session attributes the same on every rank
shengliangxu Oct 13, 2026
e9f72db
Let a synced session attribute combine every rank's value
shengliangxu Oct 13, 2026
048dbc2
Agree on a one-off value across ranks without keeping it on the session
shengliangxu Oct 14, 2026
70c007b
Declare a synced session attribute by its type: Synced[T]
shengliangxu Oct 14, 2026
4ddd7fd
Set a synced fact through its handle: session.plan.decide(fn, rank)
shengliangxu Oct 14, 2026
b9080c4
Name the session's collectives after what they do: broadcast_from and…
shengliangxu Oct 14, 2026
6797d3a
Drop the session object; keep the process group's lifetime and the co…
shengliangxu Oct 14, 2026
82a4e2b
Record a model's distribution plan on it, and make the plan a value w…
shengliangxu Oct 14, 2026
b11d6d8
Yield this process's DistributedState from distributed_run
shengliangxu Oct 14, 2026
7f46ef6
Own every group a run sets up, and set up FSDP2's one-rank group here
shengliangxu Oct 14, 2026
8bf70e2
Drop setup_fsdp2_process_group for the FSDP2 load to set its group up
shengliangxu Oct 14, 2026
bd31620
Refuse to record a second, different plan on a model
shengliangxu Oct 15, 2026
67cd9bf
Reject a plan degree below 1
shengliangxu Oct 15, 2026
967342d
Record a model's plan once, refusing a reload as well as a redistribu…
shengliangxu Oct 15, 2026
cccb71d
Leave a process group the caller set up before the run, under a launc…
shengliangxu Oct 15, 2026
05695ae
Exit a rank whose body raised in a distributed run instead of destroy…
shengliangxu Oct 16, 2026
693ef98
Test broadcast_from, all_reduce and the plan guard on two real ranks
shengliangxu Oct 16, 2026
ad2999e
Trim the distribution plan's and distributed_run's tests
shengliangxu Oct 16, 2026
7a92514
Destroy the bounded-timeout group when its block exits
shengliangxu Oct 16, 2026
629afeb
Date test_distributed_run 2026, as LICENSE_HEADER does
shengliangxu Oct 22, 2026
8fce8c9
Refuse a non-CUDA device under a launch, and a rank outside the run
shengliangxu Oct 24, 2026
d90a8f8
Join the collective on a SystemExit too, and bound a barrier named gr…
shengliangxu Oct 24, 2026
0acb6bc
Say that a module added after a plan is recorded must read it off its…
shengliangxu Oct 24, 2026
07c0fdb
Abort a run only on its body's own error, not on one its caller handles
shengliangxu Oct 25, 2026
dcde6a2
Read a run's rank from its live group; barrier unless its body raised
shengliangxu Oct 25, 2026
b1dd500
Exit every rank on a collective's SystemExit; fail unpicklable values
shengliangxu Oct 25, 2026
ad2689c
Fail every rank of an all_reduce alike when ranks fail differently
shengliangxu Oct 25, 2026
a99099d
Bind the rank's CUDA device under a launch, also to a group the calle…
shengliangxu Oct 27, 2026
1a53333
Refuse the collectives on a launch whose process group is not set up
shengliangxu Oct 27, 2026
def9c44
Say why the collectives check and unpickle values as they do
shengliangxu Oct 27, 2026
23549a6
Serialize a collective's value once, where the failure is still sent
shengliangxu Oct 27, 2026
f405e48
Test that a collective's value the other ranks could not load fails o…
shengliangxu Oct 27, 2026
dcf5628
Connect a bounded group on every rank before a job can destroy it
shengliangxu Oct 27, 2026
1d3d063
Say what the collectives' load-back check proves, and why it loads fully
shengliangxu Oct 27, 2026
ae4fd1b
Give broadcast_from's sending rank the loaded copy too
shengliangxu Oct 27, 2026
18f8f8e
Say a single process gets fn's value as it is, unserialized
shengliangxu Oct 27, 2026
f0153ac
Say how distribution_plan reads a launch before its group, and what t…
shengliangxu Oct 27, 2026
356bbcd
Say which runs abort a rank whose body raised
shengliangxu Oct 27, 2026
4a68b4a
Say a job names its group by keyword under bounded_collectives
shengliangxu Oct 27, 2026
df5f78a
Clear WORLD_SIZE in the single-process collectives test
shengliangxu Oct 27, 2026
f433158
Send a failure whose text cannot be formatted, rather than fail befor…
shengliangxu Oct 27, 2026
1d06a1a
Point a model without a plan at set_distribution_plan only
shengliangxu Oct 27, 2026
8f722de
Say inline why _decoded loads with weights_only=False, as SECURITY.md…
shengliangxu Oct 27, 2026
4520ab6
Refuse a plan recorded in a run of another number of ranks
shengliangxu Oct 27, 2026
d2c8924
Keep the collective helpers' docstrings to their contracts
shengliangxu Oct 28, 2026
4cb61ec
Format distributed_run's own error report through _text too
shengliangxu Oct 28, 2026
73172fb
Say distributed_run's exit path assumes every rank exits
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