Repository navigation
fix(store): respect PyTorch tensor strides - #5392
Conversation
Codecov Report❌ Patch coverage is
❌ Your project check has failed because the head coverage (67.70%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #5392 +/- ##
==========================================
+ Coverage 67.67% 67.70% +0.02%
==========================================
Files 1258 1258
Lines 194975 195257 +282
==========================================
+ Hits 131953 132200 +247
- Misses 63022 63057 +35 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Reference tracel-ai/burn#5392 (fixes #5383) in upstream-issues §4 — the open PR that will let us drop the contiguous-preprocess workaround for SPAN/RealPLKSR/SCUNet once merged + released. Co-authored-by: DeepSeek V4 Flash <41898282+github-copilot[bot]@users.noreply.github.com>
antimora
left a comment
antimora
left a comment
There was a problem hiding this comment.
Reviewed the stride handling end to end, and verified findings against the branch rather than by reading alone: disassembled the new fixture with pickletools and independently re-derived its expected values, probed the helper functions directly, and built corrupted .pt files to exercise the failure paths.
The core fix is correct. The gather handles offset, permute and zero-stride broadcast uniformly, is_contiguous_stride's size-1-dim skipping matches PyTorch's own is_contiguous, and the new fixture is a genuine regression test that fails on main. Baseline is green: 238 lib tests plus 37 in pytorch-tests, clippy --all-targets -D warnings clean, cargo fmt --check clean.
Replacing *i as usize with usize::try_from also closes real wrapping bugs, and the fallback this removes was worse than merely lenient: it returned vec![0.0f32; n] regardless of dtype, so a truncated i64 tensor came back as f32 and failed later with a confusing dtype error.
Three things I would like to see before merge, all of them one-liners or net line reductions:
- Drop
.max(1)atpickle_reader.rs:755. Verified:empty.ptthen loads correctly, scalars are unchanged, 239/239 tests pass, and the three newshape.contains(&0)special cases become unnecessary. - Propagate the read error at line 819 instead of returning zeros. A deleted storage member currently loads as an all-zeros tensor while a storage truncated by 8 bytes is a hard error.
- Scope the storage-extent check at line 757 to one tensor. As written it aborts
PytorchReader::newfor the whole checkpoint with a message that names no tensor.
Details inline, along with some smaller doc and test notes.
One process note: the truncation change is user-visible. Anyone whose checkpoint previously "loaded" with zero padding now gets an error, so it probably deserves a release note even though the checklist marks no user-facing change.
| "Tensor byte length calculation overflows usize".to_string(), | ||
| ) | ||
| })?; | ||
| let mut reordered = Vec::with_capacity(byte_len); |
There was a problem hiding this comment.
This allocates before validating any index, whereas contiguous_tensor_bytes validates against data.len() first and never allocates on the failure path.
The new storage_extent > total check at line 757 does not bound this. storage_extent sums (dim - 1) * stride, so a dim with stride 0 contributes nothing: a pickle declaring shape (2^21, 2^21) with stride (0, 0) over a 1-element storage has extent == total == 1 and passes, then dies in the allocator. Vec::with_capacity routes to handle_alloc_error, which aborts rather than unwinds, so the catch_unwind in to_data() does not contain it. That makes it an uncatchable process kill where the contiguous path would have produced a recoverable DataError.
main aborts on the same input via its own Vec::with_capacity(num_elements) in the dtype arms, so this is pre-existing in kind. Worth closing here because this is the hardening pass, and because the extent check reads as protection against exactly this while not providing it. Comparing num_elements against data.len() / element_size up front is one line and matches what the contiguous path already does.
There was a problem hiding this comment.
I did not add num_elements <= data.len() / element_size: legitimate zero-stride views violate that relation (the existing expanded fixture has 6 logical elements backed by 3 stored elements). b0f55ed instead validates the maximum referenced storage index against actual bytes before reordering and uses fallible try_reserve_exact for the intermediate buffer. A general cap on materialized logical tensor size would need a repository-wide resource policy, so I am leaving this thread open for maintainer preference.
There was a problem hiding this comment.
Reopening this one. try_reserve_exact does not close it.
I ran the case end to end: shape (2^21, 2^21), stride (0, 0), 32-element storage. Extent is 1 so the bounds check passes, try_reserve_exact for 17592186044416 bytes succeeds, and the gather loop then runs 2^42 iterations. I killed it at 100% CPU and 34% of system RAM.
Isolated probe on macOS:
PROBE_RESERVE: SUCCEEDED for 17592186044416 bytes (cap=17592186044416)
try_reserve_exact only reports what the allocator rejects at reservation time. macOS, and Linux with overcommit, hand back address space lazily, so the failure moves to first touch, which here is the per-element extend_from_slice. The guard turns an immediate abort into slow RAM exhaustion.
Your objection to num_elements <= data.len() / element_size was right and my suggestion was wrong: expand legitimately produces more logical elements than stored ones, as your own expanded fixture shows. There is no bound derivable from metadata alone, so this does need the resource policy you described.
Not blocking, and not something to fix here. main hits the same wall on the same input via values.resize(num_elements, 0.0), so severity is unchanged. Tracked in #5506.
The one thing worth changing is the wording above, since as written it reads as resolved and the next person will trust it.
There was a problem hiding this comment.
Thanks for testing this end to end and for the correction. To clarify my earlier reply: try_reserve_exact does not prevent the zero-stride overcommit / first-touch exhaustion case, so this PR does not resolve that resource-exhaustion risk. The severity remains unchanged from main, and the repository-wide materialization policy is tracked in #5506. I’m leaving this thread open for that follow-up.
# Conflicts: # crates/burn-store/src/pytorch/pickle_reader.rs
|
@laggui Thank you. I merged current upstream main (b8693f3) into this branch in 64e992e and resolved the only conflict in crates/burn-store/src/pytorch/pickle_reader.rs by retaining the upstream burn_pack::Tensor / bridge::deferred migration while porting the cumulative stride reordering, storage validation, empty/scalar, and legacy-storage behavior. Validation on the resulting exact head:
GitHub now reports the PR mergeable against that main head, with the new current-head CI running. Could you please take another look and confirm whether it can be merged once those checks finish? |
antimora
left a comment
antimora
left a comment
There was a problem hiding this comment.
Rechecked at 64e992e. Verified against a build rather than by reading: 260 unit tests, 6 integration and 36 doctests pass, cargo clippy -p burn-store --all-targets -- -D warnings clean, cargo fmt --check clean, and git merge-tree against current main (9542363) merges clean.
Every thread is addressed except the allocation one, where I replied inline. Short version: try_reserve_exact does not fire on macOS or on Linux with overcommit, so that input still exhausts RAM. Severity is unchanged from main and it does not belong in this PR. Tracked in #5506. Worth adjusting the reply wording there though, since it currently reads as resolved.
Two things I checked independently rather than taking on trust.
The merge resolution. No TensorSnapshot leftovers, the upstream zeros() helper from a467575 is fully removed along with its call sites and leaves no dangling references, and every removed line traces to the stride work.
Issue #5383's actual reproducer, which you skipped a fixture for. I synthesized shape [1, 3, 3, 3] with strides (27, 1, 9, 3) against the fixture storage and hand-derived the expected gather:
[0, 3, 6, 9, 12, 15, 18, 21, 24, 1, 4, 7, 10, 13, 16, 19, 22, 25, 2, 5, 8, 11, 14, 17, 20, 23, 26]
Exact match, so the gather really is rank-generic and skipping the extra binary fixture is fine.
Dropping the zero-fill fallback is the most valuable part of this change. A missing storage member loading as plausible all-zero weights is the kind of failure that costs someone a week.
Approving.
Pull Request Template
Checklist
cargo run-checkshas been executed.Related Issues/PRs
Fixes #5383.
Changes
Release note
Corrupt PyTorch checkpoints whose storage is missing or truncated now return a contextual data error. They are no longer accepted by padding or replacing tensor values with zeros.
Testing
cargo run-checks(initial implementation)cargo xtask check lint(current update)cargo test -p burn-store(250 unit tests and 35 doctests passed; 17 doctests ignored)cargo clippy -p burn-store --all-targets -- -D warningscargo fmt --all -- --check