[Reward Model Training] Inconsistent accuracy caused by flash-attention about alpaca_farm

[lxuechen]

Thanks for raising this issue. The difference in accuracy you get there is more likely due to the randomness in training than any other reason.

Due to finite arithmetic, results with and without FlashAttention will differ slightly. This is because the underlying algorithm and implementation differ for FlashAttention and standard attention. The backward pass of FlashAttention is implemented in the original flash-attn codebase and has been appropriately tested; PyTorch's autograd engine handles the rest. So it's unlikely there's a bug with the backward pass.

I did another ablation with and without FlashAttention. From the figure below, you can see that the difference is small and not statistically significant. Each curve is based on three independent runs (note I fixed the validation set to be the same for all runs). The full experimental log can be found here.