DeepSpeed

35c8b03c - Make pass contracts cover the passes DeepCompile actually schedules; Renaming files (#8251)

Make pass contracts cover the passes DeepCompile actually schedules; Renaming files (#8251) # [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: ```python 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.py` → `zero_1_and_2_compile.py`, and `init_z1.py` → `init_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>