Fusing the Attention Training Dataflow with Local Safe Softmax and Model-Independent Tiling
The demand for efficient training of Transformer language models is rapidly increasing, yet the resource-intensive attention mechanism severely bottlenecks the process. Previous attempts to fuse attention operations managed to reduce off-chip memory access, thereby improving memory bandwidth. Unfortunately, these techniques are impractical in restricted environments, as they exhibit resource utilization that scales with model dimensions and do not consider the entire training dataflow. To address this, we propose operation fusion for the attention training dataflow using local safe softmax and model-independent tiling. By taking both the forward and the backward passes into account, local safe softmax reduces I/O access under small-cache conditions. Furthermore, model-independent tiling ensures that the required on-chip memory footprint remains independent of model dimensions, enabling scalability across diverse models. Software simulations and hardware evaluation show that our method reduces off-chip access, on-chip memory usage, FLOPs, training runtime, and energy cost compared to conventional approaches, confirming its suitability for efficiently training attention in resource-restricted hardware.