Thanks so much for the PR! This looks great and I'm testing it today. Though @satpalsr can you merge main into your branch and resolve the conflicts against the recent changes that add Flash Attention triton support?

I find this commit is relevant to my issue (#883).
The general idea to tackle the problem is the same.
I think we do not need to add a new parameter layer_past in the function of flash_attention, but simply use if self.training == True to check whether it is training or inference stage.
High five for the same solution!!!

@dashstander removed conflicts. Please test out now.
Thanks for the suggestion @DaoD

There are a few outstanding bugs that I have pointed out, though the general gist is ensuring that this works for both training and inference and that the code switches between those modes correctly. Some of that is actually not an issue with this PR and seems to be an issue with the generate.py script, as was mention in the (now closed) #883. Actively working on this now, should be able to iron this stuff out.

Just had a productive conversation with @satpalsr . Currently the implementation relies on the torch.nn.Module attribute training which isn't used elsewhere in the code base. There's also an issue with the three different FlashAttention functions getting assigned to the same name at initialization. Since models are always initialized with self.training == True even resetting it by calling inference_mode won't change the FlashAttention function.

Since models are always initialized with self.training == True even resetting it by calling inference_mode won't change the FlashAttention function.

Where does this happen? I was expecting it to be in setup_for_inference_or_eval but it doesn't seem to be.

Where does this happen? I was expecting it to be in setup_for_inference_or_eval but it doesn't seem to be.

@StellaAthena It's an attribute built in to torch.nn.Module and just set to True by default. Part of the changes if we want to use this particular attribute will be setting it in the train_mode and inference_mode methods (here and here). An alternative might be to specifically add an inference vs training configuration to NeoxArgs, but that would definitely need to be part of a larger discussion.

I also just realized that this doesn't add the Triton kvpacked function , so technically inference with Triton FlashAttention + ALiBi would be broken.

Edit: Nevermind, this is wrong. The Triton function we import has the Q / K / V matrices split up.

I tested and confirmed this all worked with these changes.

Thanks @dashstander for completing this.