Conversation
- Move the `SDM_CHUNK_MEMORY_FRACTION` budget of attention and TabFM cell embedding chunks into one helper, `sdm._memory.chunk_memory_limit`. - Expose the automatic attention batch size limit as `TransformerBlock.auto_batch_size_limit`, so callers can plan passes that align with its chunks. Signed-off-by: Jingang Qu <jqu@nvidia.com>
- Find finite values without an `abs()` copy and count them without an int64 copy of the mask, and compute Standardize, ClipSigma and PowerTransform statistics with fewer full-size temporaries. - Find the Yeo-Johnson bounds before allocating the PowerTransform workspaces. - Compute RobustScale quantiles over column chunks and ClipSigma transforms over row chunks within the chunk memory limit. - Transform in Standardize, ClipSoft and RobustScale with fewer temporaries. - Select evenly spaced ensemble members as views, select shared DropConstantColumns groups before re-stacking members, and select Choice members one option at a time. Signed-off-by: Jingang Qu <jqu@nvidia.com>
- Without gradients on CUDA, embed the context rows once and the query rows in balanced passes that replay the recorded context state, when the query cells would exceed the chunk memory limit. - Align passes with the chunks of the row attention, so every row runs in a chunk of the same size as in a single pass. Passes then match a single pass up to rare rounding differences in small passes. - Free the label embedding and each layer's full key/value early in the ICL block. Signed-off-by: Jingang Qu <jqu@nvidia.com>
- Transform queries and post-process member outputs in passes over rows within the chunk memory limit in `RecipeExecution`. - Keep member outputs in the model output dtype until post-processing, which casts them to the query dtype and inverts numerical targets pass by pass. - Post-process the benchmark adapter's outputs through `transform_output` after freeing the transformed queries. Signed-off-by: Jingang Qu <jqu@nvidia.com>
JingangQu
requested review from
RBendias,
akihironitta,
aw471 and
rusty1s
as code owners
September 27, 2026 09:50
JingangQu
added this pull request to stack #997
September 27, 2026 13:35
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/structured-data-models/.coderabbit.yaml Review profile: QUIET Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe pull request adds memory-aware row chunking for model and recipe execution, updates numerical processors to use finite-value statistics and chunked operations, and adds RankGaussian and balanced categorical shifts. Tests cover chunked execution, non-finite values, and the new processor methods. ChangesMemory-aware execution
Numerical processor updates
Additional processor options
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant RowEmbedding
participant TransformerLayer
participant Cache
RowEmbedding->>TransformerLayer: Process context rows
TransformerLayer->>Cache: Record context state
RowEmbedding->>Cache: Freeze context state
RowEmbedding->>TransformerLayer: Process query row chunks with cached state
Merge Risk: ⚪ Minimal · up to The reviewed test adjustment does not change CUDA execution, and no concrete user-impacting regression is established in the supplied context. No specific merge blocker remains. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
- Keep every fitted value as a knot while a column has at most max_knots rows, and every distinct value while it has at most max_knots of them, so outputs are unchanged in both cases. - Otherwise keep the fitted values whose mid-ranks come closest to normal quantiles evenly spaced between the extremes, which keeps outputs within about two knot spacings in normal scores. - Fit columns and transform rows in chunks within the chunk memory limit. - Keep loaded state in double precision, and support contexts and queries without rows. Signed-off-by: Jingang Qu <jqu@nvidia.com>
JingangQu
changed the title
`TableTensor.cuda()` is typed to return `Tensor`, so the type check rejected passing it to `MemberQuery._replace`. Cast it back to `TableTensor`. Signed-off-by: Jingang Qu <jqu@nvidia.com>
JingangQu
changed the title
This branch has not been deployed
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.
Stacked on #994.
Without the KV cache,
KumoTabularembeds all context and query rows in one forward and post-processes all estimator outputs at once. On BeyondArena, Kumo-Tabular-L (16 estimators, contexts capped at 200k rows and 500 columns) runs out of memory on an 80 GB H100 for the largest tables, e.g.delivery_eta_1min the forward andmaps_router_eta_1min the recipe fit. This PR runs the query side of the no-KV path in passes sized from the chunk memory limit thatTransformerBlockalready uses forbatch_size_limit="auto", and cuts full-size temporaries in the recipe fit. All 142 BeyondArena core datasets now run without the KV cache, at unchanged speed, with predictions equal to #994 up to rare floating-point rounding.Changes
sdm/_memory.py):chunk_memory_limit(device)is the budget that attention"auto"batching already used (device memory × per-process fraction ×SDM_CHUNK_MEMORY_FRACTION); attention and the TabFM cell embedding now call it, andsplit_sizesizes balanced chunks within it. All new chunk and pass sizes derive from this limit or from the context size; none is hard-coded.row_embedding.py): without gradients on CUDA, when the query rows' cell buffer exceeds the limit,RowEmbeddingembeds the context rows once and the query rows in balanced passes that replay the recorded context throughsdm.cache.Cache, the same record and replay asfit/predict. Passes align with the row-attention chunks, so every row runs in a chunk of the same size as in a single pass; otherwise FlashAttention would round long rows differently in small chunks. The ICL block still runs once and now frees the label embedding and each layer's key/value early.execution.py,base.py):RecipeExecution.transformandtransform_output(outputs, dtype)process rows in passes.transform_outputnow also casts member outputs todtypeand inverts numerical targets pass by pass, so outputs stay in the model dtype until then;ICLModelpasses the dtype of the first transformed query, before callbacks.Recipedocuments the row independence this relies on.ensemble.py,choice.py,numerical/): evenly spaced ensemble members, such asChoice's round-robin options, are selected as views instead of copies;nanmean,isfiniteand bool counts no longer make full-size copies;PowerTransformfinds its bounds before allocating workspaces;RobustScalequantiles andClipSigmatransforms run in chunks within the limit. Processor outputs are bitwise unchanged.Results
maps_router_eta_1m, and Elo is unchanged (1409 vs 1410 in one pool).Known limitation
The recipe fit is not chunked. On
maps_router_eta_1mit still sets the 74.8 GiB peak, becauseChoicere-stacks its members while its input and option outputs are alive. This fits on an H100 80 GB and an RTX Pro 6000 96 GB.RankGaussian from #993, memory-bounded
This PR also includes #993 (rank-Gaussian views and balanced class shifts in the Kumo recipe) and bounds the memory of its
RankGaussian, which kept every fitted value and ran out of memory onmaps_router_eta_1mwith Kumo-Tabular-L:max_knots(default 8192) knots. Columns with at mostmax_knotsrows or distinct values keep all of them, so outputs are unchanged; larger columns keep the fitted values whose mid-ranks come closest to normal quantiles evenly spaced between the extremes, which stays within about two knot spacings (max |Δz| ≤ 2e-3 on large BeyondArena tables).