feat(recompute): coordinate checkpoint-without-output recomputation

[FDecaYed]

What

Extend CheckpointWithoutOutput so it can safely coordinate several sequential checkpointed operations:

• preserve mixed tensor and non-tensor arguments through recompute;
• return None gradients for metadata arguments;
• add a CheckpointManager that discards all managed outputs and registers one ordered recompute hook;
• preserve CUDA-graph capture behavior;
• add focused unit coverage for argument restoration, hook ordering, storage reuse, gradients, and edge cases.

This reusable recompute primitive is the root of the HybridModel mHC lane, but contains no mHC, model, or GPTModel code itself.

Provenance

• Reconstructs the checkpoint-management portion of the mHC work from [dev] feat(mHC): Add basic pytorch implementation of manifold hyper connection(mHC). #2943 and the current main proposal [Split 1/N of #3430] feat(mHC): basic pytorch implementation of manifold hyper connection #4531.
• Refactors the same implementation captured in [Tracking] DeepSeek-v4 HybridModel merge plan #5795 (commit 07c0d08).
• Original implementation by the [dev] feat(mHC): Add basic pytorch implementation of manifold hyper connection(mHC). #2943/[Split 1/N of #3430] feat(mHC): basic pytorch implementation of manifold hyper connection #4531 authors; DSv4 adaptation by @hxbai.
• Reconstructed against current main; feature cutoff: 2026-07-21.
• GPTModel integration is intentionally excluded.

Testing

• uv run isort --check-only on both changed Python files
• uv run ruff check on both changed Python files
• uv run python -m compileall -q on both changed Python files
• Focused tests are included; local execution requires the project Torch runtime and is left to CI.

Stack follow-up

• feat(transformer): add native manifold hyper-connections #5936 contains the next logical slice and temporarily targets this PR mirror.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[FDecaYed]

/ok to test 18e842a

[FDecaYed]

Closing this draft because the generic checkpoint-without-output/mHC contract is supplied by external prerequisite #4531. The later P2 PR will be a focused HybridModel/fused-mHC delta after #4531 and P1 merge. Tracking PR: #5795.