TE Gemma tutorial attempt#2

[sudhakarsingh27]

Description

TLDR;
Adds a tutorial to showcase how to use Transformer Engine to accelerate generation with HuggingFace Gemma model.

Specifically, it showcases following features:

1. How to use TE's TransformerLayer layer in place of HF's GemmaDecoderLayer in Gemma models. Monkey-patching and mapping correct module names.
2. How to use non-paged and paged KV cache from TE
3. How to speedup up generation with CUDA Graphs and fp8_model_init that keeps model params in FP8 precision.

Attempt#1 @ #829

Type of change

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

[cyanguwa]

Do we need to do the same for group_linear?

[cyanguwa]

I missed this the first time, but why are we copying pre_step_seqlens_temp to CPU and then to GPU?

I think I wanted to say this while we were working on the tutorial, but we could extend our pre_step API to support cases where users provide a GPU tensor for step_dict too. But let's do that in a future PR.

[sudhakarsingh27]

self.pre_step_seqlens[: len(pre_step_seqlens_temp)].copy_(
pre_step_seqlens_temp, non_blocking=False
)

This is for CUDA graphs to work correctly since pre_step_seqlens could have different lengths. Since I'm referencing pre_step_seqlens_temp twice in this statement, I created a temporary CPU tensor. I could try checking if I could do that directly with a list but that'd mean creating a temporary list

[cyanguwa]

Could you read through this file again? I feel it's not as polished as te_gemm.py. There are some mistakes that can be easily spotted.

Also, could you please run the CI again after all changes are done? Just wanted to make sure the fp8_calibration, MHA changes are not breaking anything.

Nitpicks:

• Animation 1 -> Figure 1, are we calling it animation in other tutorials?
• Building on this foundation, the current objective is -> Building on that foundation, this tutorial is
• ushered an era -> ushered in an era?
• tutorial 2, etc. -> tutorial 2.
• for both those cases -> for both use cases
• CUDA Graph is not just used to avoid CPU launch overhead. It helps remove any other CPU overhead too - it only replays the GPU activities.
• "Weight calibration involves calculating FP8 scaling factors from higher precision forward passes." -> Does it not require backward passes?
• This eliminates the need to cast from higher precision to BF16, -> This eliminates the need to cast from higher precision to BF16 every time,
• other artefacts used in the following tutorial. -> other artifacts used in this tutorial.
• CUDA Graphs to function -> CUDA Graphs to function.
• Blue blobs in the above figure are GPU kernels and whitespace b/w those -> Blue blobs in the top figure are GPU kernels and whitespace between those
• The second big blob seems to be in different proportions between no-CG and CG, or is it just me.
• It is highly recommended to familiarize oneself with the tutorial on FP8 precision to understand the necessity of scaling. -> Didn't we say this already? I feel we're repeating ourselves quite a bit in the fp8_autocast and fp8_model_init sections.
• during the forwards in higher precision -> Is it just the forward that we need to run?
• This approach is beneficial during training: one can perform one cast for both backward and forward passes, leading to speedups. However, performing a single cast for each forward pass introduces too much overhead to achieve a speedup. -> This paragraph doesn't quite make sense.
• assert type(linear_fp8.weight.data) is te.float8_tensor.Float8Tensor -> Is .data a Float8Tensor, or .weight? Also, should we use Float8TensorBase?
• Using less memory during generation (by storing weights in FP8 precision using fp8_model_init) -> Didn't we talk about this in the tutorial? Why "not extensively talked about"?
• playing around -> playing around with?

[sudhakarsingh27]

• Could you read through this file again? I feel it's not as polished as te_gemm.py. There are some mistakes that can be easily spotted.

• Also, could you please run the CI again after all changes are done? Just wanted to make sure the fp8_calibration, MHA changes are not breaking anything.

Nitpicks:

• Animation 1 -> Figure 1, are we calling it animation in other tutorials?
  • I’m not sure if other tutorials have animations. Llama didn’t.
• Building on this foundation, the current objective is -> Building on that foundation, this tutorial is
• ushered an era -> ushered in an era?
• tutorial 2, etc. -> tutorial 2.
  • Seems to be grammatically correct but improved grammar of the context around it
• for both those cases -> for both use cases
• CUDA Graph is not just used to avoid CPU launch overhead. It helps remove any other CPU overhead too - it only replays the GPU activities.
  • Updated the text
• "Weight calibration involves calculating FP8 scaling factors from higher precision forward passes." -> Does it not require backward passes?
  • Backward pass is not needed to calibrate weights and inputs.
• This eliminates the need to cast from higher precision to BF16, -> This eliminates the need to cast from higher precision to BF16 every time,
• other artefacts used in the following tutorial. -> other artifacts used in this tutorial.
• CUDA Graphs to function -> CUDA Graphs to function.
• Blue blobs in the above figure are GPU kernels and whitespace b/w those -> Blue blobs in the top figure are GPU kernels and whitespace between those
• The second big blob seems to be in different proportions between no-CG and CG, or is it just me.
  • You are correct, I just noticed that for first generated tokens, the CUDA graphs’ attn kernel is slightly longer than non CUDA graphs case. After that it becomes the same
• It is highly recommended to familiarize oneself with the tutorial on FP8 precision to understand the necessity of scaling. -> Didn't we say this already? I feel we're repeating ourselves quite a bit in the fp8_autocast and fp8_model_init sections.
• during the forwards in higher precision -> Is it just the forward that we need to run?
  • yes, only the fwd pass is needed
• This approach is beneficial during training: one can perform one cast for both backward and forward passes, leading to speedups. However, performing a single cast for each forward pass introduces too much overhead to achieve a speedup. -> This paragraph doesn't quite make sense.
  • Updated paragraph
• assert type(linear_fp8.weight.data) is te.float8_tensor.Float8Tensor -> Is .data a Float8Tensor, or .weight? Also, should we use Float8TensorBase?
  • This is just a small dirty check. We use FakeTensor for FP8 weights since PyT doesn’t have a notion of FP8 datatype, so weight is still in BF16 but weight.data is actual Float8Tensor.
  • Updated the check now to make it more intuitive
• Using less memory during generation (by storing weights in FP8 precision using fp8_model_init) -> Didn't we talk about this in the tutorial? Why "not extensively talked about"?
  • We aren't talking about the memory usage of the models in the tutorial. I've now mentioned in there how fp8_model_init results in less memory usage.
• playing around -> playing around with?

[sudhakarsingh27]

/te-ci pytorch L0

[sudhakarsingh27]

/te-ci pytorch L0

[sudhakarsingh27]

/te-ci pytorch L0