TE Gemma tutorial attempt#2

[sudhakarsingh27]

Description

Adds a tutorial to showcase how to:

1. use TE TransformerEngine layer in place of HuggingFace's GemmaDecoderLayer in Gemma models.
2. use non-paged and paged KV cache from TE
3. use CUDA Graphs and fp8_model_init to optimize generation times

Attempt#1 @ #829

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)

[pggPL]

This tutorial is not visible in the docs.

[pggPL]

I was focusing more on documentation. The first part with HF finetuning is ok (I have some small suggestions, but if we merge it without any change I think it would be ok).

For the second part - I generally agree with high-level concepts, but it needs to be polished. So I think it will be good if you read it once again and look for minor issues and then I will have a look once again. I left some comments with areas in which something is missing or they left after some removed parts.

I think also that if we present something as an example we should be sure that it works correctly - so maybe consider copying the code from all the cells into some test file and add it to the CI pipeline. But this may be part of other PR, this is big enough.

I was not looking deeply into the code - it seems that some code is left commented and I am not sure if this is final form. If you want me to look into that, let me know.

I think this is a good moment to think more about the future of TE docs. We were adding many unrelated tutorials/examples lately (attention, export to onnx, offloading in near future) and our docs seem quite unstructured at this point. We have section "tutorials and examples" when there is code explaining some fp8 recipes and these tutorials. Maybe it is worth to split it into tutorial part and example part. In tutorial path we can have only descriptions of the most important features and in examples composition of these things like gemma/llama finetuning/generation.

We lack good descriptions of many good features we have inside - blockwise scaling is not in the docs, moe support, context parallelism. Maybe some refactoring and adding more structure and adding descriptions of these things is good idea for the future. @ptrendx

[pggPL]

This file is nice in current form - simple extension of llama tutorial.
Having llama and gemma tutorial next to each other seems to be quite weird, but I do not see better solution and we should leave it as it is imo.

[pggPL]

this picture looks wierd in html

[sudhakarsingh27]

I now think this picture is an overkill. I'll plan to remove it

[pggPL]

Why it is on CPU? I haven't tried to deeply understand what's going on here yet.

[sudhakarsingh27]

There's a self.pre_step_seqlens variable in InferenceParams which is on GPU and this temp variable is used to populate it. But I see the confusion - this call is redundant as torch.Tensor is created on CPU by default. I'll remove it