Account for grouped checkpoint, head, and enclosing forward buffers by bradhilton · Pull Request #898 · OpenPipe/ART

bradhilton · 2026-09-15T08:07:56Z

TrainerRank can admit a forward even though tensors retained through backward and temporary buffers exceed its estimate. This PR accounts for grouped activation-checkpoint buffers, routed experts, loaded LoRA adapters, output heads, and enclosing forward outputs. It can select smaller microbatches or refuse work that the old estimate admitted.

Current draft head: 05e8ebf32dd7edb6d3e7793b552e24f32aecf054. Its parent merges main 5b7acf06 without changing the original PR patch. The #927 trusted-profile/cached-width logic and this PR’s group-row/head-workspace pricing both survive the independently checked composition. Integration CI exposed three failures in the new cold-signature test: a monkeypatch returned a raw string where the estimator needs a typed slot reference. This successor deletes only that three-line monkeypatch, using the real resolver for the fixture’s already-registered slots; all assertions, parameters, and the entire production source tree are unchanged. The original 3 failures and 1,103 passing tests remain recorded. Quality checks and hosted two-H200 validation now pass for this exact resulting tree. Three independent agents cleared the final source correction. The earlier workload-specific GPU results below remain historical evidence.

Reference-model buffers stay charged alongside gradient work, and cached estimates distinguish adapter layouts and gradient modes. When the cheap estimate lacks enough information, the planner builds an exact execution plan. The latest component adds a logical checkpoint-input-gradient allowance for the supported BF16, full/uniform, one-layer recomputation configuration. It combines component and empirical estimates with a maximum, keeps backward-only bytes out of forward retention, and sums eligible child allowances without changing split execution order. This is a partial memory model, not a bound on every possible backward allocation.

Runtime changes are limited to art.trainer_rank. Public APIs, art.megatron, forward/backward operators, losses, numerical tolerances, dependency pins, and the existing physical-memory/cache-release policy are unchanged. Tests and CI routing are also updated.

The compatibility correction introduced at 935e0e3aa455801d65489a6ff8e042088d091be2 is retained. The constructor now distinguishes a deliberately unsupported MoE layout from a previously supported cache that was subsequently corrupted. Unsupported capacity/hook configurations retain the existing generic admission fallback; invalid negative/noninteger caches and a supported coefficient changed to zero still fail. This compatibility correction does not change supported-layout arithmetic.

Validation:

The new regression fails in both constructor-declined cases on unchanged 8b3e1dfb, then all nine focused cases pass with the correction, including invalid-cache rejection and supported-layout arithmetic. Ruff lint and formatting checks pass for all three changed files. At exact head 935e0e3aa, three agents cleared the source correction, quality checks passed, and hosted two-H200 validation passed. These CI results do not replace the workload-specific numerical and near-budget qualification described below. Earlier import, fixture, and one-byte test-oracle failures remain retained; no runtime arithmetic or numerical tolerance was relaxed to pass them.
The following GPU measurements were collected before this compatibility correction. They remain evidence for the supported fixture at the recorded earlier source, rather than a new GPU claim for the current head.
The current diagnostic package passed actual worker-image CPU qualification, source guards, and fresh by-value restoration before its GPU run.
A completed one-H200 cold forward/backward with eight requests and 50,640 packed rows recorded a conservative allocation-counter increase of 29.450 GiB against 30.471 GiB predicted, a 1.021 GiB margin. A private diagnostic reads the existing peak before the original reset and joins it with the immediate post-backward peak; it adds no reset. This supplies a conservative counter bound across the previous observation gap for this fixture, conditional on the audited reset sequence and no other counter owner. The endpoint is unsynchronized, the bound includes earlier counter history and loss construction, and it excludes later gradient spooling. It does not establish physical-library headroom or general backward safety. No optimizer update or numerical comparison was performed in this diagnostic.
A later diagnostic at public head 935e0e3a, with the separately reviewed private scalar observer, completed one cold forward/backward with sixteen requests and 101,280 packed/logical rows. The conservative allocation-counter increase was 58.661 GiB against 60.919 GiB predicted and 67.725 GiB usable, a 2.258 GiB prediction margin. This is the same conditional counter measurement described above, not physical-library or numerical qualification. The overall attempt still failed: a retained pre-kubelet exec HTTP 504 interrupted final result transfer, and the final copy/acknowledgment did not complete. Original native/container exits remain 1; all exact resources and recorded host processes were independently closed. The duplicated retained counter records describe one backward, zero optimizer updates. The narrow collector correction has since passed its focused mock controls; its live adoption is pending.
A separate 24-view no-gradient test, using retained canonical views and three resident slots rather than replaying the original 052 checkpoint/prefetch history, consumed all 538,616 logical rows in both cold and repeated calls without a recorded OOM or admission refusal. It nevertheless failed numerical qualification: 31,835 of 49,152 pooled hidden-state elements differed beyond the unchanged tolerance. The planner selected different groupings across the two calls; neither a packing cause nor a regression from this PR is established. Original task/container exit 1 is preserved. Source review found no concrete caller indexing/reference-copy defect; investigation continues under the no-gradient issue Investigate large hidden-state differences when independent requests are co-packed (Qwen3.6, no-grad) #901.
Earlier matched main/candidate and observer-off backward diagnostics also failed strict pointwise repeatability checks while forward hashes/losses and separate relative-L2 checks passed. Those failures remain unresolved under Investigate output and gradient variability between matched same-source Qwen TrainerRank actors #902; they are not waived by the memory results. No forced garbage collection or tensor-copy workaround is added by this PR.

A newer restored-state qualification attempt stopped before the first actor identity callback, optimizer update or target backward. The local actor rejected its preloaded callback with HTTP 400. Fresh-process CPU controls reproduce an import-order cycle in the private test helper; relocating its existing import preserves all function bodies and module ownership and passes the controls. The original HTTP response did not retain the underlying exception, so this remains a strongly supported diagnosis rather than a recovered worker traceback. The private correction is being packaged and qualified separately; it does not change this ART PR or supply a new memory result. The failed native/container exits remain 1, and all four exact Kubernetes resources plus recorded host processes are closed.

The retained GPU attempts have independently reconciled resource and host-process cleanup. The earlier raw failures and detailed measurements remain retained. This PR remains a draft with a significant-admission-behavior merge hold. It does not close #848/#870, the numerical investigations, or total-overhead qualification.

bradhilton