$ the-wire · showcase
TensorDict store scalars, TorchTitan's one-line GC fix
By RepoJournal · Filed · About PyTorch · Composed from the cited sources · methodology
11 people shipped this
TensorDict closed three long-standing gaps in typing, export, and store writes, while TorchTitan traced a killed GB300 run to a single leaked reference in its rollout loop.
[rl] controller: drop each finished rollout group (full GCs froze the controller) (#5215) pytorch/torchtitan
Each rollout-loop task held its last finished group for about an hour; with 1,152 such tasks, full Python GCs froze the controller for 7 to 30 seconds, and a roughly 30 second freeze killed a GB300 run when Monarch dropped links at its 30 second delivery timeout. A single `del` in `Controller._rollout_loop` brings the same recipe to 1 to 4 second GCs at step 10.
[BugFix] Make modules exported with torch.export (non-strict) accept TensorDict inputs pytorch/tensordict
`torch.export.export(M(), (td,))` succeeds with its default `strict=False`, but calling the returned module with a TensorDict or tensorclass raises a tree-spec flattening `ValueError` on torch 2.13, 2.14, and nightly. Strict export was never affected, so the workaround is to pass `strict=True` until this lands.
[Quality] Type the tensorclass constructor from its fields pytorch/tensordict
`tensordict/tensorclass.pyi` typed the constructor as `(*args, batch_size=..., device=..., **kwargs)`, so mypy and pyright accepted missing required fields, unknown keywords, and non-tensor values. The constructor is now typed from the fields via `@dataclass_transform(eq_default=False)`, which means existing tensorclass call sites may start reporting type errors that were always real.
[Feature] Write scalars into stores with store[index] = value pytorch/tensordict
`store[done] = 0` raised `TypeError: 'float' object is not iterable` even though the same assignment works on a TensorDict. `TensorDictStore`, `LazyStackedTensorDictStore`, and elements of a lazy stack now accept scalars at an index, writing them into every tensor entry.
[Performance] Read the tensordict.nn modes without a lock on every module call pytorch/tensordict
Every `tensordict.nn` module call read global modes through a `threading.Lock` and a `torch.compiler.is_compiling()` call, with the skip-existing wrapper paying for its reads twice; `_ContextManager.get_mode` and `set_mode` no longer take the lock, and the wrapper's mode restore is fixed. TensorDictStore scalar writes, FA4 bump, and the rest of the day's smaller changes round out the list.
Action items
- → Pass strict=True to torch.export.export for modules taking TensorDict or tensorclass inputs until the non-strict path... pytorch/tensordict [plan]
- → Require FA4 >= 4.0.0b34 for RL head_dim 256 split-KV decode pytorch/torchtitan [plan]
