FlashMask-3 speeds structured attention across GPUs
In plain words
Long-context training can use masks to control which tokens attend to which others. But applying those masks and dividing the work across GPUs can slow training, especially when some GPUs receive more work than others.
FlashMask-3 is a method that makes the mask guide both GPU computation and how work and data are shared across devices. In tests, it increased full-training speed over a standard balancing and communication setup for one large model and three document-mask datasets. The result is limited to that tested setup.
The research
When training uses long sequences split across GPUs, can structured attention masks be applied without making mask processing and uneven workloads a new bottleneck? The question matters because distributing the sequence adds both computation and communication to the work that the masks shape.
FlashMask-3 combines a Hopper attention kernel with rank-local masks, sparse-work balancing, and overlapped communication for distributed training.
FlashMask-3 combines a Hopper GPU attention kernel with distributed methods for localizing masks, balancing sparse work, and overlapping computation with communication.
The comparison measured full training against Dual Chunk balancing and all-gather communication for a 118B-A14B model at 128K context on 128 H100 GPUs.
Training throughput was higher on all three document-mask datasets, by 7.3% to 11.0% across datasets with 70% to 86% average mask sparsity. This supports a gain for the combined method in that setup; attention benchmarks alone would not establish a full-training improvement.
The full-training gain was measured for one 118B-A14B model, document masks, three datasets, and a specific 128-H100 setup. The paper also reports that disabling hierarchical routing reduced the combined method’s gain to below the baseline on two of those datasets.
Analysis
Deep
- Original paper authors
- Guoxia Wang, Qianyue He, Siming Wu, Haoyang Xie, Qiwen Bao
- Submitted to arXiv
- License
- arXiv non-exclusive