$ the-wire · showcase
int64 offsets in mem-efficient SDPA, FSDP policy fix for chunked loss
By RepoJournal · Filed · About PyTorch · Composed from the cited sources · methodology
15 people shipped this
peterdsharpe
@peterdsharpe
2 cited
Ziming Wang
@ZenAlexa
1 cited
drisspg
@drisspg
1 cited
sramireddy2
@sramireddy2
1 cited
Sherlock Huang
@SherlockNoMad
Felipe Mello
@felipemello1
fegin
@fegin
Andrey Talman
@atalman
Ivan Zaitsev
@izaitsevfb
Bita Bayestehparvin
@bbayeste
novak-vaclav
@novak-vaclav
Scott Roy
@metascroy
Erik-Lundell
@Erik-Lundell
abm-77
@abm-77
AdrianLundell
@AdrianLundell
The day's consequential work is two correctness fixes: mem-efficient attention no longer returns wrong results past 2^31 key/value offsets, and torchtitan's chunked loss stops clobbering the caller's FSDP policies.
[SDPA] Fix 32-bit key/value offset overflow in mem-efficient attention (#203345) pytorch/pytorch
The per-key-block offsets to_key_start * k/v_strideM (three sites) and the head offsets head_id * q_strideH and kv_head_id * k/v_strideH now compute in int64; they wrapped past 2^31 in int32. Long-context workloads on the mem-efficient kernel were silently getting wrong numbers, so pin to a build containing this commit.
[Loss] Preserve caller FSDP policies across chunked backward (#5563) pytorch/torchtitan
ChunkedLossWrapper restored head sync unconditionally on the final chunk, overwriting the caller's FSDP and HSDP accumulation policies, and an inner root callback could reduce earlier decoder gradients before the current microbatch arrived. The wrapper now follows the schedule's manual-finalization owner from the first head pass across GPipe, 1F1B and interleaved schedules.
[BugFix] Raise instead of losing an entry when a "/" key clashes in to_h5 and to_zarr pytorch/tensordict
to_h5, to_zarr and PersistentTensorDict.from_dict split paths on '/', so a key like "c/d" landed where the nested key ("c", "d") goes and one entry silently vanished, leaving only a UserWarning about replacing an array. Leaf keys are now checked after the write and a clash raises instead of dropping data.
[BugFix] Warn when memmap or save writes an entry named shape, device or _type pytorch/tensordict
Saving through memmap, memmap_, save or a .tdz archive with an entry named shape, device or _type at any depth loses that entry on load_memmap with no error, because meta.json stores bookkeeping under the same names. Datasets with a shape label and shadow=True tensorclasses with a device field are exposed; the fix warns at write time.
[BugFix] Let test_trainer.py collect without Hydra so tests-optdeps runs again pytorch/rl
tests-optdeps has failed on every push to main since Sept 25, because test_trainer.py took module-level Hydra imports and collection aborted with ModuleNotFoundError, so nothing ran. Import guards let collection proceed, and TdMpc2Planner tests are skipped when Hydra is absent.
Action items
- → Update PyTorch past 7c2c0e7 to get the int64 offset fix for mem-efficient SDPA pytorch/pytorch [plan]
