DeepSpeed

92269f43 - Fix AutoTP + deep compile collectives silently drop when AC is on (#8355)

Fix AutoTP + deep compile collectives silently drop when AC is on (#8355) Fix https://github.com/deepspeedai/DeepSpeed/issues/8357 ## The fix Two independent changes. **1. Walk every graph, then verify the result.** `iter_graphs()` yields the top-level graph and every nested `GraphModule`, so checkpointed regions are covered. The pass then compares the tensor-parallel layers it *reached* in the graph against the ones it gave a collective to, and raises if any are missing. The check matters more than the traversal fix. It catches the next variant of this failure — a matmul traced to an op `_MATMUL_TARGETS` does not list, a future change to `nn_module_stack`, another higher-order op — instead of only the checkpointing case fixed here. Given the failure mode is silent wrong numbers, a loud error is the right default. The check is fail-closed. The filter against the deferred-layer names is applied only when the two naming schemes demonstrably overlap, because intersecting mismatched names would quietly empty the set and turn the check into a no-op — the exact failure it exists to catch. **2. Fail if a deferred layer runs eagerly.** `defer_collectives_to_compiler` is set during `engine.compile()`, but `nn.Module.compile()` is lazy: compilation happens at the first forward. With `TORCH_COMPILE_DISABLE=1` or `torch._dynamo.config.disable` the backend never runs, the model executes eagerly forever with its collectives switched off, and training is silently wrong. The affected layers now assert they are inside a compiled region. ## Tests | Test | Covers | |---|---| | `TestAutoTPCompileActivationCheckpointing` | the bug above; verified to fail against the pre-fix pass | | `test_unhandled_parallel_layer_raises` | a layer reached but not rewritten must stop the compile | | `test_deferred_layer_run_eagerly_raises` | the eager-execution guard | Existing tests in the file are unchanged and still pass (2x/4xH200, `-m sequential`: 6 passed). The regression test was checked against the pre-fix pass to confirm it actually catches the bug: with `iter_graphs`/the coverage check reverted and the new tests kept, it fails. The magnitude of the divergence in that configuration (forward off by 2.7e-01) is from the standalone reproducer above rather than the test itself, since the distributed-test harness reports only the worker exit code. Note that these tests, like all the AutoTP compile tests, are `@pytest.mark.sequential`, which `tests/pytest.ini` excludes by default. They need the marker to run at all: ```bash pytest tests/unit/compile/test_tp_compile.py -m sequential ``` ### With this PR applied The same comparison, run against pristine `master` and against this branch **in a single job on one node**, so both arms see identical hardware. Ten Adam steps rather than one, so anything that drifts rather than breaking outright has somewhere to show up. Reported as the worst per-step difference against the module-level AutoTP path: | config | | `master` (ba3246de) | **this PR** | |---|---|---|---| | activation checkpointing **on**, tp=2 | loss | 1.10e-01 **FAIL** | **0.00e+00 PASS** | | | grad norm | 1.05e-01 | **0.00e+00** | | activation checkpointing **on**, tp=4 | loss | 1.42e-01 **FAIL** | **1.19e-07 PASS** | | | grad norm | 9.67e-02 | **5.96e-08** | | activation checkpointing **off**, tp=2 | loss | 0.00e+00 PASS | 0.00e+00 PASS | | | grad norm | 0.00e+00 | 0.00e+00 | --------- Signed-off-by: pengdurice <pengduhit@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>