Repository navigation
fix(training): reclaim fused-optimizer activations per step (#1624, #1640) - #1641
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Warning Review limit reached
More reviews will be available in 2 minutes and 9 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, the refill rate gradually slows as usage increases. The highest same-day bursts are limited more strictly. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
Walkthrough
ChangesFused Optimizer Arena Boundedness Fix
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 7590-7593: Replace the thread-static GetFusedStepCount() check
with the per-instance _fusedTrainingCommitted flag to gate the TensorArena
creation, since GetFusedStepCount() is shared across threads and can cause
incorrect arena allocation for multi-model scenarios. Additionally, restructure
the arena allocation from a using statement spanning the entire method body to
one that creates the arena conditionally based on _fusedTrainingCommitted and
disposes it immediately after TryStepWithFusedOptimizer completes in a finally
block, ensuring the transient arena is released before any fallback code paths
like TrainWithTapeStreaming can execute.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: a4d41dc2-0a0e-4472-87e6-657f83b603eb
📒 Files selected for processing (2)
src/NeuralNetworks/NeuralNetworkBase.cstests/AiDotNet.Tests/UnitTests/NeuralNetworks/FusedTrainingArenaBoundednessTests.cs
3935e9b to
915b642
Compare
|
Deployment failed with the following error: Learn More: https://vercel.com/franklins-projects-02a0b5a0?upgradeToPro=build-rate-limit |
ooples
changed the title
…1624/#1640) The default training route for a float model with plain Adam/SGD is the fused-optimizer path (TryTrainWithFusedOptimizer). Unlike the eager tape and streaming paths it opened no per-step TensorArena, so every step ran under whatever arena the caller had open. When a caller wraps the whole loop in ONE un-Reset outer arena (the shared ModelFamilyTests base, or any cross-step buffer pooling), each step's forward/backward activations piled into that outer arena's ring and accumulated -> linear growth -> OOM. On ViT-Base this was +403 MB/step (#1640's 100 GB over a long run); on SimCSE ~290 MB/step (#1624). Fix: open a per-step arena around the fused step, matching how the eager and streaming paths already scope -- and how PyTorch bounds training memory. PyTorch's caching allocator returns each iteration's freed blocks to a reuse pool so the process reaches a steady state, while optimizer state persists outside the per-iteration churn. The arena's _persistent pool is that block cache; Dispose returns the step's transients for the next step to reuse. The compiled plan's Adam m/v moments + persistent input/target are GC-heap-held by the plan's own references (not arena-ring allocated), so the scope recycles only transients and leaves optimizer state intact -- exactly like PyTorch's optimizer.state. Unconditional (every step, including the first), so there is no one-time leak. The arena is disposed in a finally immediately after the fused step returns, so it never spans the result handling or the OOM -> TrainWithTapeStreaming fallback in the else-if (which degrades to the memory-bounded path precisely because we are out of memory -- holding the failed step's transient ring across that switch would waste memory under pressure and nest streaming inside a transient scope). Verified: ViT-Base (86M) under one outer arena, 40 steps, 40 GB cap -- WITHOUT: +403 MB/step climb (2412->19898 MB, would hit 100 GB at ~step 245); WITH: flat ~2.6 GB from step 1. SimCSE under a 16 GB cap plateaus flat. Fused Adam step-for-step parity + convergence and broader fused suite 29/29 green; ~789 ms/step, no perf regression. The regression test asserts BOTH the post-warmup plateau AND that the memorization loss keeps falling (the persistent Adam moments survive the per-step scope). Closes #1624 Closes #1640 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
915b642 to
259b88c
Compare
This branch was successfully deployed
1 participant
Closes #1624
Closes #1640
Root cause
The default training route for a float model with a plain Adam/SGD optimizer is the fused-optimizer path (
NeuralNetworkBase.TryTrainWithFusedOptimizer). Unlike the eager tape (TrainWithTape) and streaming paths — which each open their own per-stepTensorArena— it opened none, so every step ran under whatever arena the caller had open. When a caller wraps a whole multi-step training loop in one un-Reset outerTensorArena— which the sharedModelFamilyTestsbase does, and which any cross-step buffer pooling does — each step's forward/backward activations piled into that outer arena's ring and accumulated. The cursor only advances, so the resident set grew by one step's working set per step → linear growth → OOM.This is not the streaming path and not an autotuner misfire:
ShouldUseStreamingTraining()correctly keeps these models on the eager/fused path (their weights+grads+Adam-moments footprint is tiny vs. available RAM). The leak was purely the missing per-step scope on the fused route.Fix — the way PyTorch bounds training memory
PyTorch reaches a steady state because its caching allocator returns each iteration's freed blocks to a reuse pool (memory never grows with step count), while optimizer state (
exp_avg/exp_avg_sqinoptimizer.state) is created once and persists outside the per-iteration churn.This change applies that exact model: open a per-step
using var arena = TensorArena.Create()around the fused step. The arena's cross-arena_persistentpool is the block cache —Disposereturns the step's transient activations/gradients to it for the next step to reuse. It's unconditional (every step, including the first), so there's no one-time leak.Safe because the compiled plan's persistent state — Adam m/v moments + persistent input/target — is held by the plan's own managed references on the GC heap, not arena-ring allocated, so the per-step scope recycles only transients and leaves optimizer state intact (mirroring
optimizer.statepersisting across iterations). Verified by fused Adam step-for-step parity + a convergence assertion.Validation
#1640 — the named victim, VisionTransformer (ViT-Base, ~86M params,
[1,3,224,224]), one outer arena, 40 steps, Release, 40 GB cap — clean A/B:Without the fix it climbs linearly ~403 MB/step, which extrapolates to #1640's observed 100 GB at ~step 245 — matching its 300 s multi-step window. So #1640's blow-up is the same cross-step arena accumulation, not a within-step activation peak; this one fix resolves both issues.
#1624 — SimCSE under
DOTNET_GCHeapHardLimit=0x400000000(16 GB): live heap plateaus flat across 10 / 20 / 30 steps (was 2.73× and climbing → OOM).Correctness / no regression: fused Adam-convergence parity + LR-schedule mapping + compiled-step suite 16/16, broader fused/transformer-training suite 29/29 green; ~789 ms/step (no perf regression). The added regression test
FusedTrainingArenaBoundednessTestsasserts both the post-warmup heap plateau and that the memorization loss keeps falling (so a fix that bounded memory by corrupting Adam's moments would fail it).Builds against master's pinned
AiDotNet.Tensors(0.86.1) — no package bump required.🤖 Generated with Claude Code