Skip to content

perf(training): schedule training graphs to minimise live memory - #49

Open
runwangdl wants to merge 1 commit into
develfrom
feat/memory-minimising-scheduler
Open

perf(training): schedule training graphs to minimise live memory#49
runwangdl wants to merge 1 commit into
develfrom
feat/memory-minimising-scheduler

Conversation

@runwangdl

Copy link
Copy Markdown
Owner

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.

Describe the intent of your PR here.

Added

Changed

Fixed

PR Merge Checklist

  1. The PR is rebased on the latest devel commit and pointing to devel.
  2. Your PR reviewed and approved.
  3. All checks are passing.
  4. The CHANGELOG.md file has been updated.
  5. If the docker was modified, change back its link after review.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant