pyTorch — Transformer Engine 1.2.1 documentation
论文:https://arxiv.org/pdf/2310.06825.pdf
transformer engine 的 sliding window 是用了 flash atttention(新版本2以上,这里用的最新版本2.5.2)里对sliding_window的实现。所以不需要用transformer engine。直接用flash attention即可。
传入参数为
(4096,0):前面看4096个位置。
(-1,-1):全局注意力
(-1,0):自回归注意力,只看前面所有位置。是无sliding window 的实现。