Skip to content

Make pass contracts cover the passes DeepCompile actually schedules; Renaming files - #8251

Draft
pengdurice wants to merge 4 commits into
deepspeedai:masterfrom
pengdurice:peng-fix-conflict-contract-v1
Draft

Make pass contracts cover the passes DeepCompile actually schedules; Renaming files#8251
pengdurice wants to merge 4 commits into
deepspeedai:masterfrom
pengdurice:peng-fix-conflict-contract-v1

Conversation

@pengdurice

Copy link
Copy Markdown
Contributor

[DeepCompile] Make pass contracts cover the passes DeepCompile actually schedules

Follow-up to #8139. Three gaps kept PassContract from firing on real schedules:

  1. conflicts_with was ignored whenever the other pass had no registered contract.
  2. Only four pass modules were registered, so offload_parameters and the ZeRO-1/2 reduce passes
    could not be validated at all.
  3. The one incompatibility DeepSpeed enforces — one offload target per run — was checked inside
    if schedule is None: in init_z3.py, so a user-supplied schedule combining both was accepted.

The conflict fix

validate_schedule matches conflicts against applied, but an uncontracted pass hit continue
before applied.append(name), so a conflict naming it was missed in both orderings. Resolving a
missing contract to a shared empty one instead of branching around it fixes that and removes two
lines of control flow:

contract = _pass_contracts.get(name, _UNCONSTRAINED)

The pass stays unconstrained; it is simply visible to conflicts others declare. test_conflict_is_symmetric
missed this because it registers the second pass with an empty contract, which still lands in the registry.

Contracts

Registered name Callable Contract
zero1_compile add_z1_reduce conflicts with zero3_compile
zero2_compile add_z2_reduce conflicts with zero3_compile, zero1_compile
offload_parameters offload_parameter_fwd requires z3_gather_release
offload_adam_states move_opt_states requires opt_states_evicted, conflicts with the three below
offload_adam_states_sync move_opt_states_sync conflicts with the three below, and offload_adam_states
offload_adam_states_for_init offload_adam_states_for_init provides opt_states_evicted, conflicts with the three below

The three shared conflicts are offload_parameters, zero1_compile, zero2_compile.

  • offload_parameter_fwd rewrites the dc.allgather_param nodes zero3_compile inserts; without
    it the pass silently matches nothing.
  • The optimizer-state passes dereference a module global set only from init_z3.py, so naming one
    in a ZeRO-1/2 schedule validates today and then dies on None.
  • opt_states_evicted records a dependency previously implicit in schedule order: move_opt_states
    plans from profiled peaks, which only describe the run once offload_adam_states_for_init has
    taken the optimizer state off the accelerator. move_opt_states_sync reads neither
    profiling_results nor mem_budget and so takes no such requirement.

Conflicts are declared on one side only, which the fix above is what makes reliable.

Naming

Two modules serve stages 1 and 2 but were named for stage 1, matching neither
deepspeed/runtime/zero/stage_1_and_2.py nor the zero2_compile registration:
zero1_compile.pyzero_1_and_2_compile.py, and init_z1.pyinit_z1_and_2.py
(init_z1()init_z1_and_2()). Pure renames, recorded by git as such. Constants become
NAME_Z1 / CONTRACT_Z1 beside the existing _Z2 pair; registered names stay per-stage.

Testing

tests/unit/compile/test_pass_contract.py, 11 tests to 20, all CPU-only. Covers the conflict fix
in both orderings, each new rejection, schedules written with callables rather than names, and that
every schedule init_z1_and_2 and init_z3 build still validates. One guard test asserts every
name in a built-in conflicts_with is itself registered, since those are string literals.

BUILTIN_PASSES mirrors the engine's registration block by hand — reaching the real one needs a
constructed engine — so a pass added to the engine and not the test would go untested.

Compatibility

Registering the remaining passes widens what a schedule may name and narrows nothing. The new
conflicts reject only combinations that already failed at runtime, later and less legibly.

One change genuinely narrows: move_opt_states without offload_adam_states_for_init is now a
PassContractError where it previously ran. Such a run did not fail outright — it planned its
offloading from peaks inflated by the resident optimizer state. Turning that into a schedule-time
error is the intent, but it is the one thing here that rejects something that used to execute.

Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
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