DeepSpeed
b048a7dc - [DeepCompile] Add lightweight pass contracts for optimization passes (#8139)

Commit
45 days ago
[DeepCompile] Add lightweight pass contracts for optimization passes (#8139) ### Summary Based on the DeepCompile efficiency and robustness track in the Q3 2026 roadmap #8104 , this PR implements: _Formal pass contracts and validation of optimization passes: Add lightweight optimization pass contracts for automatic compatibility validation and ordering_ ### Motivation DeepCompile passes execute in the order given by the schedule (`init_schedule`, `backend.py`), with no validation of that order. Inter-pass dependencies are implicit: `prefetch` and `selective_gather` rewrite the all-gather/release ops inserted by `zero3_compile` and require it to run first (cf. the default schedule in `init_z3.py`). An invalid order currently surfaces as a downstream failure rather than a diagnostic at schedule time. ### Changes - `deepspeed/compile/passes/contract.py`: `PassContract` dataclass (`provides` / `requires` / `conflicts_with` / `phase`), a contract registry, and `validate_schedule()`, which raises `PassContractError` on the first unmet requirement or conflict. No torch dependency. - Contracts declared for the four built-in passes: `zero3_compile` provides `z3_gather_release`; `prefetch` and `selective_gather` require it. - `register_compile_pass(name, fn, contract=None)`: optional `contract` argument, backward compatible. - `validate_schedule()` invoked on user-supplied schedules in the compile path (`engine.py`). ### Design - Passes without a contract are unconstrained; existing and custom passes are unaffected. Internal default schedules are not validated, only user-supplied ones. - Dependencies are expressed as capability tags (`z3_gather_release`) rather than pass names, so passes can be added or swapped without editing dependents. - Conflicts are symmetric: either pass may declare the incompatibility, and both orderings are rejected. - Scope is limited to metadata and validation, not dependency resolution or automatic reordering. ### Testing - `tests/unit/compile/test_pass_contract.py`: unit tests for valid ordering, missing requirements, same-step providers, symmetric conflicts, uncontracted-pass pass-through, and callable resolution. --------- Signed-off-by: Jahnavi Yelamanchi <jahnaviyelamanchi03@gmail.com>
Parents
Loading