March 2026: Delivered targeted optimizations to selective activation checkpointing (SAC) for linear operations in huggingface/torchtitan, aligning behavior with aten.mm to reduce unconditional recomputation and improve training memory efficiency. Extended the SAC policy to include aten.linear.default, normalized weight shapes to match mm conventions, and ensured consistent checkpointing across backends, including cases where aten.linear decomposes to aten.mm. These changes are traceable to commit dfb1a6ad9b7025f8b776392e35c84c3047ad04e3 and deliver more predictable memory usage, improved throughput, and a clearer maintenance path across configurations.

February 2026 monthly summary for pytorch/torchtitan focused on delivering hardware-performance visibility improvements and stabilizing training on AWS Trainium/Inferentia. Implemented metrics enhancements for BF16 TFLOPS and hardened initialization to prevent NaN losses, improving reliability and MFU accuracy on Neuron-backed instances. These changes enhance hardware utilization visibility, contribute to more stable training, and support smoother deployments on AWS training/inference hardware.