fix(#1221): Transformer.Predict bypassed eval-mode wrapper, Dropout fired at inference - #1242
Merged
Merged
Conversation
…so eval-mode wrapper applies Root cause: Transformer<T>.Predict() overrode NeuralNetworkBase.Predict() directly, which meant it bypassed the base's wrappers introduced in 6bf657c (fix(Predict+Dense)): - SetTrainingMode(false) for the duration of the forward pass - NoGradScope to suppress autograd recording - NormalizeInputBatchDim for unbatched inputs - Output squeeze when input was promoted For a token-input Transformer (Phase A4 byte-LM config: V=256, d=64, L=2, ctx=32) the layer chain contains ~5 DropoutLayer instances (one post-embedding, two per encoder block). With training mode left on at inference time — which happened in any caller that didn't manually call SetTrainingMode(false) before Predict, including any code path that re-enabled training between train and eval — Dropout's stochastic mask reroll fired on every Predict call. After many training steps the trained Dense output projection settled into a configuration where the random Dropout masks at eval time pushed logits toward uniform-over-V (eval NLL = ln(V) exactly, top-1 = 0%). Fix: Transformer's special-case routing (DecoderLayer cross-attention, AttentionLayer mask, MultiHeadAttention encoder-output tracking) is moved from Predict() to PredictEager(). The base's Predict() now wraps the override with its standard inference scaffolding — eval mode, no-grad, batch promote, output squeeze — and Dropout becomes a no-op at inference automatically. Test: Predict_WithExternalTrainingModeOn_StillDeterministic forces training mode ON between Train and Predict, then verifies two Predict calls on the same input produce bit-identical logits. Pre-fix this test fails with max abs diff = 9.84e-3 (Dropout firing); post-fix diff < 1e-6 (Dropout properly disabled). Verification: broader Transformer test set went from 412 → 370 failures (42 tests recovered, no regressions). Remaining 370 are pre-existing #1224 cluster failures unrelated to this fix. Co-Authored-By: Claude Opus 4.7 (1M context) <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 (2)
WalkthroughTransformer moves its forward-pass override from the public ChangesTransformer Prediction Architecture
Sequence Diagram(s)(omitted — changes are internal control-flow adjustments limited to transformer forward-pass and a unit/integration test; conditions for diagram generation not met) Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Blocking code-quality notes
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)
Review rate limit: 2/5 reviews remaining, refill in 33 minutes and 31 seconds. Comment |
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Fixes #1221 where Transformer<T>.Predict() bypassed base inference wrappers (eval-mode toggle/no-grad/batch normalization), causing Dropout to run during inference and produce stochastic/uniform logits.
Changes:
- Moves Transformer’s special-case forward routing from
Predict()intoPredictEager()so basePredict()wrappers are applied. - Adds a regression test ensuring
Predict()remains deterministic even if callers force training mode on before inference.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/NeuralNetworks/Transformer.cs |
Routes Transformer’s eager forward via PredictEager() to preserve NeuralNetworkBase.Predict() inference scaffolding (fixes Dropout-at-inference). |
tests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/TransformerProductionScaleConvergenceIssue1221Tests.cs |
Adds regression test covering the “external training mode re-enabled before Predict” scenario. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![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 the current code and only fix it if needed.
Inline comments:
In
`@tests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/TransformerProductionScaleConvergenceIssue1221Tests.cs`:
- Around line 405-436: Add an assertion that Predict restores the caller's
training-mode state: after the two Predict calls (which occur after
model.SetTrainingMode(true) and the existing checks), add
Assert.True(model.IsTrainingMode) to verify the model's training-mode flag was
preserved; reference the existing methods model.SetTrainingMode(true),
model.Predict(...), and model.IsTrainingMode to locate where to insert this
post-call check so the test fails if Transformer.Predict or any override fails
to restore the prior mode.
🪄 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: ad610159-540b-4405-b2be-d6bd7bb3f8c1
📒 Files selected for processing (2)
src/NeuralNetworks/Transformer.cstests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/TransformerProductionScaleConvergenceIssue1221Tests.cs
…ger doc - TransformerProductionScaleConvergenceIssue1221Tests: Predict_WithExternalTrainingModeOn_StillDeterministic now asserts the post-Predict state restores caller's IsTrainingMode (locks in the save/restore contract). Tolerance loosened from 1e-6 to 1e-3 so GPU floating-point nondeterminism (kernel scheduling / reduction-order variance) doesn't make the test flaky on CI agents with GPU-backed Engine. Pre-fix Dropout-fired diff is ~1e-2 so 1e-3 still cleanly distinguishes the bug from float noise. - Transformer.cs PredictEager remarks scoped to encoder/decoder routing only — the wrapper-level "why" (training mode toggle, no-grad, batch-promote) belongs on Predict's docs, not on this protected override. Cross-references back to Predict for the inference-scaffolding contract. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/NeuralNetworks/Transformer.cs`:
- Around line 475-490: PredictEager currently captures encoderOutput only once
(when null) which records the first MultiHeadAttentionLayer<T> output; change
the logic to capture the last such encoder-side attention output before any
decoder layer by removing the "only-if-null" guard and instead track whether a
DecoderLayer<T> has been seen (e.g., bool decoderSeen = false), set decoderSeen
= true when you hit a DecoderLayer<T>, and whenever you encounter a
MultiHeadAttentionLayer<T> and decoderSeen is false assign/overwrite
encoderOutput with that layer's output so the final encoder block's attention
output is retained for cross-attention.
In
`@tests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/TransformerProductionScaleConvergenceIssue1221Tests.cs`:
- Around line 387-455: The test relies on BuildTransformer defaults for dropout;
make the dropout precondition explicit by passing a non-zero dropout rate into
the BuildTransformer call (e.g., dropoutRate: 0.1) or, alternatively, assert
that the constructed model contains an enabled Dropout layer before using
pred1/pred2; update the test's model construction line (the BuildTransformer
invocation) or add a pre-comparison check that inspects the model's layers for a
Dropout instance with DropoutRate > 0 to guarantee the regression is actually
exercised (this targets the BuildTransformer usage and the Transformer.Predict
path referenced in the comment).
🪄 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: 8a36bdea-5620-4015-81fd-f6fe6d5b9e3e
📒 Files selected for processing (2)
src/NeuralNetworks/Transformer.cstests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/TransformerProductionScaleConvergenceIssue1221Tests.cs
…t in determinism test Addresses two reviewer threads on PR #1242: OVkn — Transformer.PredictEager was capturing the FIRST encoder MultiHeadAttention output and freezing it, so a multi-block encoder stack always fed the decoder the first block's output instead of the fully-encoded representation. Replace the once-only `encoderOutput is null` check with a `seenDecoder` flag: keep updating encoderOutput on every encoder MultiHead until the first decoder is reached, then freeze. OVkt — Predict_WithExternalTrainingModeOn_StillDeterministic relied on TransformerArchitecture's framework-default dropoutRate (0.1) to produce stochastic Predict output pre-fix. If that default ever drops to 0.0 the test silently no-ops. Pin dropoutRate=0.2 explicitly via a new BuildTransformer parameter, and add a precondition assertion that the model contains at least one DropoutLayer with SupportsTraining=true before comparing pred1/pred2. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
3 of 5 tasks
1 of 3 tasks
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.
Closes #1221.
Summary
Root cause for #1221's "Transformer training completes but eval logits are uniform-over-V" reproducer:
Transformer<T>.Predict()overrodeNeuralNetworkBase.Predict()directly, bypassing the wrappers introduced in6bf657c45:SetTrainingMode(false)for the forward passNoGradScopeto suppress autogradNormalizeInputBatchDimFor a Phase A4 byte-LM config (V=256, d=64, L=2, ctx=32) the layer chain contains ~5
DropoutLayerinstances. With training mode left on at inference (which happened any time a caller re-enabled training between train and eval), Dropout's stochastic mask reroll fired on every Predict call, pushing the trained output Dense into a configuration where eval logits collapse to uniform-over-V (NLL = ln(V), top-1 = 0%).Fix
Move the special-case layer routing (DecoderLayer cross-attention, AttentionLayer mask, MultiHeadAttention encoder-output tracking) from
Predict()toPredictEager(). The basePredict()now wraps the override with its standard inference scaffolding — eval mode, no-grad, batch promote, output squeeze — and Dropout becomes a no-op at inference automatically.Test plan
Predict_WithExternalTrainingModeOn_StillDeterministicfails pre-fix with max abs diff = 9.84e-3 (Dropout firing), passes post-fix with diff < 1e-6🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests