hf_olmo: support flash attn 2

[wade3han]

#460, tested with a simple snippet as below:

import transformers, torch

model = transformers.AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B-Instruct", use_flash_attention_2="flash_attention_2", trust_remote_code=True).cuda()
tokenizer = transformers.AutoTokenizer.from_pretrained("allenai/OLMo-7B-Instruct", trust_remote_code=True)

print(tokenizer.decode(model.generate(torch.tensor(tokenizer.encode("Hello World! My name is")).unsqueeze(0).cuda())[0]))
# Hello World! My name is Emily and I am a second year student at the University of California,

[epwalsh]

Someone familiar with transformers internals should review this (maybe @AkshitaB). I'm not sure what transformers does with this, but I'd be very cautious if they're monkey-patching our attention mechanism since flash-attn expends a different input shape (the head and sequence dimensions are flipped compared to what we normally do).