[DimensionSTP]

Fixes #8874

Problem

ZeRO-3 currently discovers a parameter accessed outside its owning module only when it is NOT_AVAILABLE. A resident weight can therefore be missing from the recorded external dependencies. If it is subsequently partitioned, the same access registers a new dependency after the trace is complete. Checkpoint recomputation then fails because the parameter fetch queue no longer matches the module's dependencies.

I encountered this during full LoRA checkpoint resume using Hugging Face PEFT and TRL, with gradient accumulation and reentrant activation checkpointing. It also reproduces with a small PyTorch adapter and DeepSpeed alone, without either library, on master 68faa67c.

Change

Discover eager external parameter accesses independently of residency, while leaving the existing condition for actually gathering a parameter unchanged. Accesses from the parameter's own module and the separate DeepCompile eager-fallback path retain their behavior.

For ZenFlow, use the existing active-consumer set in the forward and backward hooks so that nested consumers share the compute layout and only the last consumer restores the storage layout. A parameter discovered after the consumer's pre-hook must not receive an unmatched post-hook transpose.

This does not lower the persistence threshold, remove trace checks, or invalidate the trace on every step. Production changes are confined to parameter_offload.py; regression tests use three existing test files. The getter is specific to ZeRO-3; ZeRO-1/2 do not use this offload path.

Validation

Hardware: main1 H100 host, 1 GPU for the ZenFlow tests, 2 GPUs for the resume regression, and 4 for the earlier real-model check. The cloud exposes the device name as CUDA GPU, compute capability 9.0, with 81,559 MiB per device. Driver 535.161.08, PyTorch 2.13.0+cu129, CUDA runtime 12.9, nvcc 12.8. DeepSpeed was built from 68faa67c without prebuilt ops; existing tests may JIT the operators they use.

• The new CPU getter regression fails for external access on unmodified master; both cases pass with the fix.
• The real two-rank checkpoint regression fails on unmodified master with the same adapter/base fetch-queue mismatch. All eight patched cases pass, with no errors or skips on either rank. They cover parameter offload none/CPU, persistent/non-persistent adapters, and reentrant/non-reentrant checkpointing, all with GA 4.
• The regression saves a checkpoint, creates a new engine, and restores optimizer and scheduler state. It checks each microstep's loss and the weights after both optimizer updates against a non-checkpointed plain PyTorch reference.
• All eight resume cases were rerun after the ZenFlow follow-up and still pass on both ranks. The CPU getter and eager-fallback/guard checks also pass: eight cases, without errors or skips.
• Six real-engine ZenFlow cases pass against plain PyTorch: persistent/non-persistent matrices, with no external access, a metadata-only access, or an external access followed by the owning Linear. They check loss, transposed storage values, and weights after three optimizer updates using FP32, synchronous CPU optimizer offload, and the warmup phase. The nested-owner case fails in backward with the previous ZenFlow hooks; a metadata-only control also passes there. The existing BF16 ZeRO-3 ZenFlow selective-update training test passes as well.
• The standalone issue reproducer fails on the original and completes full resume on the patched source.
• The existing attribute-delegating wrapper and dangling external-parameter training tests pass. Six existing CPU eager-fallback/guard checks also pass.
• Before the ZenFlow follow-up, the original 27B checkpoint resumed from global step 2 to 6 and completed two evaluations with stock PEFT 0.21.2 and stage3_param_persistence_threshold=1000000. This used BF16 base weights, FP32 LoRA, GA 4, reentrant checkpointing, and TRL fused NLL, with Liger and ZenFlow disabled. No PEFT source or training setting was changed to make it pass. The 27B run was not repeated for the ZenFlow-only follow-up.
• All applicable project pre-commit hooks pass for the four changed files.

The small regression is in tests/unit/v1/zero/test_zero_activation_checkpoint_lifecycle.py::TestZero3ActivationCheckpointLifecycle::test_resume_with_external_persistent_adapter. Locally I ran the official test bodies and fixtures through an external torchrun launcher, rather than the fork-based launcher. No model, optimizer, or collective was mocked. The local commands were:

python ../run_official_tests.py --case final_official_fixed
python ../run_official_tests.py --case standalone_fixed --selection standalone_reproducer
python ../run_official_tests.py --case review_resume_final
python ../run_official_tests.py --case review_zenflow_final --world-size 1 --selection tests/unit/runtime/zenflow/test_zf.py::TestZenFlowExternalParameter
python ../run_official_tests.py --case review_zenflow_existing_selective --world-size 1 --selection 'tests/unit/runtime/zenflow/test_zf.py::TestZenFlowSingleGPU::test_zenflow_single_gpu[step-10-3-False-0-3]'
CUDA_VISIBLE_DEVICES='' DS_ACCELERATOR=cpu python -m pytest -q tests/unit/runtime/zero/test_zero_context.py -k available_parameter_access
CUDA_VISIBLE_DEVICES='' DS_ACCELERATOR=cpu python -m pytest -q tests/unit/v1/compile/test_z3_eager_fallback.py -k 'guard or still_gathers_outside_guard or releases_leftover'
pre-commit run --files deepspeed/runtime/zero/parameter_offload.py tests/unit/runtime/zero/test_zero_context.py tests/unit/runtime/zenflow/test_zf.py tests/unit/v1/zero/test_zero_activation_checkpoint_lifecycle.py

The small resume and ZenFlow regressions peak at approximately 1.93 GiB allocated on rank 0. They are correctness checks, not a throughput benchmark or a long-run convergence claim; the failing baseline cannot complete the corresponding training run. NVMe, optimizer CPU offload outside the tested ZenFlow paths, FP16, non-CUDA devices, multi-node/TP/PP/MoE, asynchronous ZenFlow, ZenFlow checkpoint resume, and compiled training were not validated.

Related Work

I checked open external-parameter, LoRA, and trace-related PRs and did not find this fix. #8431 changes prefetch sizing, not dependency discovery. This preserves the frozen-parameter checkpoint support from #8148 and the DeepCompile getter behavior from #8328. #6718 discusses a related getter mechanism under compilation, but this is an eager resume failure, not a claim to fix that issue.

AI Assistance

AI assistance was used during the investigation and preparation of the fix and tests.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-21T11:27:47.545942Z	0380e42	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0380e42e27

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Keep ZenFlow transposes symmetric during registration

With ZeRO-3 ZenFlow enabled, a resident 2-D weight first touched externally inside a consumer's forward (such as the added LoRA dtype lookup) is inserted into that consumer's external set here after pre_sub_module_forward_function() already enumerated the set. post_sub_module_forward_function() then sees the new weight and applies its ZenFlow .t() (parameter_offload.py:562-565) without the matching pre-transpose, leaving the parameter in the compute layout when partitioning/optimizer code expects ZenFlow's transposed storage layout; this can silently corrupt the next update or forward. Snapshot the pre-hook set or otherwise make the registration's transpose accounting symmetric, and cover this configuration.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

[DimensionSTP]

Fixed using the existing active-consumer set in the forward and backward hooks, so nested consumers share the compute layout and only the last consumer restores the storage layout. Added real-engine regressions in the existing ZenFlow test file for resident/partitioned weights and metadata-only/nested access. All six cases match plain PyTorch on H100; the nested-owner case fails with the previous hooks. The eight existing full-resume cases also pass.