For a more comprehensive tutorial, check out our `Quickstart Notebook <https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/quickstart.ipynb>`_.
For a more comprehensive tutorial, check out our `Getting Started Guide <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/getting_started.html>`_.
.. overview-end-marker-do-not-remove
.. overview-end-marker-do-not-remove
...
@@ -175,15 +175,22 @@ For example to use the NGC PyTorch container interactively,
...
@@ -175,15 +175,22 @@ For example to use the NGC PyTorch container interactively,
.. code-block:: bash
.. code-block:: bash
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.08-py3
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.01-py3
For example to use the NGC JAX container interactively,
For example to use the NGC JAX container interactively,
.. code-block:: bash
.. code-block:: bash
docker run --gpus all -it --rm nvcr.io/nvidia/jax:25.08-py3
docker run --gpus all -it --rm nvcr.io/nvidia/jax:26.01-py3
Where 25.08 (corresponding to August 2025 release) is the container version.
Where 26.01 (corresponding to January 2026 release) is the container version.
We recommend updating to the latest NGC container available here:
If you run any examples, please ensure you are using a matching version of TransformerEngine. TransformerEngine is pre-built and packaged inside the containers with examples available at ``/opt/transformerengine`` or ``/opt/transformer-engine``. If you would like to use examples from TE main branch and are running into import errors, please try the latest pip package or building from source, although NGC containers are recommended for ease-of-use for most users.
***Solution:**ThiscanoccurwhenTEisbuiltagainstthecontainer's system installation of cuDNN, but pip packages inside the virtual environment pull in pip packages for ``nvidia-cudnn-cu12/cu13``. To resolve this, when building TE from source please specify the following environment variables to point to the cuDNN in your virtual environment.
* **Symptoms:** Regular TE installs work correctly but UV wheel builds fail at runtime.
* **Solution:** Ensure that ``uv build --wheel --no-build-isolation -v`` is used during the wheel build as well as the pip installation of the wheel. Use ``-v`` for verbose output to verify that TE is not pulling in a mismatching version of PyTorch or JAX that differs from the UV environment'sversion.
***Solution:**Ensure``--no-build-isolation``isusedduringinstallation.Ifpre-buildingwheels,ensurethatthewheelisbothbuiltandinstalledwith``--no-build-isolation``.See"Problems using UV or Virtual Environments"aboveifusingUV.
"# Accelerating Hugging Face Llama 2 and 3 Fine-Tuning with Transformer Engine\n",
"\n",
"<div class=\"alert alert-info\">\n",
"\n",
"<b>Goal</b>\n",
"\n",
"This tutorial showcases how to accelerate finetuning a full [Llama 2](https://huggingface.co/meta-llama/Llama-2-7b-hf) or [Llama 3](https://huggingface.co/meta-llama/Meta-Llama-3-8B) models from Hugging Face by using `TransformerLayer` from the [Transformer Engine library](https://github.com/NVIDIA/TransformerEngine) in `BF16` and `FP8` precisions.\n",
"\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"id": "331f476a",
"metadata": {},
"source": [
"## Dependencies for this tutorial\n",
"\n",
"Following files and media are necessary to effectively run this tutorial:\n",
"\n",
"1. `te_llama.py`\n",
" - This file contains the code to load a Hugging Face Llama 2 or Llama 3 checkpoint in Transformer Engine's `TransformerLayer` instead of Hugging Face's `LlamaDecoderLayer`. This is used in the following two sections of the tutorial - \"Improvement 1\" and \"Improvement 2\".\n",
"2. `utils.py`\n",
" - This file contains the code related to dataloading, hyperparameters, setting up model/optimizers/accelerator, model training and other miscellaneous tasks like restarting the jupyter notebook from within the cell. \n",
"3. `media/`\n",
" - This directory contains the images used in the following tutorial.\n",
"\n",
"These packages are necessary to run this tutorial:\n",
"<b>Note on running the tutorial with Llama 3 weights</b>\n",
"\n",
"This tutorial shows the cell outputs when run with Llama 2 7B weights. It can be run with Llama 3 8B weights simply by providing the directory with those weights (in Hugging Face format) instead of Llama 2 7B weights. These two models are almost identical, the biggest difference being the model dimension (the smallest Llama 3 model has 8B parameters, whereas the smallest Llama 2 has 7B), which enables this tutorial to work for both of them.\n",
" <figcaption> Fig 1: Llama visualized as a transformer. (generated with [Nvidia's AI-foundation models](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/sdxl))</figcaption>\n",
"</figure>\n",
"\n",
"A flashback:\n",
"\n",
"- 2017: [\"Attention Is All You Need\"](https://arxiv.org/abs/1706.03762) paper introduced pioneering \"Transformer\" architecture and changed the NLP field forever.\n",
"- 2018-2020: Emergence of GPT model series that showed causal decoder architectures are great fit for pretraining, few-shot and zero-shot learning.\n",
"- Fast forward to 2023-2024: Following GPT-3/GPT-4 success stories, researchers and companies raced to produce the next best pretrained model that could further be finetuned for application-specific use-cases.\n",
"- February 2023: Meta releases [Llama 2](https://llama.meta.com/llama2) models (Large Language Model Meta AI). \n",
" - These models range from 7B to 70B parameters.\n",
" - LLaMA 2 was pretrained on 2 trillion tokens.\n",
"- April 2024: Meta releases [Llama 3](https://llama.meta.com/llama3) models.\n",
" - These models range from 8B to 70B parameters.\n",
" - LLaMA 3 was pretrained on 15 trillion tokens.\n",
"\n",
"For more information on Llama 2 consider reading the [Huggingface tutorial](https://huggingface.co/blog/llama2). As a quick summary, here are some of the important differences b/w the conventional transformer decoder architecture vs Llama 2 architecture:\n",
"\n",
"1. Decoder only model (causal language modeling and next word prediction)\n",
"2. RMSNorm in place of the LayerNorm\n",
"3. SwiGLU activation function\n",
"4. RoPE as positional embeddings \n",
"5. Grouped Query Attention for the 70B model\n",
"6. Trained on 4K context length\n",
"\n",
"Hugging Face also released a [tutorial about Llama 3](https://huggingface.co/blog/llama3). The key points are:\n",
"\n",
"1. Use of bigger tokenizer - 128256 vs 32K.\n",
"2. Grouped Query Attention is used also by smaller 8B model.\n",
"3. The context length increased to 8K for all models.\n",
"3. Llama 3 was trained on 8x more data than Llama 2.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/transformer_vs_llama.svg\">\n",
" <figcaption> Fig 2: Comparing GPT and Llama architectures. </figcaption>\n",
"</figure>"
]
},
{
"cell_type": "markdown",
"id": "a110de1a",
"metadata": {},
"source": [
"## Hugging Face's `LlamaModel`\n",
"Hugging Face provides an open-source implementation of `Llama` model in [modeling_llama.py](https://github.com/huggingface/transformers/blob/3d2900e829ab16757632f9dde891f1947cfc4be0/src/transformers/models/llama/modeling_llama.py#L4).\n",
"\n",
"Here's a block diagram that shows how Llama model is implemented in the Hugging Face repo. Notice the modular encapsulated form and `LlamaDecoderLayer` at the core of the model implementation.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/llama_for_causal_lm.svg\">\n",
" <figcaption> Fig 3: Causal Llama Model Block Diagram. </figcaption>\n",
"</figure>\n",
"\n",
"The above diagram translates to the following text output of the model in PyTorch. Notice that the core of the model has 32 `LlamaDecoderLayer`s. \n",
"Let's take a closer look at `LlamaDecoderLayer`. It is composed of `input_layernorm`, `self_attn`, `post_attention_layernorm` and `mlp` modules. Each module has associated weights as shown in the diagram.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/llama_zoom.svg\">\n",
" <figcaption> Fig 4: Causal Llama Model Block Diagram (with simplified illustration of the [LlamaDecoderLayer](https://github.com/huggingface/transformers/blob/e770f0316d2a9b787c9d1440f204fcb65e176682/src/transformers/models/llama/modeling_llama.py#L695)). </figcaption>\n",
"</figure>\n",
"\n",
"##### Self_Attn Layer\n",
"For simplicity in the block diagram illustration of the \"self_attn\" box, we omit the \"Grouped Query Attention\" operation and only showcase the modules which have associated weights.\n",
" \n",
"##### MLP Layer\n",
"\n",
"SwiGLU is an activation defined as follows in the [modeling_llama.py](https://github.com/huggingface/transformers/blob/7c4995f93d8d24aae05e1e43279c96dce736e5c8/src/transformers/models/llama/modeling_llama.py#L236) file in the Hugging Face github repo:\n",
"```\n",
"\"\"\"\n",
"1. `self.up_proj`, `self.gate_proj` and `self.down_proj` are \"Linear\" layers\n",
"It requires a set of 3 weights as compared to 2 weights in conventional \"MLP\" layers e.g. in the traditional transformer or GPT architectures. This is also illustrated in the following figure:\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/swiglu.svg\">\n",
" <figcaption> Fig 5: A look inside the feedforward layer with <code>swiglu</code> activation function. </figcaption>\n",
"Llama 2 weights are loaded into the Hugging Face native implementation `LlamaForCausalLM` (refer to [modeling_llama.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py)). \n",
"\n",
"For this and other subsequent runs, the `batch_size` is `8`. The `LlamaDecoderLayer` is left unchanged in the baseline as follows:\n",
"The baseline implementation will be run in `BF16` precision.\n",
"\n",
"</div>"
]
},
{
"cell_type": "markdown",
"id": "b38eb3ac",
"metadata": {},
"source": [
"<div class=\"alert alert-info\">\n",
"\n",
"<b>Note</b>\n",
" \n",
"This tutorial loads and trains a Llama 3 8B or a Llama 2 7B model which takes up most of the GPU memory and therefore, we need to restart the jupyter notebook each time before running the following sections. A small utility method `restart_jupyter_notebook` is defined in the accompanying `utils.py` file. This function restarts the jupyter notebook so that the GPU memory is flushed before the model is loaded again from the checkpoint in order to avoid running into OOM (Out Of Memory) errors.\n",
"\n",
"If the utility doesn't work, comment this line `restart_jupyter_notebook()` in the following cell and manually restart the jupyter notebook before running the cell. Repeat the same for other sections in this tutorial.\n",
"\n",
"</div>\n"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "2e9d7a8c",
"metadata": {},
"outputs": [
{
{
"name": "stdout",
"cell_type": "markdown",
"output_type": "stream",
"metadata": {},
"text": [
"source": [
"10 finetuning steps complete!\n",
"# Accelerating Hugging Face Llama 2 and 3 Fine-Tuning with Transformer Engine\n",
"Average time taken per step: 248 milliseconds\n"
"\n",
]
"<div class=\"alert alert-info\">\n",
}
"\n",
],
"<b>Goal</b>\n",
"source": [
"\n",
"# Restart the notebook (to flush the GPU memory)\n",
"This tutorial showcases how to accelerate finetuning a full [Llama 2](https://huggingface.co/meta-llama/Llama-2-7b-hf) or [Llama 3](https://huggingface.co/meta-llama/Meta-Llama-3-8B) models from Hugging Face by using `TransformerLayer` from the [Transformer Engine library](https://github.com/NVIDIA/TransformerEngine) in `BF16` and `FP8` precisions.\n",
"from utils import restart_jupyter_notebook\n",
"\n",
"restart_jupyter_notebook()\n",
"</div>\n"
"\n",
],
"\n",
"id": "6a5b2993"
"# Import necessary packages, methods and variables\n",
},
"from utils import *\n",
"\n",
"\n",
"# Provide Huggingface Access Token\n",
"hyperparams.hf_access_token = \"\"\n",
"assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n",
"\n",
"# Provide a directory to cache weights in to avoid downloading them every time.\n",
"# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n",
"hyperparams.weights_cache_dir = \"\"\n",
"\n",
"# For Llama 2, uncomment this line (also set by default)\n",
"In addition to basic layers like `Linear` and `LayerNorm`, Transformer Engine offers larger modules like `MultiheadAttention` (combines \"LayerNorm\" and \"Self Attention\") and `LayerNormMLP` (combines \"LayerNorm\" and \"MLP\") that could replace their counterparts in the `LlamaDecoderLayer` and potentially provide a speedup. Transformer Engine also offers a full `TransformerLayer` (which further combines `MultiheadAttention` and `LayerNormMLP` layers) which could replace `LlamaDecoderLayer` and provide a speedup (with careful mapping of the weights since the name of the weights are different for those two layers). Let's take a closer look at Transformer Engine's `TransformerLayer`. \n",
"\n",
"#### Transformer Engine's `TransformerLayer`\n",
"\n",
"At a higher level, TE's `TransformerLayer` could be visualized as an apt replacement for the `LlamaDecoderLayer`. But the internals of the `TransformerLayer` are organized a bit differently. \n",
"Just like Hugging Face's `LlamaDecoderLayer`, Transformer Engine's `TransformerLayer` encapsulates `self_attention` (as `MultiheadAttention`) and `mlp` (as `LayerNormMLP`). A major difference is that the two `Norm`s are included in the `MultiheadAttention` and `LayerNormMLP` layers as shown in the following output prompt:\n",
"\n",
"```\n",
"TransformerLayer(\n",
" (self_attention): MultiheadAttention(\n",
" (layernorm_qkv): LayerNormLinear()\n",
" (core_attention): DotProductAttention()\n",
" (proj): Linear()\n",
" )\n",
" (layernorm_mlp): LayerNormMLP()\n",
")\n",
"```\n",
"\n",
"Another difference is that Transformer Engine implements an efficient version of feedforward layer with SwiGLU in which the weights from the `up_proj` and `gate_proj` modules are merged together and SwiGLU is applied using a custom fused kernel. This is done so that only one big and efficient Matrix Multiplication operation is issued to the GPU instead of two smaller ones.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/swiglu_te.svg\">\n",
" <figcaption> Fig 8: Abstract illustration of the SwiGLU implementation in Transformer Engine. </figcaption>\n",
"</figure>\n",
"\n",
"#### `TransformerLayer` options explained\n",
"\n",
"<div class=\"alert alert-info\">\n",
"\n",
"<b>Note</b>\n",
" \n",
"Here, we go over some of the options in `TransformerLayer` that are needed for the tutorial. For a complete list of options, refer the [TransformerLayer API documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html?highlight=transformerlayer#transformer_engine.pytorch.TransformerLayer).\n",
"\n",
"</div>\n",
"\n",
"In the accompanying `te_llama.py` file, `TELlamaDecoderLayer` is defined as a wrapper over TE's `TransformerLayer` with a few needed options that make `TransformerLayer` a plug-in replacement for the HF's `LlamaDecoderLayer`.\n",
"Here's a list summarizing each option briefly:\n",
"\n",
"1. `hidden_size`: size of each input sample.\n",
"2. `ffn_hidden_size`: intermediate size to which samples are projected.\n",
"3. `num_attention_heads`: number of attention heads in the transformer layer.\n",
"4. `bias`: switch to add additive biases to the submodule layers.\n",
"5. `layernorm_epsilon`: a value added to the denominator of layer normalization for numerical stability. Default is `1e-5`.\n",
"6. `hidden_dropout`: dropout probability for the dropout op after FC2 layer (fully connected layer no. 2). Default is `0.1`.\n",
"7. `attention_dropout`: dropout probability for the dropout op during multi-head attention. Default is `0.1`. \n",
"8. `fuse_qkv_params`: if set to True, TransformerLayer module exposes a single fused parameter for query-key-value. This enables optimizations such as QKV fusion without concatentations/splits and also enables the argument fuse_wgrad_accumulation.\n",
"9. `normalization`: type of normalization applied. Default is `LayerNorm`.\n",
"10. `activation`: type of activation used in the MLP block. Default is `gelu`.\n",
"11. `attn_input_format`: controls whether the dimensions of the intermediate hidden states is 'batch first' ('bshd') or 'sequence first' ('sbhd'). `s` stands for the sequence length, `b` batch size, `h` the number of heads, `d` head size. Note that these formats are very closely related to the `qkv_format` in the `MultiHeadAttention` and `DotProductAttention` modules.\n",
"12. `num_gqa_groups`: number of GQA groups in the transformer layer. Grouped Query Attention is described in [this paper](https://arxiv.org/pdf/2305.13245.pdf). This only affects the keys and values, not the querys. GQA-1 is equivalent to Multi-Query Attention ([MQA](https://arxiv.org/pdf/1911.02150.pdf)), while GQA-H is equivalent to MultiHead Attention, i.e. `num_gqa_groups = num_attention_heads`.\n",
"\n",
"\n",
"Further, note that `RotaryPositionEmbedding` is defined as part of the `TELlamaDecoderLayer` (wrapper around TE's `TransformerLayer`) itself since it expects this rope cache if RoPE is used in the model. \n",
"\n",
"Let's revisit how `LlamaDecoderLayer`s form the core of the decoder layer stack in HF's llama implementation:\n",
"A major portion of the Hugging Face model implementation (32 `LlamaDecoderLayer` layers) could be potentially replaced with Transformer Engine's `TransformerLayer` layers. Let's see how it is made possible.\n",
"\n",
"\n",
"#### Mapping weights from HF's `LlamaDecoderLayer` to TE's `TransformerLayer`\n",
"\n",
"Refer the accompanying file `te_llama.py` which provides a reference to create a Llama 2 model with TE's `TransformerLayer` after replacing HF's `LlamaDecoderLayer`.\n",
"\n",
"Briefly, following pieces of code are put together:\n",
"\n",
"1. `TELlamaDecoderLayer` is added as a wrapper for `TransformerLayer`. \n",
"2. Before creating a `LlamaForCausalLM`, `replace_decoder` context manager is used to monkey-patch `LlamaDecoderLayer` with `TELlamaDecoderLayer`.\n",
"\n",
"```\n",
"@contextmanager\n",
"def replace_decoder(te_decoder_cls):\n",
" \"\"\"\n",
" Replace `LlamaDecoderLayer` with custom `TELlamaDecoderLayer`.\n",
"3. A custom `pretrained_from_local` method is added that copies the weights from the checkpoint (which is meant for HF Llama implementation) to the modified `TELlamaForCausalLM` by carefully mapping the weights from the `LlamaDecoderLayer` (HF) to `TransformerLayer` (TE). The method `replace_params` maps and copies apt weights from `LlamaDecoderLayer` to the `TransformerLayer`. Refer to the following diagram for more details.\n",
"The following figure shows how the weights get mapped from the HF's `LlamaDecoderLayer` to TE's `TransformerLayer`.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/weight_swap.svg\">\n",
" <figcaption> Fig 9: Replace `LlamaDecoderLayer` with `TransformerLayer`. </figcaption>\n",
"</figure>\n",
"\n",
"After initializing the modified Llama model this way, the core decoder layers get changed to `TELlamaDecoderLayer` (wrapper around `TransformerLayer`) as shown in the following output:\n",
"In summary, the model gets changed as follows with a large chunk of the implementation (core decoder layers) coming from Transformer Engine.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/model_change.svg\">\n",
" <figcaption> Fig 10: Language model after the HF's `LlamaDecoderLayer`s are replaced with TE's `TransformerLayer`s. </figcaption>\n",
"</figure>\n",
"\n",
"\n",
"<div class=\"alert alert-info\">\n",
"<b>Note</b>\n",
"\n",
"Let's first run this \"TELlama\" implementation in `BF16` precision.\n",
"</div>"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "bdb34b91",
"metadata": {},
"outputs": [
{
{
"name": "stdout",
"cell_type": "markdown",
"output_type": "stream",
"metadata": {},
"text": [
"source": [
"10 finetuning steps complete!\n",
"## Dependencies for this tutorial\n",
"Average time taken per step: 185 milliseconds\n"
"\n",
]
"Following files and media are necessary to effectively run this tutorial:\n",
}
"\n",
],
"1. `te_llama.py`\n",
"source": [
" - This file contains the code to load a Hugging Face Llama 2 or Llama 3 checkpoint in Transformer Engine's `TransformerLayer` instead of Hugging Face's `LlamaDecoderLayer`. This is used in the following two sections of the tutorial - \"Improvement 1\" and \"Improvement 2\".\n",
"# Restart the notebook (to flush the GPU memory)\n",
"2. `utils.py`\n",
"from utils import restart_jupyter_notebook\n",
" - This file contains the code related to dataloading, hyperparameters, setting up model/optimizers/accelerator, model training and other miscellaneous tasks like restarting the jupyter notebook from within the cell. \n",
"restart_jupyter_notebook()\n",
"3. `requirements.txt`\n",
"\n",
" - This file contains the necessary Python packages for this tutorial.\n",
"\n",
"4. `media/`\n",
"# Import necessary packages, methods and variables\n",
" - This directory contains the images used in the following tutorial.\n",
"from utils import *\n",
"\n",
"\n",
"\n",
"\n",
"<div class=\"alert alert-info\">\n",
"# Provide Huggingface Access Token\n",
"\n",
"hyperparams.hf_access_token = \"\"\n",
"<b>Note on running the tutorial with Llama 3 weights</b>\n",
"assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n",
"\n",
"\n",
"This tutorial shows the cell outputs when run with Llama 2 7B weights. It can be run with Llama 3 8B weights simply by providing the directory with those weights (in Hugging Face format) instead of Llama 2 7B weights. These two models are almost identical, the biggest difference being the model dimension (the smallest Llama 3 model has 8B parameters, whereas the smallest Llama 2 has 7B), which enables this tutorial to work for both of them.\n",
"# Provide a directory to cache weights in to avoid downloading them every time.\n",
"\n",
"# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n",
"</div>\n",
"hyperparams.weights_cache_dir = \"\"\n",
""
"\n",
],
"# For Llama 2, uncomment this line (also set by default)\n",
"# Uncomment and run this cell when running the tutorial for the first time\n",
},
"# %pip install -r requirements.txt"
{
],
"cell_type": "markdown",
"id": "099697e2",
"id": "0c9fbd65",
"execution_count": null,
"metadata": {},
"outputs": []
"source": [
},
"Compared to the \"baseline\" implementation, we see that using Transformer Engine's `TransformerLayer` in place of Huggging Face's `LlamaDecoderLayer` gives a speedup of **34%** even when using only BF16 precision!\n",
"\n",
"| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n",
"Now that most of the HF Llama model implementation (`LlamaDecoderLayer`s) has been swapped with Transformer Engine implementation (`TELlamaDecoderLayer` or `TransformerLayer`), let's see how finetuning in `FP8` precision helps improve performance.\n",
"\n",
"#### How to run the model in `FP8` precision\n",
"\n",
"After the substitution, the model can be run in `FP8` precision by the following change over the previous BF16 runs. (For more information, refer the corresponding `wrap_with_accelerator` function in the accompanying `utils.py` file).\n",
"\n",
"```\n",
"# Specify the `FP8RecipeKwargs` (additional argument required to run in `fp8` precision)\n",
" <figcaption> Fig 1: Llama visualized as a transformer. (generated with [Nvidia's AI-foundation models](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/sdxl))</figcaption>\n",
"</figure>\n",
"\n",
"A flashback:\n",
"\n",
"- 2017: [\"Attention Is All You Need\"](https://arxiv.org/abs/1706.03762) paper introduced pioneering \"Transformer\" architecture and changed the NLP field forever.\n",
"- 2018-2020: Emergence of GPT model series that showed causal decoder architectures are great fit for pretraining, few-shot and zero-shot learning.\n",
"- Fast forward to 2023-2024: Following GPT-3/GPT-4 success stories, researchers and companies raced to produce the next best pretrained model that could further be finetuned for application-specific use-cases.\n",
"- February 2023: Meta releases [Llama 2](https://llama.meta.com/llama2) models (Large Language Model Meta AI). \n",
" - These models range from 7B to 70B parameters.\n",
" - LLaMA 2 was pretrained on 2 trillion tokens.\n",
"- April 2024: Meta releases [Llama 3](https://llama.meta.com/llama3) models.\n",
" - These models range from 8B to 70B parameters.\n",
" - LLaMA 3 was pretrained on 15 trillion tokens.\n",
"\n",
"For more information on Llama 2 consider reading the [Huggingface tutorial](https://huggingface.co/blog/llama2). As a quick summary, here are some of the important differences b/w the conventional transformer decoder architecture vs Llama 2 architecture:\n",
"\n",
"1. Decoder only model (causal language modeling and next word prediction)\n",
"2. RMSNorm in place of the LayerNorm\n",
"3. SwiGLU activation function\n",
"4. RoPE as positional embeddings \n",
"5. Grouped Query Attention for the 70B model\n",
"6. Trained on 4K context length\n",
"\n",
"Hugging Face also released a [tutorial about Llama 3](https://huggingface.co/blog/llama3). The key points are:\n",
"\n",
"1. Use of bigger tokenizer - 128256 vs 32K.\n",
"2. Grouped Query Attention is used also by smaller 8B model.\n",
"3. The context length increased to 8K for all models.\n",
"3. Llama 3 was trained on 8x more data than Llama 2.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/transformer_vs_llama.svg\">\n",
" <figcaption> Fig 2: Comparing GPT and Llama architectures. </figcaption>\n",
"</figure>"
],
"id": "e37e2cc1"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Hugging Face's `LlamaModel`\n",
"Hugging Face provides an open-source implementation of `Llama` model in [modeling_llama.py](https://github.com/huggingface/transformers/blob/3d2900e829ab16757632f9dde891f1947cfc4be0/src/transformers/models/llama/modeling_llama.py#L4).\n",
"\n",
"Here's a block diagram that shows how Llama model is implemented in the Hugging Face repo. Notice the modular encapsulated form and `LlamaDecoderLayer` at the core of the model implementation.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/llama_for_causal_lm.svg\">\n",
" <figcaption> Fig 3: Causal Llama Model Block Diagram. </figcaption>\n",
"</figure>\n",
"\n",
"The above diagram translates to the following text output of the model in PyTorch. Notice that the core of the model has 32 `LlamaDecoderLayer`s. \n",
"Let's take a closer look at `LlamaDecoderLayer`. It is composed of `input_layernorm`, `self_attn`, `post_attention_layernorm` and `mlp` modules. Each module has associated weights as shown in the diagram.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/llama_zoom.svg\">\n",
" <figcaption> Fig 4: Causal Llama Model Block Diagram (with simplified illustration of the [LlamaDecoderLayer](https://github.com/huggingface/transformers/blob/e770f0316d2a9b787c9d1440f204fcb65e176682/src/transformers/models/llama/modeling_llama.py#L695)). </figcaption>\n",
"</figure>\n",
"\n",
"#### Self_Attn Layer\n",
"For simplicity in the block diagram illustration of the \"self_attn\" box, we omit the \"Grouped Query Attention\" operation and only showcase the modules which have associated weights.\n",
" \n",
"#### MLP Layer\n",
"\n",
"SwiGLU is an activation defined as follows in the [modeling_llama.py](https://github.com/huggingface/transformers/blob/7c4995f93d8d24aae05e1e43279c96dce736e5c8/src/transformers/models/llama/modeling_llama.py#L236) file in the Hugging Face github repo:\n",
"```\n",
"\"\"\"\n",
"1. `self.up_proj`, `self.gate_proj` and `self.down_proj` are \"Linear\" layers\n",
"It requires a set of 3 weights as compared to 2 weights in conventional \"MLP\" layers e.g. in the traditional transformer or GPT architectures. This is also illustrated in the following figure:\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/swiglu.svg\">\n",
" <figcaption> Fig 5: A look inside the feedforward layer with <code>swiglu</code> activation function. </figcaption>\n",
"Llama 2 weights are loaded into the Hugging Face native implementation `LlamaForCausalLM` (refer to [modeling_llama.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py)). \n",
"\n",
"For this and other subsequent runs, the `batch_size` is `8`. The `LlamaDecoderLayer` is left unchanged in the baseline as follows:\n",
"The baseline implementation will be run in `BF16` precision.\n",
"\n",
"</div>"
],
"id": "c9529229"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<div class=\"alert alert-info\">\n",
"\n",
"<b>Note</b>\n",
" \n",
"This tutorial loads and trains a Llama 3 8B or a Llama 2 7B model which takes up most of the GPU memory and therefore, we need to restart the jupyter notebook each time before running the following sections. A small utility method `restart_jupyter_notebook` is defined in the accompanying `utils.py` file. This function restarts the jupyter notebook so that the GPU memory is flushed before the model is loaded again from the checkpoint in order to avoid running into OOM (Out Of Memory) errors.\n",
"\n",
"If the utility doesn't work, comment this line `restart_jupyter_notebook()` in the following cell and manually restart the jupyter notebook before running the cell. Repeat the same for other sections in this tutorial.\n",
"\n",
"</div>\n"
],
"id": "b38eb3ac"
},
{
"cell_type": "code",
"metadata": {},
"source": [
"# Restart the notebook (to flush the GPU memory)\n",
"from utils import restart_jupyter_notebook\n",
"restart_jupyter_notebook()\n",
"\n",
"\n",
"# Import necessary packages, methods and variables\n",
"from utils import *\n",
"\n",
"\n",
"# Provide Huggingface Access Token\n",
"hyperparams.hf_access_token = \"\"\n",
"assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n",
"\n",
"# Provide a directory to cache weights in to avoid downloading them every time.\n",
"# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n",
"hyperparams.weights_cache_dir = \"\"\n",
"\n",
"# For Llama 2, uncomment this line (also set by default)\n",
"In addition to basic layers like `Linear` and `LayerNorm`, Transformer Engine offers larger modules like `MultiheadAttention` (combines \"LayerNorm\" and \"Self Attention\") and `LayerNormMLP` (combines \"LayerNorm\" and \"MLP\") that could replace their counterparts in the `LlamaDecoderLayer` and potentially provide a speedup. Transformer Engine also offers a full `TransformerLayer` (which further combines `MultiheadAttention` and `LayerNormMLP` layers) which could replace `LlamaDecoderLayer` and provide a speedup (with careful mapping of the weights since the name of the weights are different for those two layers). Let's take a closer look at Transformer Engine's `TransformerLayer`. \n",
"\n",
"### Transformer Engine's `TransformerLayer`\n",
"\n",
"At a higher level, TE's `TransformerLayer` could be visualized as an apt replacement for the `LlamaDecoderLayer`. But the internals of the `TransformerLayer` are organized a bit differently. \n",
"Just like Hugging Face's `LlamaDecoderLayer`, Transformer Engine's `TransformerLayer` encapsulates `self_attention` (as `MultiheadAttention`) and `mlp` (as `LayerNormMLP`). A major difference is that the two `Norm`s are included in the `MultiheadAttention` and `LayerNormMLP` layers as shown in the following output prompt:\n",
"\n",
"```\n",
"TransformerLayer(\n",
" (self_attention): MultiheadAttention(\n",
" (layernorm_qkv): LayerNormLinear()\n",
" (core_attention): DotProductAttention()\n",
" (proj): Linear()\n",
" )\n",
" (layernorm_mlp): LayerNormMLP()\n",
")\n",
"```\n",
"\n",
"Another difference is that Transformer Engine implements an efficient version of feedforward layer with SwiGLU in which the weights from the `up_proj` and `gate_proj` modules are merged together and SwiGLU is applied using a custom fused kernel. This is done so that only one big and efficient Matrix Multiplication operation is issued to the GPU instead of two smaller ones.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/swiglu_te.svg\">\n",
" <figcaption> Fig 8: Abstract illustration of the SwiGLU implementation in Transformer Engine. </figcaption>\n",
"</figure>\n",
"\n",
"### `TransformerLayer` options explained\n",
"\n",
"<div class=\"alert alert-info\">\n",
"\n",
"<b>Note</b>\n",
" \n",
"Here, we go over some of the options in `TransformerLayer` that are needed for the tutorial. For a complete list of options, refer the [TransformerLayer API documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html?highlight=transformerlayer#transformer_engine.pytorch.TransformerLayer).\n",
"\n",
"</div>\n",
"\n",
"In the accompanying `te_llama.py` file, `TELlamaDecoderLayer` is defined as a wrapper over TE's `TransformerLayer` with a few needed options that make `TransformerLayer` a plug-in replacement for the HF's `LlamaDecoderLayer`.\n",
"Here's a list summarizing each option briefly:\n",
"\n",
"1. `hidden_size`: size of each input sample.\n",
"2. `ffn_hidden_size`: intermediate size to which samples are projected.\n",
"3. `num_attention_heads`: number of attention heads in the transformer layer.\n",
"4. `bias`: switch to add additive biases to the submodule layers.\n",
"5. `layernorm_epsilon`: a value added to the denominator of layer normalization for numerical stability. Default is `1e-5`.\n",
"6. `hidden_dropout`: dropout probability for the dropout op after FC2 layer (fully connected layer no. 2). Default is `0.1`.\n",
"7. `attention_dropout`: dropout probability for the dropout op during multi-head attention. Default is `0.1`. \n",
"8. `fuse_qkv_params`: if set to True, TransformerLayer module exposes a single fused parameter for query-key-value. This enables optimizations such as QKV fusion without concatentations/splits and also enables the argument fuse_wgrad_accumulation.\n",
"9. `normalization`: type of normalization applied. Default is `LayerNorm`.\n",
"10. `activation`: type of activation used in the MLP block. Default is `gelu`.\n",
"11. `attn_input_format`: controls whether the dimensions of the intermediate hidden states is 'batch first' ('bshd') or 'sequence first' ('sbhd'). `s` stands for the sequence length, `b` batch size, `h` the number of heads, `d` head size. Note that these formats are very closely related to the `qkv_format` in the `MultiHeadAttention` and `DotProductAttention` modules.\n",
"12. `num_gqa_groups`: number of GQA groups in the transformer layer. Grouped Query Attention is described in [this paper](https://arxiv.org/pdf/2305.13245.pdf). This only affects the keys and values, not the querys. GQA-1 is equivalent to Multi-Query Attention ([MQA](https://arxiv.org/pdf/1911.02150.pdf)), while GQA-H is equivalent to MultiHead Attention, i.e. `num_gqa_groups = num_attention_heads`.\n",
"\n",
"\n",
"Further, note that `RotaryPositionEmbedding` is defined as part of the `TELlamaDecoderLayer` (wrapper around TE's `TransformerLayer`) itself since it expects this rope cache if RoPE is used in the model. \n",
"\n",
"Let's revisit how `LlamaDecoderLayer`s form the core of the decoder layer stack in HF's llama implementation:\n",
"A major portion of the Hugging Face model implementation (32 `LlamaDecoderLayer` layers) could be potentially replaced with Transformer Engine's `TransformerLayer` layers. Let's see how it is made possible.\n",
"\n",
"\n",
"### Mapping weights from HF's `LlamaDecoderLayer` to TE's `TransformerLayer`\n",
"\n",
"Refer the accompanying file `te_llama.py` which provides a reference to create a Llama 2 model with TE's `TransformerLayer` after replacing HF's `LlamaDecoderLayer`.\n",
"\n",
"Briefly, following pieces of code are put together:\n",
"\n",
"1. `TELlamaDecoderLayer` is added as a wrapper for `TransformerLayer`. \n",
"2. Before creating a `LlamaForCausalLM`, `replace_decoder` context manager is used to monkey-patch `LlamaDecoderLayer` with `TELlamaDecoderLayer`.\n",
"\n",
"```\n",
"@contextmanager\n",
"def replace_decoder(te_decoder_cls):\n",
" \"\"\"\n",
" Replace `LlamaDecoderLayer` with custom `TELlamaDecoderLayer`.\n",
"3. A custom `pretrained_from_local` method is added that copies the weights from the checkpoint (which is meant for HF Llama implementation) to the modified `TELlamaForCausalLM` by carefully mapping the weights from the `LlamaDecoderLayer` (HF) to `TransformerLayer` (TE). The method `replace_params` maps and copies apt weights from `LlamaDecoderLayer` to the `TransformerLayer`. Refer to the following diagram for more details.\n",
"The following figure shows how the weights get mapped from the HF's `LlamaDecoderLayer` to TE's `TransformerLayer`.\n",
"\n",
"<figure align=\"center\">\n",
"<img src=\"media/weight_swap.svg\">\n",
" <figcaption> Fig 9: Replace `LlamaDecoderLayer` with `TransformerLayer`. </figcaption>\n",
"</figure>\n",
"\n",
"After initializing the modified Llama model this way, the core decoder layers get changed to `TELlamaDecoderLayer` (wrapper around `TransformerLayer`) as shown in the following output:\n",
"Compared to the \"baseline\" implementation, we see that using Transformer Engine's `TransformerLayer` in place of Huggging Face's `LlamaDecoderLayer` gives a speedup of **34%** even when using only BF16 precision!\n",
"\n",
"| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n",
"Now that most of the HF Llama model implementation (`LlamaDecoderLayer`s) has been swapped with Transformer Engine implementation (`TELlamaDecoderLayer` or `TransformerLayer`), let's see how finetuning in `FP8` precision helps improve performance.\n",
"\n",
"### How to run the model in `FP8` precision\n",
"\n",
"After the substitution, the model can be run in `FP8` precision by the following change over the previous BF16 runs. (For more information, refer the corresponding `wrap_with_accelerator` function in the accompanying `utils.py` file).\n",
"\n",
"```\n",
"# Specify the `FP8RecipeKwargs` (additional argument required to run in `fp8` precision)\n",
"| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 217 | 1.24 |\n",
"| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | FP8 | 185 | 1.46 |\n",
"\n",
"For Llama 3 8B, we get the most speedup of **46%** with FP8 precision!\n",
"\n"
],
"id": "e7cf9c3a"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Conclusion\n",
"\n",
"Using `TransformerLayer` module from Transformer Engine as a substitute for Hugging Face's `LlamaDecoderLayer` provides a speedup over Hugging Face's native Llama 2 and Llama 3 implementations. This needs careful initialization of the model such that the model weights (which are meant for `LlamaDecoderLayer`) are correctly mapped to their counterparts in TE's `TransformerLayer`. Even with `BF16` precision, `TransformerLayer` provides a speedup over the baseline implementation. With `FP8` precision, the speed up is even more pronounced!"
],
"id": "95d6c42b"
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
}
],
"source": [
"# Restart the notebook (to flush the GPU memory)\n",
"from utils import restart_jupyter_notebook\n",
"restart_jupyter_notebook()\n",
"\n",
"\n",
"# Import necessary packages, methods and variables\n",
"from utils import *\n",
"\n",
"\n",
"# Provide Huggingface Access Token\n",
"hyperparams.hf_access_token = \"\"\n",
"assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n",
"\n",
"# Provide a directory to cache weights in to avoid downloading them every time.\n",
"# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n",
"hyperparams.weights_cache_dir = \"\"\n",
"\n",
"# For Llama 2, uncomment this line (also set by default)\n",
"| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 217 | 1.24 |\n",
"| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | FP8 | 185 | 1.46 |\n",
"\n",
"For Llama 3 8B, we get the most speedup of **46%** with FP8 precision!\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "95d6c42b",
"metadata": {},
"source": [
"## Conclusion\n",
"\n",
"Using `TransformerLayer` module from Transformer Engine as a substitute for Hugging Face's `LlamaDecoderLayer` provides a speedup over Hugging Face's native Llama 2 and Llama 3 implementations. This needs careful initialization of the model such that the model weights (which are meant for `LlamaDecoderLayer`) are correctly mapped to their counterparts in TE's `TransformerLayer`. Even with `BF16` precision, `TransformerLayer` provides a speedup over the baseline implementation. With `FP8` precision, the speed up is even more pronounced!"
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
FP8 Current Scaling
===================================
FP8 current scaling recipe is the simplest low precision recipe provided by Transformer Engine.
To understand how this recipe works, we first need to examine what the FP8 data type is and how it differs from other floating point formats.
FP8 data type
-------------
The FP8 datatype, introduced in Hopper architecture, is actually 2 distinct datatypes, useful in different parts of the training of neural networks:
* E4M3 -- consists of 1 sign bit, 4 exponent bits and 3 bits of mantissa. It can store values up to +/-448 and ``nan``.
* E5M2 -- consists of 1 sign bit, 5 exponent bits and 2 bits of mantissa. It can store values up to +/-57344, +/- ``inf`` and ``nan``. The tradeoff of the increased dynamic range is lower precision of the stored values.
.. raw:: html
:file: img/fp8_formats.svg
*Figure 1: Structure of the floating point datatypes. All of the values shown (in FP16, BF16, FP8 E4M3 and FP8 E5M2) are the closest representations of value 0.3952.*
**E4M3 and E5M2 usage in training**
By default, Transformer Engine uses a hybrid approach:
* *Forward pass* - activations and weights require more precision, so E4M3 datatype is used to store them.
* *Backward pass* - gradients are less susceptible to precision loss but require higher dynamic range, so E5M2 datatype is preferred.
The user can configure this behavior via the ``fp8_format`` parameter of the recipe.
Scaling factors
---------------
Limited dynamic range of FP8 datatype is insufficient for many tensors.
To address this, values in the tensor are scaled. FP8 Current Scaling recipe uses one **FP32** scale factor per tensor. The representation of a tensor element ``x`` in FP8 precision is given by:
.. code-block:: python
x = x_fp8 * s
where
* ``x_fp8`` is the FP8 value (E4M3 or E5M2),
* ``s`` is a global **FP32** scaling factor applied to the entire tensor.
**FP8 Current Scaling quantization**
Let's take a closer look at how quantization to FP8 with scaling factor is implemented in
the FP8 Current Scaling recipe.
.. raw:: html
:file: img/fp8_scaling_concept.svg
*Figure 3: Quantization to FP8 consists of amax (absolute maximum) computation, scaling to fit the FP8 range and casting to the respective FP8 format.*
Quantization to FP8 consists of 3 steps:
1. Computation of the absolute maximum value of the tensor - we refer to it as ``amax``.
2. Applying the scaling factor of ``fp8_max / amax`` to the tensor, to fit it into the FP8 range
3. Casting into the respective FP8 format using *Round To Nearest Even (RTNE)*. Values round to the nearest representable FP8 value. When exactly halfway between two values, rounds to the one with even mantissa to minimize systematic bias.
**Performance analysis**
Quantization is a memory-bound operation that requires reading the tensor twice:
* First read: compute ``amax`` across all elements.
* Second read: apply the scaling factor and cast to FP8.
This is a significant overhead compared to other recipes, which typically require only a single memory read.
.. raw:: html
:file: img/fp8_cast_process.svg
*Figure 4: FP8 quantization with current scaling recipe - two tensor reads are needed, one to compute amax and one to apply the scaling factor and cast to FP8.*
Transpose handling
------------------
*Ada and Hopper*
On Ada and Hopper, the backward pass requires a transposed FP8 tensor.
The columnwise layout is physically different from the rowwise layout, so a transpose operation is needed.
All 3 options from :ref:`Performance Considerations Transpose handling section <handling_transposes>` are supported.
*Blackwell and later*
Blackwell hardware supports multiple GEMM layouts natively, eliminating the need for explicit transposes.
The rowwise and columnwise tensors share the same physical memory layout.
