Make pass contracts cover the passes DeepCompile actually schedules; Renaming files - #8251
Draft
pengdurice wants to merge 4 commits into
Draft
Make pass contracts cover the passes DeepCompile actually schedules; Renaming files#8251pengdurice wants to merge 4 commits into
pengdurice wants to merge 4 commits into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[DeepCompile] Make pass contracts cover the passes DeepCompile actually schedules
Follow-up to #8139. Three gaps kept
PassContractfrom firing on real schedules:conflicts_withwas ignored whenever the other pass had no registered contract.offload_parametersand the ZeRO-1/2 reduce passescould not be validated at all.
if schedule is None:ininit_z3.py, so a user-supplied schedule combining both was accepted.The conflict fix
validate_schedulematches conflicts againstapplied, but an uncontracted pass hitcontinuebefore
applied.append(name), so a conflict naming it was missed in both orderings. Resolving amissing contract to a shared empty one instead of branching around it fixes that and removes two
lines of control flow:
The pass stays unconstrained; it is simply visible to conflicts others declare.
test_conflict_is_symmetricmissed this because it registers the second pass with an empty contract, which still lands in the registry.
Contracts
zero1_compileadd_z1_reducezero3_compilezero2_compileadd_z2_reducezero3_compile,zero1_compileoffload_parametersoffload_parameter_fwdz3_gather_releaseoffload_adam_statesmove_opt_statesopt_states_evicted, conflicts with the three belowoffload_adam_states_syncmove_opt_states_syncoffload_adam_statesoffload_adam_states_for_initoffload_adam_states_for_initopt_states_evicted, conflicts with the three belowThe three shared conflicts are
offload_parameters,zero1_compile,zero2_compile.offload_parameter_fwdrewrites thedc.allgather_paramnodeszero3_compileinserts; withoutit the pass silently matches nothing.
init_z3.py, so naming onein a ZeRO-1/2 schedule validates today and then dies on
None.opt_states_evictedrecords a dependency previously implicit in schedule order:move_opt_statesplans from profiled peaks, which only describe the run once
offload_adam_states_for_inithastaken the optimizer state off the accelerator.
move_opt_states_syncreads neitherprofiling_resultsnormem_budgetand 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.pynor thezero2_compileregistration:zero1_compile.py→zero_1_and_2_compile.py, andinit_z1.py→init_z1_and_2.py(
init_z1()→init_z1_and_2()). Pure renames, recorded by git as such. Constants becomeNAME_Z1/CONTRACT_Z1beside the existing_Z2pair; registered names stay per-stage.Testing
tests/unit/compile/test_pass_contract.py, 11 tests to 20, all CPU-only. Covers the conflict fixin both orderings, each new rejection, schedules written with callables rather than names, and that
every schedule
init_z1_and_2andinit_z3build still validates. One guard test asserts everyname in a built-in
conflicts_withis itself registered, since those are string literals.BUILTIN_PASSESmirrors the engine's registration block by hand — reaching the real one needs aconstructed 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_stateswithoutoffload_adam_states_for_initis now aPassContractErrorwhere it previously ran. Such a run did not fail outright — it planned itsoffloading 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.