added query-key norm to accomodate OLMo2

[ysjprojects]

query_states = self.q_norm(self.q_proj(hidden_states))
key_states = self.k_norm(self.k_proj(hidden_states))
value_states = self.v_proj(hidden_states)

https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmo2/modeling_olmo2.py

OLMo2 applies RMSNorm to the q and k matrices in its attention layer, something that is not yet supported by litgpt's architecture.

To support the addition of OLMo2, this PR adds an option to norm the q and k matrices via the config.norm_qk option which defaults to False.

Currently, the method for qk norm is assumed to follow the overall norm class.

[Andrei-Aksionov]

Looks good now.
Thank you.