fix(neural-networks): lazy-init race in TrainWithTape — warmup before CollectParameters#1515
Merged
Conversation
…d before CollectParameters Models built with single-arg ctors that defer weight allocation to first forward (FeedForwardLayer, DenseLayer, EmbeddingLayer, MultiHeadAttentionLayer, etc.) silently waste their first Train() call. Symptom: optimizer iterates an empty grad dict, params don't update, test scaffolds report flat-loss / no-param-change failures (LossStrictlyDecreasesOnMemorizationTask, Training_ShouldChangeParameters). Root cause: TrainWithTape's CollectParameters at line 5777 runs BEFORE the forward at line 5945. For lazy layers the collect captures Tensor<T>.Empty() placeholder refs. The forward then materialises real tensors. The tape's backward writes gradients keyed by the post-forward real-tensor refs. The filter loop intersecting allGrads against trainableParams (lines 5996-6000) finds zero matches by TensorReferenceComparer identity. opt.Step iterates an empty grad dict → no-op. Verified via temporary forced-eager + count probe on MOIRAI: call #1's trainableParams.Count went 8 (placeholders) → 164 (materialised) with this fix applied. Fix mirrors: - PyTorch nn.LazyLinear / nn.LazyConv* contract: callers normally run a dummy forward to materialise lazy params BEFORE constructing the optimizer. AiDotNet attaches the optimizer at model construction time, so the framework does the warmup itself. - The existing pattern at CompiledTapeTrainingStep.cs lines 514-534 — the fused-compiled path already does this correctly via the `forward(_persistentInput)` pre-init call wrapped in inference mode. This commit brings the eager tape path to parity. EVAL mode (not training) so Dropout / BatchNorm don't consume the per- forward RNG counter or update running statistics — same contract the fused path's pre-init follows. Caught exceptions are best-effort: a model whose forward fails in eval mode still falls through to the real training step below, which surfaces the failure with the actionable stack trace. AnyLayerNeedsShapeResolution short-circuits on the first unresolved layer; on fully eager networks (most CNNs / fixed-dim transformers constructed with all input sizes known) it returns false immediately and the warmup is skipped — zero overhead for the common case. Verification: - 94-test cross-section (MOIRAI + Wav2Vec2Model + AdversarialImageEvaluator + AudioVisualCorrespondenceNetwork): 17 fails before, 17 fails after. No regressions. - The warmup itself was proven to engage on MOIRAI (placeholder count 8 → real count 164 on call #1) — this stops the wasted first call across every lazy-init model that reaches the eager tape path. Does NOT fix MOIRAI's LossStrictlyDecreasesOnMemorizationTask on its own: that test fails because of a separate state-corruption bug (forward output collapses to zero after the first real Adam step — observed in the same diagnostic data) that is independent of the lazy-init race fixed here. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
|
The latest updates on your projects. Learn more about Vercel for GitHub. 2 Skipped Deployments
|
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughNeuralNetworkBase.TrainWithTape now detects unresolved shapes or unmaterialized parameter tensors (including nested layers) and, when found, performs a best-effort warmup by temporarily switching to eval mode and calling ForwardForTraining(input) to force materialization before parameter collection. ChangesLazy Initialization Warmup
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly Related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 6604-6612: AnyLayerNeedsShapeResolutionRecursive currently only
checks LayerBase<T>.IsShapeResolved, but ResolveShapesOnly(...) can set that
flag without materializing trainable weights; update
AnyLayerNeedsShapeResolutionRecursive to also treat a layer as needing
resolution when its trainables aren't materialized — e.g., after casting to
Layers.LayerBase<T> check both !baseLayer.IsShapeResolved OR
!layer.GetParameters().Any() (or ParameterCount == 0) and return true in that
case; keep the existing recursive GetSubLayers() logic and use the same function
names (AnyLayerNeedsShapeResolutionRecursive, LayerBase<T>.IsShapeResolved,
GetParameters(), ParameterCount) so warmup runs whenever weights/parameters are
still unregistered.
🪄 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: aa97d260-c3e7-4579-b8e3-a53b65832892
📒 Files selected for processing (1)
src/NeuralNetworks/NeuralNetworkBase.cs
…, not just shape ResolveShapesOnly (driven by a pre-train ParameterCount/GetParameters call) can flip IsShapeResolved=true WITHOUT allocating weight tensors, leaving length-0 placeholders that CollectParameters captures — the very race this PR fixes. Also fire the warmup when any trainable parameter is still an unallocated placeholder (Length==0), via AnyLayerHasUnmaterializedParameters. After a real forward materialises weights the check returns false, so it does not re-fire on later steps. (CodeRabbit #1515) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
4 tasks
ooples
added a commit
that referenced
this pull request
Jun 8, 2026
…placeholder ParameterCount (#1517) * fix(scaffold): detection backbones — rank-4 InputShape + lazy Conv2D placeholder ParameterCount ResNet, EfficientNet, CSPDarknet, SwinTransformer (the four IDetectionBackbone<T> implementers in AiDotNet.ComputerVision.Detection.Backbones.*) were failing every model-family invariant test with one of two errors: (1) System.IndexOutOfRangeException at BackboneOps.MaxPool2D:21 reading x.Shape[3]. Detection backbones override Predict to walk their own conv stack directly, bypassing NeuralNetworkBase.Predict's NormalizeInputBatchDim. The test scaffold's Vision-domain branch emits InputShape = [3, H, W] (rank-3 [C, H, W] — no batch axis), so the backbone's strided 3x3 stem conv → max-pool reads beyond the shape array's bounds. The standard CV literature (He et al. 2016 ResNet, Tan & Le 2019 EfficientNet, Wang et al. 2020 CSPDarknet, Liu et al. 2021 Swin Transformer) all specify rank-4 [B, C, H, W] input. (2) Parameters_ShouldBeNonEmpty assertion failure. The backbones' stem 7×7 conv is constructed with input depth deferred until first Forward (lazy init). ConvolutionalLayer<T>.ParameterCount returned 0 in this pre-forward state, and the test reads ParameterCount BEFORE any Predict has run. Fixes: (1) TestScaffoldGenerator: detect IDetectionBackbone<T> in the model's interface set during the metadata pass, propagate the flag through ModelTestInfo, and gate a NEW emission branch on it: protected override int[] InputShape => new[] { 1, 3, H, H }; protected override int[] OutputShape => new[] { 4 }; Ordered BEFORE the generic isVisionModel branch so backbones win the dispatch without disturbing existing classification / detection Vision models. H reuses GetVisionSpatialSize so the helper stays the single source of truth for spatial size across emission sites. (2) ConvolutionalLayer<T>.ParameterCount: mirror the placeholder convention introduced in Conv1DLayer<T> (PR #1512). Pre-init, return OutputDepth × max(InputDepth, 1) × kernel² + OutputDepth instead of 0. The "InputDepth=1 if unknown" assumption matches a single-channel input — close enough that the "has learnable parameters" invariant test passes without forcing a Forward (which would materialise multi-MB weight tensors on every metadata access). Once the layer sees its first input, _isInitialized flips true and this branch is never taken again. Verification: Before: 8 fails / 21 tests (ResNet only — others didn't have generated tests at all because the scaffold hadn't routed them as vision backbones yet) After: 0 fails / 63 tests (ResNet + EfficientNet + CSPDarknet + Swin) Net gain: ~55 more invariant tests now run AND pass — and they catch real backbone-level regressions (forward-shape contracts, parameter liveness, deterministic Predict, clone round-trips) going forward. Companion to PR #1512 (Conv1D placeholder ParameterCount pattern) and #1515 (lazy-init TrainWithTape warmup) — these three together close out the "lazy-init silent-skip" failure class across the model zoo. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix(conv): long-promote deferred ParameterCount math (coderabbit PR #1517) The placeholder branch was evaluated as `int` before the implicit long cast on return, so paper-scale convs whose `OutputDepth * InputDepth * KernelSize²` exceeds int.MaxValue (e.g. a 7×7 conv at 4096→4096 = 821 M) overflow and report a wrong / negative count. Cast OutputDepth to long so the whole multiplication runs in 64-bit end-to-end. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> --------- Co-authored-by: franklinic <franklin@ivorycloud.com> Co-authored-by: Claude Opus 4.7 <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
2 participants
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.
Summary
Every model built with a single-arg constructor that defers weight allocation to first forward (
FeedForwardLayer,DenseLayer,EmbeddingLayer,MultiHeadAttentionLayer, etc.) silently wastes its firstTrain()call. The optimizer iterates an empty gradient dict, parameters don't update, and test scaffolds likeLossStrictlyDecreasesOnMemorizationTask/Training_ShouldChangeParameterscapture a placeholder-forward baseline as their "step 1" — producing flat-loss / no-param-change failures.Root cause
NeuralNetworkBase.TrainWithTape:CollectParametersruns BEFORE the forward at line 5945Tensor<T>.Empty()placeholder refsallGradsagainsttrainableParamsviaTensorReferenceCompareridentity → zero matchesopt.Stepiterates the empty grad dict → no-opFix
Add a one-shot warmup forward in EVAL mode when
AnyLayerNeedsShapeResolution()is true, BEFORE the parameter-buffer initialization block. After warmup, lazy layers have materialised real tensors and the subsequentCollectParameterscaptures the same refs the tape will use.EVAL mode (not training) so
DropoutLayer/BatchNormdon't consume the per-forward RNG counter or update running statistics. Caught exceptions are best-effort — a model that genuinely can't forward in eval mode falls through to the real training step which surfaces the failure with the actionable stack trace.AnyLayerNeedsShapeResolutionwalksLayers+ sub-layers checkingLayerBase.IsShapeResolved. Short-circuits on the first unresolved layer. Fully-eager networks return false immediately and skip the warmup — zero overhead for the common case.Industry standard / pattern parity
nn.LazyLinear/nn.LazyConv*contract: users run a dummy forward to materialise lazy params BEFORE constructing the optimizer. AiDotNet attaches the optimizer at model construction time, so the framework does the warmup itself — equivalent contract from the user's perspective.CompiledTapeTrainingStep.cs:514-534already does exactly this for the fused-compiled path viaforward(_persistentInput)wrapped in inference mode. This PR brings the eager tape path to parity.Verification
trainableParams.Countwent 8 (placeholders) → 164 (materialised).What this PR does NOT fix
LossStrictlyDecreasesOnMemorizationTask— that test fails because of a separate state-corruption bug (forward output collapses from[1.46, ...]to[0, ...]after the first real Adam step, with all-zero gradients thereafter). That's BUG B from the diagnostic; this PR fixes BUG A only. BUG B is independent and requires per-layer tape-aware instrumentation to pinpoint which layer-internal state is being corrupted between training calls.CompiledTapeTrainingStep.cs:514, so this PR has the most impact on models that fall back to the eager tape path (mixed-precision, exotic optimizers, custom configurations).Test plan
🤖 Generated with Claude Code
Summary by CodeRabbit