Add fp32 to the set of valid inputs to attention layer

[j316chuck]

Description

Add fp32 to the set of valid inputs for attention layer.

Note:

• flash attn is incompatible with torch.float32 🔴
• torch attn is compatible with torch.float32 ✅

Tests:

• Before: full-eval-fp32-train-fp8-llama3-8b-metamath-4ep-ObqlFj 🔴
  • mixed_precision: full
• After: torch-attn-full-eval-fp32-train-fp8-metamath-4ep-pmGJKN ✅
  • mixed_precision: full

[j316chuck]

Tested manually, there is no unit test for this.

[dakinggg]

@j316chuck I don't think this is correct. Flash attention does not support fp32 (unless that changed recently?)