perf(training): schedule training graphs to minimise live memory - #49
Open
runwangdl wants to merge 1 commit into
Open
perf(training): schedule training graphs to minimise live memory#49runwangdl wants to merge 1 commit into
runwangdl wants to merge 1 commit into
Conversation
The training entry points passed `_mockScheduler`, which emits nodes in the order the exporter produced them -- roughly the whole forward pass, then the whole backward pass. That keeps every saved activation live across the entire backward pass even where its gradient consumer runs early, and the arena has to be sized for the resulting peak. Add `_memoryMinimisingScheduler`, the standard greedy list schedule: among the ready nodes, run the one maximising (bytes of inputs whose last remaining use is this node) - (bytes it allocates), ties broken on the original index. The result is always a valid topological order, so this changes when tensors are live, never what is computed; `InPlaceAccumulatorV2` is charged no allocation because it writes through to its accumulator. Measured on GAP9, --defaultMemLevel L3, L3 peak, address-deduplicated: ResNet8 1420 KB -> 1316 KB (-7.3%) MobileNetV1 3903 KB -> 3667 KB (-6.0%) CCT 3183 KB -> 3151 KB (-1.0%) with no rematerialisation and no change to the arithmetic. CCT gains little because 2317 of its 3151 KB is persistent (frozen weights, bias, positional embedding, optimiser state, input/label), which ordering cannot touch. Six schedulers were compared per model. An eager-free variant ties this one on ResNet8 and MobileNetV1 but does nothing on CCT, so they are not interchangeable and this is the one that wins on all three. `_tensorSize` assumes 4 bytes per element rather than reading `tensor.dtype`. That is deliberate and is documented at the constant: it is the weight of a tie-break heuristic, not an allocation, and CCT carries tensors whose ONNX dtype is UNDEFINED, which surfaces as 0 and makes `np.dtype()` raise. Verified end to end on gvsoc, not merely compiled: ResNet8 Errors: 0 (49.28M train cycles, unchanged from 49.31M) and MobileNetV1 Errors: 0 (52.87M). `_mockScheduler` is left in place for the inference entry points. Co-Authored-By: Claude Opus 5 <noreply@anthropic.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.
The training entry points passed
_mockScheduler, which emits nodes in the order the exporter produced them -- roughly the whole forward pass, then the whole backward pass. That keeps every saved activation live across the entire backward pass even where its gradient consumer runs early, and the arena has to be sized for the resulting peak.Add
_memoryMinimisingScheduler, the standard greedy list schedule: among the ready nodes, run the one maximising (bytes of inputs whose last remaining use is this node) - (bytes it allocates), ties broken on the original index. The result is always a valid topological order, so this changes when tensors are live, never what is computed;InPlaceAccumulatorV2is charged no allocation because it writes through to its accumulator.Measured on GAP9, --defaultMemLevel L3, L3 peak, address-deduplicated:
ResNet8 1420 KB -> 1316 KB (-7.3%)
MobileNetV1 3903 KB -> 3667 KB (-6.0%)
CCT 3183 KB -> 3151 KB (-1.0%)
with no rematerialisation and no change to the arithmetic. CCT gains little because 2317 of its 3151 KB is persistent (frozen weights, bias, positional embedding, optimiser state, input/label), which ordering cannot touch.
Six schedulers were compared per model. An eager-free variant ties this one on ResNet8 and MobileNetV1 but does nothing on CCT, so they are not interchangeable and this is the one that wins on all three.
_tensorSizeassumes 4 bytes per element rather than readingtensor.dtype. That is deliberate and is documented at the constant: it is the weight of a tie-break heuristic, not an allocation, and CCT carries tensors whose ONNX dtype is UNDEFINED, which surfaces as 0 and makesnp.dtype()raise.Verified end to end on gvsoc, not merely compiled: ResNet8 Errors: 0 (49.28M train cycles, unchanged from 49.31M) and MobileNetV1 Errors: 0 (52.87M).
_mockScheduleris left in place for the inference entry points.Describe the intent of your PR here.
Added
Changed
Fixed
PR Merge Checklist
develcommit and pointing todevel.CHANGELOG.mdfile has been updated.