$ the-wire · showcase
TorchTitan reshapes gradient accumulation for CUDA graphs, TensorRT 11.3 lands
By RepoJournal · Filed · About PyTorch · Composed from the cited sources · methodology
11 people shipped this
TorchTitan's graph_trainer now supports in-graph inplace gradient accumulation so the whole training step can be captured as a CUDA graph, while PyTorch's MPS triangular solve drops the sequential kernel that ran a single right-hand side on one GPU thread.
TorchTitan's graph_trainer gains an in-graph inplace gradient accumulation mode, with GradState extracted to a dataclass and the accumulation added into the graph as an add_ node so the whole thing is CudaGraph-able [1]. The same effort adds GraphRunner, which caches the parameter enumeration once at creation instead of repeating it on every run; that overhead previously multiplied by N under GradientAccumulation, and GraphRunner is positioned as the replacement for the run_traced() entry point with both kept for now [2]. Validation currently rejects parameter aliasing, pipeline parallelism, and Precompile.
The MPS backend replaced its triangular solve with a block-based kernel and dropped MPSMatrixSolveTriangular [3]. The old kernel gave each right-hand side one thread walking all n steps, so complex64 n=2048 k=1 took 151ms against 1.3ms on CPU, and with many right-hand sides traffic reached roughly 34GB. In Inductor, the CUDA float32 GELU property now passes, so its stale expected-failure marker was removed; the xfail had turned a success into an XPASS failure [4].
TorchRL's ParallelEnv(use_buffers=False) no longer drops the state of workers it does not reset, a bug that also affected metadata_from_workers=True and that previously produced incomplete observations and ragged lazy stacks on partial resets [5]. The environment now caches each worker's latest step or reset output to fill untouched workers' observations and done flags while preserving caller-supplied values. Separately, TransformedEnv(TransformedEnv(env, t)) no longer raises a ValueError about a None transform: the type check in the unwrap branch ran before the transform is not None guard, making the guard dead code, and None is now handled first [6]. TorchRL also shipped TransformerModule, a causal transformer that runs over a [B, T] window in training or one step against a key/value cache during collection, selected with set_recurrent_mode like the existing recurrent modules [7].
pytorch/TensorRT moved to TensorRT 11.3 from 11.2 and added cu134 support [8]. On TensorRT-RTX, models containing nonzero now fall back to PyTorch instead of stopping compilation: aten.nonzero.default had no capability_validator, so the partitioner placed the node in a TensorRT block and the only guard was a decorator that raises at conversion time where the converter registry cannot see it [9]. A separate fix avoids an IShuffleLayer squeeze on a genuinely empty dimension that tripped a TensorRT-internal squeezeDims assertion on platforms such as DGX Spark, a 26.09 regression, by emitting an empty constant of the target shape instead [10].
Action items
- → Upgrade pytorch/TensorRT to TensorRT 11.3 if you need cu134 support pytorch/TensorRT [plan]
- → Pull pytorch/rl for the ParallelEnv(use_buffers=False) partial-reset state fix pytorch/rl [plan]
- → Monitor TorchTitan graph_trainer if you plan to CUDA-graph gradient accumulation; parameter aliasing, PP, and Precompile remain unsupported pytorch/torchtitan [monitor]
References
- [1] [graph_trainer] in-graph inplace gradient accumulation mode, to make it CudaGraph friendly ↗ pytorch/torchtitan
- [2] [graph_trainer] GraphRunner to reduce non-Graph runtime overhead ↗ pytorch/torchtitan
- [3] [MPS] Solve triangular systems by blocks, drop MPSMatrixSolveTriangular (#196241) ↗ pytorch/pytorch
- [4] [inductor] Remove stale GELU eager-equivalence xfail (#196544) ↗ pytorch/pytorch
- [5] [BugFix] ParallelEnv without buffers keeps the state of workers it does not reset ↗ pytorch/rl
- [6] [BugFix] TransformedEnv wraps a transformed env without a transform ↗ pytorch/rl
- [7] [Feature] Transformer Module ↗ pytorch/rl
- [8] upgrade to tensorrt 11.3 ↗ pytorch/TensorRT
- [9] fix(dynamo): fall back to PyTorch for nonzero on TensorRT-RTX ↗ pytorch/TensorRT
- [10] fix select empty tensor squeeze dgx spark ↗ pytorch/TensorRT
