实现 Flash Attention 的梯度计算，通过 Recomputation 实现内存高效的训练。

配套代码

Grouped Query Attention

实现 GQA/MQA 支持,让多个 Query Head 共享 KV,优化 KV Cache 内存占用。

分布式训练

从数据并行到多维混合并行，理解大模型训练的核心并行策略