Regarding the accuracy issues of deepspeed evoformer optimization

[WhenMelancholy]

Hello, I have tested the optimization effects of the evoformer attention in deepspeed-evo-attention branch, but encountered an accuracy issue. After enabling the optimization, the output of openfold differs significantly from when it is not enabled (max abs error > 0.5), and with optimization enabled, the output fluctuates with each forward pass. In contrast, without optimization, the fluctuation in openfold's forward passes is <0.0001. (Modified from this test script)

Is there a way to eliminate this fluctuation? I would be very grateful for any advice you can provide!