Fix of bug of kv_channels in TransformerLayer and added Gemma tutorial

[pggPL]

I'm working on tutorial similar to tutorial with Llama, but with Gemma model. Most parts are similar, but I have encountered few differences:

1. Official weights are in Safetensor format, not as in the case of Llama - torch checkpoints. I modified the loading function.

2. There is geglu activation function instead of swiglu, but this is also supported by the TE. It was one simple change in config.

3. Gemma hidden dimension is different than attention dimension - attention projections change dimension from hidden(3072) to key/query/value(4096). There is parameter kv_channels of the TransformerLayer in TE, which seems to be responsible for exactly that. Nevertheless, changing it to value different than hidden_dim / num_heads causes Assertion Error and is not changing number of parameters.

I propose a change of argument kv_channels to the attention_hidden_dims. I changed the description and number of neurons in attention projection layers in transforemer.py and attention.py files.

The results of Gemma with bf16, bf16 and TE, fp8 and TE are proportional to the results of the Llama.

[pggPL]

I was thinking about adding two parameters attention_hidden_dims_v and attention_hidden_dims_kq, because it is mathematically possible that hidden dimension of key vector and query vector is not equal to dimension of value vectors.

Nevertheless, FlashAttention assumes that this numbers are equal, so I gave up this idea.

[ptrendx]

Hmm, can we get the same effect without changing the API? That would be a breaking change and we would like to avoid that if possible.

[sudhakarsingh27]

Suggested change

	3 * attention_hidden_size,
	self.hidden_size_per_attention_head * num_attention_heads + 2 * self.hidden_size_kv

@pggPL, we'd still need the separate hidden size (total) for k and v heads to handle group query attention

[sudhakarsingh27]

Here's what we could do alternately:

# We still keep this line
self.hidden_size_per_attention_head = kv_channels

# Use `kv_channels` to calculate the `attention_hidden_size`
self.attention_hidden_size = self.hidden_size_per_attention_head * num_attention_heads

# Use `self.attention_hidden_size` to calculate `hidden_size_kv` (which is used in case of GQA)
self.hidden_size_kv = int(self.attention_hidden_size* self.num_gqa_groups // num_attention_heads)

...

LayerNormLinear(
                    hidden_size,
                    self.attention_hidden_size+ 2 * self.hidden_size_kv,
                    ...
)

We could probably also rename attention_hidden_size to hidden_size_q because essentially that's what it's used for in the layer/weight creation (in the call to LayerNormLinear)

[pggPL]

I see the point, but I'm not sure whether using kv_channels is best option. Currently kv_channels argument does exactly nothing if it is set to correct value. If it is set to incorrect value, the error is raised. Changing the behavior of this argument is an option, but it will be unrelated to the name - we change also q_channels.

I propose to add optional argument attn_hidden_size - legacy code will work correctly without it. We can leave kv_channels in the code - but for example only print some "this argument is deprecated" warning. Or even simply do nothing. I believe (maybe wrongly) that this will not lead to less problems with legacy code, but the argument's names will be more accurate. Let me know what do you think.

[pggPL]

I split this PR into 2:
#833 (kv_channels part)
#829 (gemma part)
Thus I close it.