fix(strategy): size reconstruct() row-block from compute_dtype, not out_dtype#117
Open
boskodev790 wants to merge 1 commit into
Open
Conversation
…ut_dtype fix / bug lane — CPU-safe validation only, no GPU scorecard. subspace.reconstruct() streams its (n x n) output row-by-row and sized the row-block from out_dtype's itemsize. But the loop tensors (Q, Ctil, QT and the (rb, n) product) all live in compute_dtype, which for fp16 inputs is bumped to fp32 for accuracy. Sizing from out_dtype therefore under-budgets the row-block by 2x at --dtype fp16, risking OOM on large-N / tight-vram_fraction runs. Fix: add an optional compute_dtype param (defaults to None -> out_dtype for back-compat) and size _row_block from it; pass cfg.compute_dtype at the call site. Mirrors how compress()/stream_gemm already size from compute_dtype. Adds tests/test_reconstruct_row_block.py (CPU-only, no GPU): asserts the row-block is sized from compute_dtype (fp32) when out_dtype is fp16, that compute_dtype=None falls back to out_dtype, and that the math is unchanged. Validation: uv run python -m strategy.smoke; uv run --extra test python -m pytest tests/ strategy/tests/ eval/tests/ -q (81 passed, 23 skipped). Fixes zeokin#66
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.
Lane
fix/bug— corrects repository behavior. CPU-safe validation only; no GPU scorecard.Problem
strategy.subspace.reconstruct()streams its(n × n)output row-by-row to stay inside a VRAM budget, sizing the row-block from the output dtype's itemsize:But every device tensor inside the loop —
Q,Ctil,QT, and the(rb, n)product frombackend.matmul— lives incfg.compute_dtype, not the output dtype. Fordtype="fp16"those differ:compute_dtypeis bumped to fp32 for accuracy (seeConfig.compute_dtype). So the row-block is sized for half the real per-row footprint, and on an actual GPU at--dtype fp16with largeN/ tightvram_fractionthe streamed path can OOM. This is the same class of under-budget as theauto_tileoperand-upcast fixes, on the reconstruct path.Fix
Size the block from the dtype the loop actually computes in — mirroring how
compress()/stream_gemm_*already passcompute_dtype:and pass
cfg.compute_dtypeat the call site.compute_dtype=Nonedefaults toout_dtype, so the existing 5-arg call intest_subspace.pyis unchanged. Onlystrategy/subspace.pyandtests/change — no scorer/eval/docs/CI paths.Tests
Added
tests/test_reconstruct_row_block.py(CPU-only, no GPU — pure sizing arithmetic via a fake backend + a_row_blockspy):compute_dtype=Nonefalls back toout_dtype(back-compat).Q @ Ctil @ Qᵀ) is unchanged.Fixes #66