[Pytorch] Implement fp32 accumulation for attention with context parallel in both forward and backward pass.#821

Open

Yuxin-CV wants to merge 2 commits intoNVIDIA:mainfrom Yuxin-CV:main

+35-17

Commits on Apr 28, 2024

Modify the attention forward pass to accumulate attention outputs using fp32 precision when employing context parallelism
Yuxin-CV
committed
Modify the attention backward pass to accumulate dq & dkv using fp32 precision buffer when employing context parallelism
Yuxin-CV
committed