Repository navigation
Conversation
Cold compiled FP32 force training can cache a grid-pair launch that exceeds per-block shared memory. wrap_triton records a call during fake-tensor tracing without loading the kernel, so the eager launch loop mistakes a successful trace for a resource-validated configuration. The actual OutOfResources exception arrives later in Inductor's backward compilation, outside the retry loop. Reducing the training batch cannot fix the tile's fixed resource footprint. For preset Pro on RTX PRO 6000 Blackwell, the second-order kernel initially requests 118784 bytes against a 101376-byte device limit. Eager execution finds a smaller grid tile requiring 79872 bytes, but cold graph capture cannot take that retry. Prewarming eagerly can consequently hide the bug. Resolve launch feasibility before recording or executing a kernel: - Separate candidate selection from the runtime launch and use the same specialization parameters for validation and execution. - Compile candidates with Triton warmup using operand dtypes instead of tensor allocations, then load the compiled handles to enforce the target device's resource limits without executing the kernel. - Represent runtime pair counts and physical strides by non-unit scalar placeholders during validation. Keep their actual values in the launch ABI so dynamic batches and degree-major views retain one compiled graph. - Cache validated choices by kernel, device, static geometry, complete operand signature and differentiation-order accumulator count. Cache an unsuccessful search as well, avoiding repeated compilation for layouts served by the reference composition. - Retain exact-device tuned configurations as preferred candidates, while subjecting them to the same validation and deduplicating fallback choices. Forward, first backward and force-loss second backward share the selector. - Preserve compiler-visible Triton operators and the reference einsum path when no candidate fits. Only resource-limit errors reject a candidate; other compilation failures remain visible. Add cold compiled force-loss regressions for the two production Pro grid layouts in FP32 and BF16 autocast. Exercise batch-dependent physical strides with changing node counts, compare outputs and second-order gradients with eager execution, and verify that only one graph is compiled. Clear launch selection state and disable the FX graph cache so prior eager tests cannot mask an invalid first choice. Cover reference fallback through compilation when no fused launch is available. Validation on RTX PRO 6000 Blackwell Server Edition with PyTorch 2.13.0+cu130 and Triton 3.7.1: - All 20 grid-pair tests pass, including existing FP32/BF16 forward and first/second-gradient checks against the float64 reference, dynamic layout tests, and the five cold-compile/fallback regressions. - A cold FP32 second-order reproducer with one graph node passes without an eager warmup. The validated and final compiled backward kernels agree on their resource footprint; the Pro second-order tile uses 79872 bytes. - A preset Pro fine-tuning smoke run completes distributed precompilation and two optimizer steps with one DDP rank, ZeRO Stage 1, a 300-atom batch budget, AMP/TF32 disabled, and both CUDA/Triton training gates enabled. All 278 floating-point model tensors in the saved checkpoint are finite. - An independent review compares 120 actual layout signatures across five geometries, two dtypes, three differentiation orders and four storage layouts. Their shared-memory usage matches the validation signatures. - Ruff and whitespace checks pass. Other GPU architectures and the project-pinned PyTorch 2.10 environment were not exercised.
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing. Review the Copilot code review billing settings.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 64.29% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 14 functions across 2 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: validating Triton grid-pair launches before graph capture. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Agent: dot
Source LGTM at 1351a4f. Reviewed both changed files and the forward/first-backward/second-backward integration, including the no-viable-config reference paths. Resource selection now occurs before recording wrap_triton; tuned choices receive the same validation, only OutOfResources triggers retries, and successful/unsuccessful choices are cached without putting node counts or physical strides into the selector key. Shared launch kwargs keep the static geometry and launch settings consistent. The five added cold-compile cases cover both layouts, FP32/BF16, changing node counts, force-loss differentiation and a forced reference path.
Bounded independent checks: exact-source helper extraction with test doubles passed preferred-tile rejection at handle loading, positive/negative caching, shared validation kwargs, propagation of unrelated compiler errors, and slot splitting. Also inspected Triton 3.6.0 upstream warmup and _init_handles implementations for the dtype-only compilation/resource-check interfaces. These checks are not GPU execution or an end-to-end pinned-environment qualification.
Reusing the PR's reported fresh-process 20-test Blackwell / PyTorch 2.13 / Triton 3.7.1 results rather than requesting another full GPU run. My executor has no PyTorch installed; I did not independently reproduce GPU numerics, pinned PyTorch 2.10 behavior, other architectures or throughput. Those remain explicitly unverified, not newly observed failures. Current CUDA CI is successful; Python and C++ test workflows are still running. No actionable source defect found in this scope.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Agent: dot
Qualification clarification for my review above: the Pass testing on CUDA aggregate is successful, but the actual Test Python on CUDA and Test C++ on CUDA jobs were skipped. Thus this CI result is not an independently executed GPU test. GPU evidence remains the author's reported Blackwell results, with the independent-validation limitations already stated. The source LGTM is unchanged.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6095 +/- ##
==========================================
- Coverage 78.08% 77.78% -0.31%
==========================================
Files 1171 1171
Lines 141661 141677 +16
Branches 5079 5078 -1
==========================================
- Hits 110620 110199 -421
- Misses 29103 29540 +437
Partials 1938 1938 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Summary
Validate Triton grid-pair launch resources before graph capture so cold compiled force training can select a feasible tile or retain the reference composition.
Previously,
wrap_tritoncould record a call during fake-tensor tracing without loading its kernel. The launch retry loop treated successful tracing as successful resource validation and cached an invalid tile. The actual shared-memory error then arrived during Inductor backward compilation, outside the retry loop. An eager warmup could hide this failure.Implementation
warmupinputs and load their handles to check device resource limits without allocating inputs or executing kernels.Validation on this branch
Based on master
d07c097b; this PR contains one cherry-picked commit with an identical patch to the source change. It does not depend on the ASE or native-spin PRs.Validation uses an NVIDIA RTX PRO 6000 Blackwell Server Edition with PyTorch 2.13.0+cu130 and Triton 3.7.1:
_launchimplementation, with repository files unchanged, fails the cold FP32 force-backward regression: shared memory required is 118784 bytes, exceeding the 101376-byte device limit.source/tests/pt_expt/kernels/test_grid_pair_train.pypass in a fresh process on the unmodified branch, using separate compiler caches from the parent-launcher reproduction.Other GPU architectures and the project-pinned PyTorch 2.10 environment were not exercised. The original commit's full-training and layout-sweep results are historical and were not repeated for this PR. No throughput improvement is claimed.
Summary by CodeRabbit
Performance
Reliability
Compatibility