diff --git a/include/cudnn_frontend/node/dln.h b/include/cudnn_frontend/node/dln.h index ad7ad1f1c..1cfa4f308 100644 --- a/include/cudnn_frontend/node/dln.h +++ b/include/cudnn_frontend/node/dln.h @@ -29,54 +29,41 @@ class DLNNode : public NodeCRTP { attributes.fill_from_context(context); - // TODO: Only inferencing from X works today. - auto X = attributes.inputs[Layernorm_backward_attributes::input_names::X]; - auto const x_tensor_dim = X->get_dim(); + auto X = attributes.inputs[Layernorm_backward_attributes::input_names::X]; + auto const x_tensor_dim = X->get_dim(); + auto const x_tensor_stride = X->get_stride(); auto DY = attributes.inputs[Layernorm_backward_attributes::input_names::DY]; auto dy_tensor_dim = DY->get_dim(); // Only infer dims and strides if user did not set them if (dy_tensor_dim.empty()) { - dy_tensor_dim.resize(x_tensor_dim.size()); DY->set_dim(x_tensor_dim); } if (DY->get_stride().empty()) { - auto const& DY_dim = DY->get_dim(); - // Default to NHWC - auto const& stride_order = detail::generate_NHWC_stride_order(DY_dim.size()); - DY->set_stride(detail::generate_stride(DY_dim, stride_order)); + DY->set_stride(x_tensor_stride); } auto DX = attributes.outputs[Layernorm_backward_attributes::output_names::DX]; auto dx_tensor_dim = DX->get_dim(); // Only infer dims and strides if user did not set them if (dx_tensor_dim.empty()) { - dx_tensor_dim.resize(x_tensor_dim.size()); DX->set_dim(x_tensor_dim); } if (DX->get_stride().empty()) { - auto const& DX_dim = DX->get_dim(); - // Default to NHWC - auto const& stride_order = detail::generate_NHWC_stride_order(DX_dim.size()); - DX->set_stride(detail::generate_stride(DX_dim, stride_order)); + DX->set_stride(x_tensor_stride); } - auto scale_bias_dim = X->get_dim(); - scale_bias_dim[0] = 1; + auto scale = attributes.inputs[Layernorm_backward_attributes::input_names::SCALE]; - // Set channel length tensors - auto infer_scale_bias_tensors = [&scale_bias_dim](std::shared_ptr& T) { - auto tensor_dim = T->get_dim(); + // Infer dscale/dbias from scale + auto infer_scale_bias_tensors = [&scale](std::shared_ptr& T) { // Only infer dims and strides if user did not set them - if (tensor_dim.empty()) { - T->set_dim(scale_bias_dim); + if (T->get_dim().empty()) { + T->set_dim(scale->get_dim()); } if (T->get_stride().empty()) { - auto const& T_dim = T->get_dim(); - // Default to NHWC - auto const& stride_order = detail::generate_NHWC_stride_order(T_dim.size()); - T->set_stride(detail::generate_stride(T_dim, stride_order)); + T->set_stride(scale->get_stride()); } }; diff --git a/include/cudnn_frontend/node/rmsnorm.h b/include/cudnn_frontend/node/rmsnorm.h index bc1f37d22..3381e57dc 100644 --- a/include/cudnn_frontend/node/rmsnorm.h +++ b/include/cudnn_frontend/node/rmsnorm.h @@ -226,54 +226,41 @@ class DRMSNormNode : public NodeCRTP { attributes.fill_from_context(context); - // TODO: Only inferencing from X works today. - auto X = attributes.inputs[Rmsnorm_backward_attributes::input_names::X]; - auto const x_tensor_dim = X->get_dim(); + auto X = attributes.inputs[Rmsnorm_backward_attributes::input_names::X]; + auto const x_tensor_dim = X->get_dim(); + auto const x_tensor_stride = X->get_stride(); auto DY = attributes.inputs[Rmsnorm_backward_attributes::input_names::DY]; auto dy_tensor_dim = DY->get_dim(); // Only infer dims and strides if user did not set them if (dy_tensor_dim.empty()) { - dy_tensor_dim.resize(x_tensor_dim.size()); DY->set_dim(x_tensor_dim); } if (DY->get_stride().empty()) { - auto const& DY_dim = DY->get_dim(); - // Default to NHWC - auto const& stride_order = detail::generate_NHWC_stride_order(DY_dim.size()); - DY->set_stride(detail::generate_stride(DY_dim, stride_order)); + DY->set_stride(x_tensor_stride); } auto DX = attributes.outputs[Rmsnorm_backward_attributes::output_names::DX]; auto dx_tensor_dim = DX->get_dim(); // Only infer dims and strides if user did not set them if (dx_tensor_dim.empty()) { - dx_tensor_dim.resize(x_tensor_dim.size()); DX->set_dim(x_tensor_dim); } if (DX->get_stride().empty()) { - auto const& DX_dim = DX->get_dim(); - // Default to NHWC - auto const& stride_order = detail::generate_NHWC_stride_order(DX_dim.size()); - DX->set_stride(detail::generate_stride(DX_dim, stride_order)); + DX->set_stride(x_tensor_stride); } - auto scale_bias_dim = X->get_dim(); - scale_bias_dim[0] = 1; + auto scale = attributes.inputs[Rmsnorm_backward_attributes::input_names::SCALE]; - // Set channel length tensors - auto infer_scale_bias_tensors = [&scale_bias_dim](std::shared_ptr& T) { - auto tensor_dim = T->get_dim(); + // Infer dscale/dbias from scale + auto infer_scale_bias_tensors = [&scale](std::shared_ptr& T) { // Only infer dims and strides if user did not set them - if (tensor_dim.empty()) { - T->set_dim(scale_bias_dim); + if (T->get_dim().empty()) { + T->set_dim(scale->get_dim()); } if (T->get_stride().empty()) { - auto const& T_dim = T->get_dim(); - // Default to NHWC - auto const& stride_order = detail::generate_NHWC_stride_order(T_dim.size()); - T->set_stride(detail::generate_stride(T_dim, stride_order)); + T->set_stride(scale->get_stride()); } }; diff --git a/samples/cpp/norm/batchnorm.cpp b/samples/cpp/norm/batchnorm.cpp index 5b4807e29..2f44149ed 100644 --- a/samples/cpp/norm/batchnorm.cpp +++ b/samples/cpp/norm/batchnorm.cpp @@ -196,7 +196,7 @@ TEST_CASE("SGBN Add Relu Graph", "[batchnorm][graph]") { SKIP("single GPU BN is not supported in cudnn versions prior to 8.7"); #endif if (check_device_arch_newer_than("ampere") == false) { - SKIP("ConvBNFprop requires Ampere and up"); + SKIP("SGBN Add Relu requires Ampere and up"); } // Create a unique_ptr for the cuDNN handle auto handle_ptr = create_cudnn_handle(); diff --git a/samples/cpp/norm/layernorm.cpp b/samples/cpp/norm/layernorm.cpp index b8951e4bf..d1de1a5df 100644 --- a/samples/cpp/norm/layernorm.cpp +++ b/samples/cpp/norm/layernorm.cpp @@ -181,17 +181,17 @@ TEST_CASE("LayerNorm Training", "[layernorm][graph]") { auto X = graph.tensor(fe::graph::Tensor_attributes() .set_name("X") - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto scale = graph.tensor(fe::graph::Tensor_attributes() .set_name("scale") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto bias = graph.tensor(fe::graph::Tensor_attributes() .set_name("bias") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); float epsilon_cpu = 1e-05f; @@ -260,17 +260,17 @@ TEST_CASE("LayerNorm Inference", "[layernorm][graph]") { auto X = graph.tensor(fe::graph::Tensor_attributes() .set_name("X") - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto scale = graph.tensor(fe::graph::Tensor_attributes() .set_name("scale") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto bias = graph.tensor(fe::graph::Tensor_attributes() .set_name("bias") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); float epsilon_cpu = 1e-05f; @@ -332,27 +332,27 @@ TEST_CASE("LayerNorm Backward", "[layernorm][graph]") { auto X = graph.tensor(fe::graph::Tensor_attributes() .set_name("X") - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto DY = graph.tensor(fe::graph::Tensor_attributes() .set_name("DY") - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto scale = graph.tensor(fe::graph::Tensor_attributes() .set_name("scale") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto mean = graph.tensor(fe::graph::Tensor_attributes() .set_name("mean") - .set_dim({batch_size * seq_length, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) + .set_dim({batch_size, seq_length, 1}) + .set_stride({seq_length, 1, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto inv_variance = graph.tensor(fe::graph::Tensor_attributes() .set_name("inv_variance") - .set_dim({batch_size * seq_length, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) + .set_dim({batch_size, seq_length, 1}) + .set_stride({seq_length, 1, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto DLN_options = fe::graph::Layernorm_backward_attributes().set_saved_mean_and_inv_variance(mean, inv_variance); @@ -362,7 +362,7 @@ TEST_CASE("LayerNorm Backward", "[layernorm][graph]") { dbias->set_output(true).set_data_type(fe::DataType_t::FLOAT); #if (CUDNN_VERSION < 8905) - SKIP("single GPU BN is not supported in cudnn versions prior to 8.7"); + SKIP("LayerNorm is not supported in cudnn versions prior to 8.9.5"); #endif if (check_device_arch_newer_than("ampere") == false) { SKIP("LayerNorm Backward requires Ampere and up"); diff --git a/samples/cpp/norm/rmsnorm.cpp b/samples/cpp/norm/rmsnorm.cpp index 529afdbcc..5da0ec204 100644 --- a/samples/cpp/norm/rmsnorm.cpp +++ b/samples/cpp/norm/rmsnorm.cpp @@ -37,12 +37,12 @@ TEST_CASE("RmsNorm Training", "[rmsnorm][graph]") { auto X = graph.tensor(fe::graph::Tensor_attributes() .set_name("X") .set_data_type(fe::DataType_t::FLOAT) - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto scale = graph.tensor(fe::graph::Tensor_attributes() .set_name("scale") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); float epsilon_cpu = 1e-05f; @@ -77,7 +77,7 @@ TEST_CASE("RmsNorm Training", "[rmsnorm][graph]") { REQUIRE(graph.warmup(handle).is_good()); Surface X_tensor(batch_size * seq_length * hidden_size); - Surface Var_tensor(batch_size * seq_length); + Surface Inv_variance_tensor(batch_size * seq_length); Surface Scale_tensor(hidden_size); Surface Y_tensor(batch_size * seq_length * hidden_size); @@ -86,7 +86,10 @@ TEST_CASE("RmsNorm Training", "[rmsnorm][graph]") { Surface workspace(workspace_size); std::unordered_map, void*> variant_pack = { - {X, X_tensor.devPtr}, {inv_variance, Var_tensor.devPtr}, {scale, Scale_tensor.devPtr}, {Y, Y_tensor.devPtr}}; + {X, X_tensor.devPtr}, + {inv_variance, Inv_variance_tensor.devPtr}, + {scale, Scale_tensor.devPtr}, + {Y, Y_tensor.devPtr}}; REQUIRE(graph.execute(handle, variant_pack, workspace.devPtr).is_good()); } @@ -103,17 +106,17 @@ TEST_CASE("RmsNorm Inference", "[rmsnorm][graph]") { auto X = graph.tensor(fe::graph::Tensor_attributes() .set_name("X") .set_data_type(fe::DataType_t::FLOAT) - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto scale = graph.tensor(fe::graph::Tensor_attributes() .set_name("scale") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto bias = graph.tensor(fe::graph::Tensor_attributes() .set_name("bias") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); float epsilon_cpu = 1e-05f; @@ -174,23 +177,23 @@ TEST_CASE("RmsNorm Backward", "[rmsnorm][graph]") { auto X = graph.tensor(fe::graph::Tensor_attributes() .set_name("X") .set_data_type(fe::DataType_t::FLOAT) - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto DY = graph.tensor(fe::graph::Tensor_attributes() .set_name("DY") .set_data_type(fe::DataType_t::FLOAT) - .set_dim({batch_size * seq_length, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size})); + .set_dim({batch_size, seq_length, hidden_size}) + .set_stride({seq_length * hidden_size, hidden_size, 1})); auto scale = graph.tensor(fe::graph::Tensor_attributes() .set_name("scale") - .set_dim({1, hidden_size, 1, 1}) - .set_stride({hidden_size, 1, hidden_size, hidden_size}) + .set_dim({1, 1, hidden_size}) + .set_stride({hidden_size, hidden_size, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto inv_variance = graph.tensor(fe::graph::Tensor_attributes() .set_name("inv_variance") - .set_dim({batch_size * seq_length, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) + .set_dim({batch_size, seq_length, 1}) + .set_stride({seq_length, 1, 1}) .set_data_type(fe::DataType_t::FLOAT)); auto DRMS_options = fe::graph::Rmsnorm_backward_attributes().has_dbias(false); @@ -221,11 +224,9 @@ TEST_CASE("RmsNorm Backward", "[rmsnorm][graph]") { Surface X_tensor(batch_size * seq_length * hidden_size); Surface DY_tensor(batch_size * seq_length * hidden_size); - Surface Mean_tensor(batch_size * seq_length); Surface Inv_variance_tensor(batch_size * seq_length); Surface Scale_tensor(hidden_size); Surface Dscale_tensor(hidden_size); - Surface Dbias_tensor(hidden_size); Surface DX_tensor(batch_size * seq_length * hidden_size); int64_t workspace_size = 0; diff --git a/samples/python/20_layernorm_forward.ipynb b/samples/python/20_layernorm_forward.ipynb index c07b5f872..c6aab3f40 100644 --- a/samples/python/20_layernorm_forward.ipynb +++ b/samples/python/20_layernorm_forward.ipynb @@ -1,460 +1,453 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# LayerNorm Forward Operation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This notebook shows how to compute a layernorm forward operation using cuDNN.\n", - "\n", - "$$\\text{LayerNorm}(x) = \\frac{x-\\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\\cdot\\gamma+\\beta$$\n", - "\n", - "Where $\\mu = E[x]$ and $\\sigma^2 = Var[x]$ are taken over all inputs in a batch." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/20_layernorm_forward.ipynb)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites and Setup\n", - "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('nvidia-smi')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If running on Colab, you will need to install the cudnn python interface." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", - "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", - "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Overview" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we are going to perform the layer norm forward pass with the following problem sizes:\n", - "\n", - "- batch size: 4\n", - "- sequence length: 1024\n", - "- embedding dimension: 768\n", - "\n", - "The tensor will be in 16-bit float format." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import cudnn\n", - "import torch\n", - "\n", - "torch.manual_seed(1)\n", - "\n", - "handle = cudnn.create_handle()\n", - "\n", - "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", - "\n", - "assert torch.cuda.is_available()\n", - "\n", - "batch, seq_size, embedding_dim = 4, 1024, 768\n", - "dtype = torch.float16\n", - "# Epsilon is a small number to prevent division by 0.\n", - "epsilon_value = 1e-3" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Wrapper" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below is how you can use the `Graph` wrapper to perform layer norm with the input tensors in PyTorch format." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# allocate random input tensors\n", - "x_gpu = torch.randn(\n", - " batch * seq_size, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype).to(\n", - " memory_format=torch.channels_last\n", - ")\n", - "bias_gpu = torch.randn(1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype).to(\n", - " memory_format=torch.channels_last\n", - ")\n", - "eps_cpu = torch.full((1, 1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", - "\n", - "# create a graph: output in half precision, but mean and inv_var in single precision\n", - "with cudnn.Graph(\n", - " io_data_type=cudnn.data_type.HALF,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - " inputs=[\"ln::input\", \"ln::scale\", \"ln::bias\", \"ln::epsilon\"],\n", - " outputs=[\"ln::Y\", \"ln::MEAN\", \"ln::INV_VARIANCE\"],\n", - ") as graph:\n", - " out, mean, inv_var = graph.layernorm(\n", - " name=\"ln\",\n", - " input=x_gpu,\n", - " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", - " scale=scale_gpu,\n", - " bias=bias_gpu,\n", - " epsilon=eps_cpu,\n", - " )\n", - " out.set_output(True).set_data_type(dtype)\n", - " mean.set_output(True).set_data_type(torch.float32)\n", - " inv_var.set_output(True).set_data_type(torch.float32)\n", - "\n", - "# execute the graph and retrieve the output tensors\n", - "out_gpu, mean_gpu, inv_var_gpu = graph(\n", - " x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle\n", - ")\n", - "\n", - "# verify the result with PyTorch operations\n", - "out_ref = torch.nn.functional.layer_norm(\n", - " x_gpu,\n", - " [embedding_dim, 1, 1],\n", - " weight=scale_gpu.squeeze(0),\n", - " bias=bias_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "mean_ref = x_gpu.to(torch.float32).mean(dim=(1, 2, 3), keepdim=True)\n", - "inv_var_ref = torch.rsqrt(\n", - " torch.var(x_gpu.to(torch.float32), dim=(1, 2, 3), keepdim=True) + epsilon_value\n", - ")\n", - "\n", - "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(mean_gpu, mean_ref, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(inv_var_gpu, inv_var_ref, atol=5e-3, rtol=3e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Python Binding APIs" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### General Setup\n", - "Create a cudnn handle, which is a per device handle used to initialize cudnn context." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### LayerNorm Training" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them easily when we calculate reference outputs." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# allocate input tensor memory, initialize them to random numbers\n", - "x_gpu = torch.randn(\n", - " batch * seq_size,\n", - " embedding_dim,\n", - " 1,\n", - " 1,\n", - " dtype=dtype,\n", - " requires_grad=True,\n", - " device=\"cuda\",\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, dtype=dtype, requires_grad=True, device=\"cuda\"\n", - ").to(memory_format=torch.channels_last)\n", - "bias_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, dtype=dtype, requires_grad=True, device=\"cuda\"\n", - ").to(memory_format=torch.channels_last)\n", - "\n", - "# Epsilon must be a scalar value on the cpu.\n", - "epsilon_cpu = torch.full(\n", - " (1, 1, 1, 1), epsilon_value, dtype=torch.float32, requires_grad=False, device=\"cpu\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Compute reference ouputs and allocate output tensor GPU buffers" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create the reference computation outputs here before the cuDNN computation, in order to use .empty_like() to create our output buffers\n", - "out_expected = torch.nn.functional.layer_norm(\n", - " x_gpu,\n", - " [embedding_dim, 1, 1],\n", - " weight=scale_gpu.squeeze(0),\n", - " bias=bias_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "\n", - "mean_expected = x_gpu.to(torch.float32).mean(dim=(1, 2, 3), keepdim=True)\n", - "\n", - "inv_var_expected = torch.rsqrt(\n", - " torch.var(x_gpu.to(torch.float32), dim=(1, 2, 3), keepdim=True) + epsilon_value\n", - ")\n", - "\n", - "# Allocate output tensor memory using PyTorch\n", - "# PyTorch has calculated their shapes already, so we can simply use .empty_like()\n", - "out_gpu = torch.empty_like(out_expected)\n", - "mean_gpu = torch.empty_like(mean_expected)\n", - "inv_var_gpu = torch.empty_like(inv_var_expected)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Create cuDNN graph\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here we assign UIDs for tensors. UIDs are a unique identifier that will allow us to provide a mapping from tensors created from cuDNN graph api calls, such as `graph.tensor_like()`, to the underlying device memory that will be used to store these tensors. Virtual tensors don't require explicit memory allocated for them, but non-vritual tensors like inputs or outputs will need to have UIDs assigned to them. \n", - "\n", - "Alternatively, one can use handles directly in the mapping, however using UIDs can be more convinient for caching of cuDNN graphs.\n", - "\n", - "For each of our inputs {X, Scale, Bias, Epsilon} and our outputs {Out, Mean, Inverse Variance}, we allocate a UID. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from enum import Enum\n", - "\n", - "\n", - "class UID(Enum):\n", - " X = 0\n", - " SCALE = 1\n", - " BIAS = 2\n", - " EPSILON = 3\n", - " OUT = 4\n", - " MEAN = 5\n", - " INV_VAR = 6" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create the cuDNN graph.\n", - "graph = cudnn.pygraph(\n", - " handle=handle,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - ")\n", - "\n", - "# Create tensor handles with the graph API, assign UIDs.\n", - "x = graph.tensor_like(x_gpu.detach()).set_name(\"X\").set_uid(UID.X.value)\n", - "scale = graph.tensor_like(scale_gpu.detach()).set_name(\"scale\").set_uid(UID.SCALE.value)\n", - "bias = graph.tensor_like(bias_gpu.detach()).set_name(\"bias\").set_uid(UID.BIAS.value)\n", - "epsilon = graph.tensor_like(epsilon_cpu).set_name(\"epsilon\").set_uid(UID.EPSILON.value)\n", - "\n", - "# Add a layernorm operation\n", - "(out, mean, inv_var) = graph.layernorm(\n", - " name=\"layernorm\",\n", - " input=x,\n", - " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", - " scale=scale,\n", - " bias=bias,\n", - " epsilon=epsilon,\n", - ")\n", - "\n", - "# Enable all outputs, by default outputs are disabled\n", - "out.set_name(\"output\").set_output(True).set_data_type(out_expected.dtype).set_uid(\n", - " UID.OUT.value\n", - ")\n", - "mean.set_name(\"mean\").set_output(True).set_data_type(mean_expected.dtype).set_uid(\n", - " UID.MEAN.value\n", - ")\n", - "inv_var.set_name(\"inv_var\").set_output(True).set_data_type(\n", - " inv_var_expected.dtype\n", - ").set_uid(UID.INV_VAR.value)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Build the graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Build the graph\n", - "graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", - "\n", - "# To run this block more than once, we need to re-run the previous block to get a new graph.\n", - "# The same instance of a graph can not be built twice." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Execute the graph" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "After validating and building a cuDNN graph, we can now execute it. To do this, we have to provide input and output buffers. We do this by using the previously allocated UIDs to associate between tensor handles generated from the graph API, and their underlying memory. \n", - "\n", - "The desired input values need to be stored in these buffers before the `graph.execute` call. Because we have done a reference computation, we can simply reuse the buffers we have allocated via PyTorch.\n", - "\n", - "Note that the EPISLON UID expects a cpu buffer, " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Mapping of (UIDs -> memory)\n", - "variant_pack = {\n", - " UID.X.value: x_gpu,\n", - " UID.SCALE.value: scale_gpu,\n", - " UID.BIAS.value: bias_gpu,\n", - " UID.EPSILON.value: epsilon_cpu,\n", - " UID.OUT.value: out_gpu,\n", - " UID.MEAN.value: mean_gpu,\n", - " UID.INV_VAR.value: inv_var_gpu,\n", - "}\n", - "\n", - "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n", - "graph.execute(variant_pack, workspace)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Test cuDNN's output against PyTorch's and check correctness" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "torch.cuda.synchronize()\n", - "\n", - "# compare to reference output\n", - "torch.testing.assert_close(out_gpu, out_expected, rtol=5e-3, atol=5e-3)\n", - "torch.testing.assert_close(inv_var_gpu, inv_var_expected, rtol=5e-3, atol=5e-3)\n", - "torch.testing.assert_close(mean_gpu, mean_expected, rtol=5e-3, atol=5e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Perform Cleanup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "cudnn.destroy_handle(handle)" - ] - } - ], - "metadata": { - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# LayerNorm Forward Operation" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This notebook shows how to compute a layernorm forward operation using cuDNN.\n", + "\n", + "$$\\text{LayerNorm}(x) = \\frac{x-\\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\\cdot\\gamma+\\beta$$\n", + "\n", + "Where $\\mu = E[x]$ and $\\sigma^2 = Var[x]$ are taken over all inputs in a batch." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/20_layernorm_forward.ipynb)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites and Setup\n", + "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('nvidia-smi')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If running on Colab, you will need to install the cudnn python interface." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", + "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", + "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Overview" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the following, we are going to perform the layer norm forward pass with the following problem sizes:\n", + "\n", + "- batch size: 4\n", + "- sequence length: 1024\n", + "- embedding dimension: 768\n", + "\n", + "The tensor will be in 16-bit float format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import cudnn\n", + "import torch\n", + "\n", + "torch.manual_seed(1)\n", + "\n", + "handle = cudnn.create_handle()\n", + "\n", + "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", + "\n", + "assert torch.cuda.is_available()\n", + "\n", + "batch, seq_size, embedding_dim = 4, 1024, 768\n", + "dtype = torch.float16\n", + "# Epsilon is a small number to prevent division by 0.\n", + "epsilon_value = 1e-3" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Wrapper" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below is how you can use the `Graph` wrapper to perform layer norm with the input tensors in PyTorch format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# allocate random input tensors in (B, S, H) layout\n", + "x_gpu = torch.randn(batch, seq_size, embedding_dim, device=\"cuda\", dtype=dtype)\n", + "scale_gpu = torch.randn(1, 1, embedding_dim, device=\"cuda\", dtype=dtype)\n", + "bias_gpu = torch.randn(1, 1, embedding_dim, device=\"cuda\", dtype=dtype)\n", + "eps_cpu = torch.full((1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", + "\n", + "# create a graph: output in half precision, but mean and inv_var in single precision\n", + "with cudnn.Graph(\n", + " io_data_type=cudnn.data_type.HALF,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + " inputs=[\"ln::input\", \"ln::scale\", \"ln::bias\", \"ln::epsilon\"],\n", + " outputs=[\"ln::Y\", \"ln::MEAN\", \"ln::INV_VARIANCE\"],\n", + ") as graph:\n", + " out, mean, inv_var = graph.layernorm(\n", + " name=\"ln\",\n", + " input=x_gpu,\n", + " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", + " scale=scale_gpu,\n", + " bias=bias_gpu,\n", + " epsilon=eps_cpu,\n", + " )\n", + " out.set_output(True).set_data_type(dtype)\n", + " mean.set_output(True).set_data_type(torch.float32)\n", + " inv_var.set_output(True).set_data_type(torch.float32)\n", + "\n", + "# execute the graph and retrieve the output tensors\n", + "out_gpu, mean_gpu, inv_var_gpu = graph(\n", + " x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle\n", + ")\n", + "\n", + "# verify the result with PyTorch operations\n", + "out_ref = torch.nn.functional.layer_norm(\n", + " x_gpu,\n", + " [embedding_dim],\n", + " weight=scale_gpu.reshape(-1),\n", + " bias=bias_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "mean_ref = x_gpu.to(torch.float32).mean(dim=-1, keepdim=True)\n", + "inv_var_ref = torch.rsqrt(\n", + " torch.var(x_gpu.to(torch.float32), dim=-1, keepdim=True) + epsilon_value\n", + ")\n", + "\n", + "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(mean_gpu, mean_ref, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(inv_var_gpu, inv_var_ref, atol=5e-3, rtol=3e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Binding APIs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### General Setup\n", + "Create a cudnn handle, which is a per device handle used to initialize cudnn context." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### LayerNorm Training" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them easily when we calculate reference outputs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# allocate input tensor memory in (B, S, H) layout\n", + "x_gpu = torch.randn(\n", + " batch,\n", + " seq_size,\n", + " embedding_dim,\n", + " dtype=dtype,\n", + " requires_grad=True,\n", + " device=\"cuda\",\n", + ")\n", + "scale_gpu = torch.randn(\n", + " 1, 1, embedding_dim, dtype=dtype, requires_grad=True, device=\"cuda\"\n", + ")\n", + "bias_gpu = torch.randn(\n", + " 1, 1, embedding_dim, dtype=dtype, requires_grad=True, device=\"cuda\"\n", + ")\n", + "\n", + "# Epsilon must be a scalar value on the cpu.\n", + "epsilon_cpu = torch.full(\n", + " (1, 1, 1), epsilon_value, dtype=torch.float32, requires_grad=False, device=\"cpu\"\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compute reference ouputs and allocate output tensor GPU buffers" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Create the reference computation outputs here before the cuDNN computation, in order to use .empty_like() to create our output buffers\n", + "out_expected = torch.nn.functional.layer_norm(\n", + " x_gpu,\n", + " [embedding_dim],\n", + " weight=scale_gpu.reshape(-1),\n", + " bias=bias_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "\n", + "mean_expected = x_gpu.to(torch.float32).mean(dim=-1, keepdim=True)\n", + "\n", + "inv_var_expected = torch.rsqrt(\n", + " torch.var(x_gpu.to(torch.float32), dim=-1, keepdim=True) + epsilon_value\n", + ")\n", + "\n", + "# Allocate output tensor memory using PyTorch\n", + "# PyTorch has calculated their shapes already, so we can simply use .empty_like()\n", + "out_gpu = torch.empty_like(out_expected)\n", + "mean_gpu = torch.empty_like(mean_expected)\n", + "inv_var_gpu = torch.empty_like(inv_var_expected)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Create cuDNN graph\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here we assign UIDs for tensors. UIDs are a unique identifier that will allow us to provide a mapping from tensors created from cuDNN graph api calls, such as `graph.tensor_like()`, to the underlying device memory that will be used to store these tensors. Virtual tensors don't require explicit memory allocated for them, but non-vritual tensors like inputs or outputs will need to have UIDs assigned to them. \n", + "\n", + "Alternatively, one can use handles directly in the mapping, however using UIDs can be more convinient for caching of cuDNN graphs.\n", + "\n", + "For each of our inputs {X, Scale, Bias, Epsilon} and our outputs {Out, Mean, Inverse Variance}, we allocate a UID. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from enum import Enum\n", + "\n", + "\n", + "class UID(Enum):\n", + " X = 0\n", + " SCALE = 1\n", + " BIAS = 2\n", + " EPSILON = 3\n", + " OUT = 4\n", + " MEAN = 5\n", + " INV_VAR = 6" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Create the cuDNN graph.\n", + "graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + ")\n", + "\n", + "# Create tensor handles with the graph API, assign UIDs.\n", + "x = graph.tensor_like(x_gpu.detach()).set_name(\"X\").set_uid(UID.X.value)\n", + "scale = graph.tensor_like(scale_gpu.detach()).set_name(\"scale\").set_uid(UID.SCALE.value)\n", + "bias = graph.tensor_like(bias_gpu.detach()).set_name(\"bias\").set_uid(UID.BIAS.value)\n", + "epsilon = graph.tensor_like(epsilon_cpu).set_name(\"epsilon\").set_uid(UID.EPSILON.value)\n", + "\n", + "# Add a layernorm operation\n", + "(out, mean, inv_var) = graph.layernorm(\n", + " name=\"layernorm\",\n", + " input=x,\n", + " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", + " scale=scale,\n", + " bias=bias,\n", + " epsilon=epsilon,\n", + ")\n", + "\n", + "# Enable all outputs, by default outputs are disabled\n", + "out.set_name(\"output\").set_output(True).set_data_type(out_expected.dtype).set_uid(\n", + " UID.OUT.value\n", + ")\n", + "mean.set_name(\"mean\").set_output(True).set_data_type(mean_expected.dtype).set_uid(\n", + " UID.MEAN.value\n", + ")\n", + "inv_var.set_name(\"inv_var\").set_output(True).set_data_type(\n", + " inv_var_expected.dtype\n", + ").set_uid(UID.INV_VAR.value)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Build the graph" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Build the graph\n", + "graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", + "\n", + "# To run this block more than once, we need to re-run the previous block to get a new graph.\n", + "# The same instance of a graph can not be built twice." + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Execute the graph" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After validating and building a cuDNN graph, we can now execute it. To do this, we have to provide input and output buffers. We do this by using the previously allocated UIDs to associate between tensor handles generated from the graph API, and their underlying memory. \n", + "\n", + "The desired input values need to be stored in these buffers before the `graph.execute` call. Because we have done a reference computation, we can simply reuse the buffers we have allocated via PyTorch.\n", + "\n", + "Note that the EPISLON UID expects a cpu buffer, " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Mapping of (UIDs -> memory)\n", + "variant_pack = {\n", + " UID.X.value: x_gpu,\n", + " UID.SCALE.value: scale_gpu,\n", + " UID.BIAS.value: bias_gpu,\n", + " UID.EPSILON.value: epsilon_cpu,\n", + " UID.OUT.value: out_gpu,\n", + " UID.MEAN.value: mean_gpu,\n", + " UID.INV_VAR.value: inv_var_gpu,\n", + "}\n", + "\n", + "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n", + "graph.execute(variant_pack, workspace)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Test cuDNN's output against PyTorch's and check correctness" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "torch.cuda.synchronize()\n", + "\n", + "# compare to reference output\n", + "torch.testing.assert_close(out_gpu, out_expected, rtol=5e-3, atol=5e-3)\n", + "torch.testing.assert_close(inv_var_gpu, inv_var_expected, rtol=5e-3, atol=5e-3)\n", + "torch.testing.assert_close(mean_gpu, mean_expected, rtol=5e-3, atol=5e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Perform Cleanup" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cudnn.destroy_handle(handle)" + ], + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/samples/python/21_layernorm_backward.ipynb b/samples/python/21_layernorm_backward.ipynb index c6834236d..a7d9eeeb8 100644 --- a/samples/python/21_layernorm_backward.ipynb +++ b/samples/python/21_layernorm_backward.ipynb @@ -1,643 +1,641 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# LayerNorm Backward Operation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This notebook shows how the backwards pass of a layernorm operation can be done using cudnn.\n", - "\n", - "$$\\text{LayerNorm}(x) = \\frac{x-\\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\\cdot\\gamma+\\beta$$" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/21_layernorm_backward.ipynb)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites and Setup\n", - "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('nvidia-smi')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If running on Colab, you will need to install the cudnn python interface." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", - "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", - "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Overview" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we are going to perform the layer norm forward and backward pass with the following problem sizes:\n", - "\n", - "- batch size: 4\n", - "- sequence length: 1024\n", - "- embedding dimension: 768\n", - "\n", - "The tensor will be in 16-bit float format." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import cudnn\n", - "import torch\n", - "\n", - "torch.manual_seed(1)\n", - "\n", - "handle = cudnn.create_handle()\n", - "\n", - "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", - "\n", - "assert torch.cuda.is_available()\n", - "\n", - "batch, seq_size, embedding_dim = 4, 1024, 768\n", - "dtype = torch.float16\n", - "# Epsilon is a small number to prevent division by 0.\n", - "epsilon_value = 1e-3" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Wrapper" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below is how you can use the `Graph` wrapper to perform forward layer norm with the input tensors in PyTorch format, then compute a loss metric and perform the backward pass.\n", - "\n", - "First is the forward pass, which is mostly the same as the one in the [forward pass notebook](20_layernorm_forward.ipynb)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# allocate random input tensors\n", - "x_gpu = torch.randn(\n", - " batch * seq_size,\n", - " embedding_dim,\n", - " 1,\n", - " 1,\n", - " device=\"cuda\",\n", - " dtype=dtype,\n", - " requires_grad=True,\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "bias_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "eps_cpu = torch.full((1, 1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", - "\n", - "# create a graph: output in half precision, but mean and inv_var in single precision\n", - "with cudnn.Graph(\n", - " io_data_type=cudnn.data_type.HALF,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - " inputs=[\"ln::input\", \"ln::scale\", \"ln::bias\", \"ln::epsilon\"],\n", - " outputs=[\"ln::Y\", \"ln::MEAN\", \"ln::INV_VARIANCE\"],\n", - ") as graph:\n", - " out, mean, inv_var = graph.layernorm(\n", - " name=\"ln\",\n", - " input=x_gpu,\n", - " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", - " scale=scale_gpu,\n", - " bias=bias_gpu,\n", - " epsilon=eps_cpu,\n", - " )\n", - " out.set_output(True).set_data_type(dtype)\n", - " mean.set_output(True).set_data_type(torch.float32)\n", - " inv_var.set_output(True).set_data_type(torch.float32)\n", - "\n", - "# execute the graph and retrieve the output tensors\n", - "out_gpu, mean_gpu, inv_var_gpu = graph(\n", - " x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle\n", - ")\n", - "\n", - "# verify the result with PyTorch operations\n", - "out_ref = torch.nn.functional.layer_norm(\n", - " x_gpu,\n", - " [embedding_dim, 1, 1],\n", - " weight=scale_gpu.squeeze(0),\n", - " bias=bias_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "mean_ref = x_gpu.to(torch.float32).mean(dim=(1, 2, 3), keepdim=True)\n", - "inv_var_ref = torch.rsqrt(\n", - " torch.var(x_gpu.to(torch.float32), dim=(1, 2, 3), keepdim=True) + epsilon_value\n", - ")\n", - "\n", - "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(mean_gpu, mean_ref, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(inv_var_gpu, inv_var_ref, atol=5e-3, rtol=3e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, we generate a random tensor and assume that is the ground truth output. Then compute the L2 loss for backward pass:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Compute gradients: Ask PyTorch not to discard the grads after use so that we can read it twice\n", - "# out_ref.grad will be used in the cudnn graph, x_gpu.grad, scale_gpu.grad, and bias_gpu.grad will\n", - "# be used to compare with the cudnn graph output.\n", - "target = torch.randn_like(out_ref)\n", - "criterion = torch.nn.MSELoss()\n", - "loss = criterion(out_ref, target)\n", - "\n", - "out_ref.retain_grad()\n", - "x_gpu.retain_grad()\n", - "scale_gpu.retain_grad()\n", - "bias_gpu.retain_grad()\n", - "\n", - "loss.backward()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Then, let's do the backward pass." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Backward pass\n", - "with cudnn.Graph(\n", - " io_data_type=cudnn.data_type.HALF,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - " inputs=[\n", - " \"ln_bwd::grad\",\n", - " \"ln_bwd::input\",\n", - " \"ln_bwd::scale\",\n", - " \"ln_bwd::mean\",\n", - " \"ln_bwd::inv_variance\",\n", - " ],\n", - " outputs=[\"ln_bwd::DX\", \"ln_bwd::DSCALE\", \"ln_bwd::DBIAS\"],\n", - ") as bwd_graph:\n", - " dx, dscale, dbias = bwd_graph.layernorm_backward(\n", - " name=\"ln_bwd\",\n", - " grad=out_ref,\n", - " input=x_gpu,\n", - " scale=scale_gpu,\n", - " mean=mean_gpu,\n", - " inv_variance=inv_var_gpu,\n", - " )\n", - " dx.set_output(True).set_data_type(dtype)\n", - " dscale.set_output(True).set_data_type(dtype)\n", - " dbias.set_output(True).set_data_type(dtype)\n", - "\n", - "# execute the graph and retrieve the output tensors\n", - "dx_gpu, dscale_gpu, dbias_gpu = bwd_graph(\n", - " out_ref.grad, x_gpu, scale_gpu, mean_gpu, inv_var_gpu, handle=handle\n", - ")\n", - "\n", - "torch.testing.assert_close(x_gpu.grad, dx_gpu, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(scale_gpu.grad, dscale_gpu, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(bias_gpu.grad, dbias_gpu, atol=5e-3, rtol=3e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The input `grad` to the backward graph is the gradient of layer norm output computed by PyTorch triggered by `loss.backward()` before. When you run the graph with PyTorch tensors, you use `detach()` to detach the tensors from the PyTorch computational graph. The output from the graph are compared against the corresponding gradients from PyTorch." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Python Binding APIs" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### LayerNorm Training" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them easily when we calculate reference outputs." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# input tensor memory, initialize them to random numbers\n", - "x_gpu = torch.randn(\n", - " batch * seq_size,\n", - " embedding_dim,\n", - " 1,\n", - " 1,\n", - " device=\"cuda\",\n", - " dtype=dtype,\n", - " requires_grad=True,\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "bias_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "# set epsilon to epsilon_value, allocate on cpu.\n", - "eps_cpu = torch.full((1, 1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create reference computation and output tensor GPU buffers using PyTorch" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# create the reference computation outputs here before the cuDNN computation, in order to use .empty_like() to create our output buffers\n", - "out_expected = torch.nn.functional.layer_norm(\n", - " x_gpu,\n", - " [embedding_dim, 1, 1],\n", - " weight=scale_gpu.squeeze(0),\n", - " bias=bias_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "\n", - "mean_expected = x_gpu.to(torch.float32).mean(dim=(1, 2, 3), keepdim=True)\n", - "\n", - "inv_var_expected = torch.rsqrt(\n", - " torch.var(x_gpu.to(torch.float32), dim=(1, 2, 3), keepdim=True) + epsilon_value\n", - ")\n", - "\n", - "# allocate output tensor memory using PyTorch\n", - "out_gpu = torch.empty_like(out_expected)\n", - "mean_gpu = torch.empty_like(mean_expected)\n", - "inv_var_gpu = torch.empty_like(inv_var_expected)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Create cuDNN Foward Graph and tensors" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create the cuDNN graph\n", - "fwd_graph = cudnn.pygraph(\n", - " handle=handle,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - ")\n", - "\n", - "# Create tensor handles with the fwd_graph\n", - "x = fwd_graph.tensor_like(x_gpu.detach()).set_name(\"X\")\n", - "scale = fwd_graph.tensor_like(scale_gpu.detach()).set_name(\"scale\")\n", - "bias = fwd_graph.tensor_like(bias_gpu.detach()).set_name(\"bias\")\n", - "epsilon = fwd_graph.tensor_like(eps_cpu).set_name(\"epsilon\")\n", - "\n", - "# Add a layernorm operation\n", - "(out, mean, inv_var) = fwd_graph.layernorm(\n", - " name=\"layernorm\",\n", - " input=x,\n", - " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", - " scale=scale,\n", - " bias=bias,\n", - " epsilon=epsilon,\n", - ")\n", - "\n", - "# Enable all outputs\n", - "out.set_name(\"output\").set_output(True).set_data_type(out_expected.dtype)\n", - "mean.set_name(\"mean\").set_output(True).set_data_type(mean_expected.dtype)\n", - "inv_var.set_name(\"inv_var\").set_output(True).set_data_type(inv_var_expected.dtype);" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Validate and build the forward graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Build the fwd_graph\n", - "fwd_graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", - "\n", - "# To run this block more than once, we need to re-run the previous block to get a new fwd_graph.\n", - "# The same instance of a fwd_graph should not be built twice." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Execute the forward graph" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Instead of mapping UIDs to memory (as in 20_layernorm.ipynb), we can directly map handles to memory. This is simpler but slightly slower to execute." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Mapping of (handles -> memory)\n", - "variant_pack = {\n", - " x: x_gpu.detach(),\n", - " scale: scale_gpu.detach(),\n", - " bias: bias_gpu.detach(),\n", - " epsilon: eps_cpu,\n", - " out: out_gpu,\n", - " mean: mean_gpu,\n", - " inv_var: inv_var_gpu,\n", - "}\n", - "\n", - "workspace = torch.empty(\n", - " fwd_graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8\n", - ")\n", - "fwd_graph.execute(variant_pack, workspace)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Test cuDNN's output against PyTorch's and check correctness" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "torch.cuda.synchronize()\n", - "\n", - "# reference output\n", - "torch.testing.assert_close(out_gpu, out_expected, rtol=5e-3, atol=5e-3)\n", - "torch.testing.assert_close(inv_var_gpu, inv_var_expected, rtol=5e-3, atol=5e-3)\n", - "torch.testing.assert_close(mean_gpu, mean_expected, rtol=5e-3, atol=5e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### LayerNorm Backwards Pass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Compute references values for backwards graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Reference backwards operation using PyTorch\n", - "target = torch.randn_like(out_expected)\n", - "criterion = torch.nn.MSELoss()\n", - "loss = criterion(out_expected, target)\n", - "\n", - "out_expected.retain_grad()\n", - "x_gpu.retain_grad()\n", - "scale_gpu.retain_grad()\n", - "bias_gpu.retain_grad()\n", - "\n", - "loss.backward()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Build backwards graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "bwd_graph = cudnn.pygraph(\n", - " handle=handle,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - ")\n", - "\n", - "# Create tensors associated with the backwards graph. DO NOT reuse tensor handles from the forward graph.\n", - "d_out = bwd_graph.tensor(\n", - " name=\"d_out\", dim=x_gpu.size(), stride=x_gpu.stride(), data_type=x_gpu.dtype\n", - ")\n", - "\n", - "x_bwd = bwd_graph.tensor_like(x, name=\"x\")\n", - "scale_bwd = bwd_graph.tensor_like(scale, name=\"scale\")\n", - "mean_bwd = bwd_graph.tensor_like(mean, name=\"mean\")\n", - "inv_var_bwd = bwd_graph.tensor_like(inv_var, name=\"inv_var\")\n", - "\n", - "# Add the layernorm backwards operation\n", - "(d_x, d_scale, d_bias) = bwd_graph.layernorm_backward(\n", - " name=\"d_layernorm\",\n", - " grad=d_out,\n", - " input=x_bwd,\n", - " scale=scale_bwd,\n", - " mean=mean_bwd,\n", - " inv_variance=inv_var_bwd,\n", - ")\n", - "\n", - "# Enable outputs.\n", - "d_x.set_output(True).set_data_type(x_gpu.dtype)\n", - "d_scale.set_output(True).set_data_type(x_gpu.dtype)\n", - "d_bias.set_output(True).set_data_type(x_gpu.dtype)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Build the bwd_graph\n", - "bwd_graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Execute the graph and check correctness against PyTorch" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create output buffers for gradients\n", - "d_x_gpu = torch.empty_like(x_gpu)\n", - "d_scale_gpu = torch.empty_like(scale_gpu)\n", - "d_bias_gpu = torch.empty_like(bias_gpu)\n", - "\n", - "workspace = torch.empty(\n", - " bwd_graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8\n", - ")\n", - "\n", - "# For the inputs of the backwards graph (x_bwd, d_out, scale_bwd, mean_bwd, inv_var_bwd), we use the outputs of the forwards graph. For d_out we use pytorches autograd .grad functionality.\n", - "bwd_graph.execute(\n", - " {\n", - " x_bwd: x_gpu.detach(),\n", - " scale_bwd: scale_gpu.detach(),\n", - " d_out: out_expected.grad,\n", - " mean_bwd: mean_gpu.detach(),\n", - " inv_var_bwd: inv_var_gpu.detach(),\n", - " d_x: d_x_gpu,\n", - " d_scale: d_scale_gpu,\n", - " d_bias: d_bias_gpu,\n", - " },\n", - " workspace,\n", - " handle=handle,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Compare results and check correctness" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "torch.cuda.synchronize()\n", - "\n", - "# compare to reference output\n", - "torch.testing.assert_close(x_gpu.grad, d_x_gpu, atol=2e-4, rtol=2e-4)\n", - "torch.testing.assert_close(scale_gpu.grad, d_scale_gpu, atol=2e-4, rtol=2e-4)\n", - "torch.testing.assert_close(bias_gpu.grad, d_bias_gpu, atol=2e-4, rtol=2e-4)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Perform Cleanup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "cudnn.destroy_handle(handle)" - ] - } - ], - "metadata": { - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# LayerNorm Backward Operation" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This notebook shows how the backwards pass of a layernorm operation can be done using cudnn.\n", + "\n", + "$$\\text{LayerNorm}(x) = \\frac{x-\\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\\cdot\\gamma+\\beta$$" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/21_layernorm_backward.ipynb)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites and Setup\n", + "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('nvidia-smi')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If running on Colab, you will need to install the cudnn python interface." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", + "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", + "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Overview" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the following, we are going to perform the layer norm forward and backward pass with the following problem sizes:\n", + "\n", + "- batch size: 4\n", + "- sequence length: 1024\n", + "- embedding dimension: 768\n", + "\n", + "The tensor will be in 16-bit float format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import cudnn\n", + "import torch\n", + "\n", + "torch.manual_seed(1)\n", + "\n", + "handle = cudnn.create_handle()\n", + "\n", + "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", + "\n", + "assert torch.cuda.is_available()\n", + "\n", + "batch, seq_size, embedding_dim = 4, 1024, 768\n", + "dtype = torch.float16\n", + "# Epsilon is a small number to prevent division by 0.\n", + "epsilon_value = 1e-3" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Wrapper" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below is how you can use the `Graph` wrapper to perform forward layer norm with the input tensors in PyTorch format, then compute a loss metric and perform the backward pass.\n", + "\n", + "First is the forward pass, which is mostly the same as the one in the [forward pass notebook](20_layernorm_forward.ipynb)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# allocate random input tensors in (B, S, H) layout\n", + "x_gpu = torch.randn(\n", + " batch,\n", + " seq_size,\n", + " embedding_dim,\n", + " device=\"cuda\",\n", + " dtype=dtype,\n", + " requires_grad=True,\n", + ")\n", + "scale_gpu = torch.randn(\n", + " 1, 1, embedding_dim, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "bias_gpu = torch.randn(\n", + " 1, 1, embedding_dim, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "eps_cpu = torch.full((1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", + "\n", + "# create a graph: output in half precision, but mean and inv_var in single precision\n", + "with cudnn.Graph(\n", + " io_data_type=cudnn.data_type.HALF,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + " inputs=[\"ln::input\", \"ln::scale\", \"ln::bias\", \"ln::epsilon\"],\n", + " outputs=[\"ln::Y\", \"ln::MEAN\", \"ln::INV_VARIANCE\"],\n", + ") as graph:\n", + " out, mean, inv_var = graph.layernorm(\n", + " name=\"ln\",\n", + " input=x_gpu,\n", + " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", + " scale=scale_gpu,\n", + " bias=bias_gpu,\n", + " epsilon=eps_cpu,\n", + " )\n", + " out.set_output(True).set_data_type(dtype)\n", + " mean.set_output(True).set_data_type(torch.float32)\n", + " inv_var.set_output(True).set_data_type(torch.float32)\n", + "\n", + "# execute the graph and retrieve the output tensors\n", + "out_gpu, mean_gpu, inv_var_gpu = graph(\n", + " x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle\n", + ")\n", + "\n", + "# verify the result with PyTorch operations\n", + "out_ref = torch.nn.functional.layer_norm(\n", + " x_gpu,\n", + " [embedding_dim],\n", + " weight=scale_gpu.reshape(-1),\n", + " bias=bias_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "mean_ref = x_gpu.to(torch.float32).mean(dim=-1, keepdim=True)\n", + "inv_var_ref = torch.rsqrt(\n", + " torch.var(x_gpu.to(torch.float32), dim=-1, keepdim=True) + epsilon_value\n", + ")\n", + "\n", + "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(mean_gpu, mean_ref, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(inv_var_gpu, inv_var_ref, atol=5e-3, rtol=3e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, we generate a random tensor and assume that is the ground truth output. Then compute the L2 loss for backward pass:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Compute gradients: Ask PyTorch not to discard the grads after use so that we can read it twice\n", + "# out_ref.grad will be used in the cudnn graph, x_gpu.grad, scale_gpu.grad, and bias_gpu.grad will\n", + "# be used to compare with the cudnn graph output.\n", + "target = torch.randn_like(out_ref)\n", + "criterion = torch.nn.MSELoss()\n", + "loss = criterion(out_ref, target)\n", + "\n", + "out_ref.retain_grad()\n", + "x_gpu.retain_grad()\n", + "scale_gpu.retain_grad()\n", + "bias_gpu.retain_grad()\n", + "\n", + "loss.backward()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Then, let's do the backward pass." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Backward pass\n", + "with cudnn.Graph(\n", + " io_data_type=cudnn.data_type.HALF,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + " inputs=[\n", + " \"ln_bwd::grad\",\n", + " \"ln_bwd::input\",\n", + " \"ln_bwd::scale\",\n", + " \"ln_bwd::mean\",\n", + " \"ln_bwd::inv_variance\",\n", + " ],\n", + " outputs=[\"ln_bwd::DX\", \"ln_bwd::DSCALE\", \"ln_bwd::DBIAS\"],\n", + ") as bwd_graph:\n", + " dx, dscale, dbias = bwd_graph.layernorm_backward(\n", + " name=\"ln_bwd\",\n", + " grad=out_ref,\n", + " input=x_gpu,\n", + " scale=scale_gpu,\n", + " mean=mean_gpu,\n", + " inv_variance=inv_var_gpu,\n", + " )\n", + " dx.set_output(True).set_data_type(dtype)\n", + " dscale.set_output(True).set_data_type(dtype)\n", + " dbias.set_output(True).set_data_type(dtype)\n", + "\n", + "# execute the graph and retrieve the output tensors\n", + "dx_gpu, dscale_gpu, dbias_gpu = bwd_graph(\n", + " out_ref.grad, x_gpu, scale_gpu, mean_gpu, inv_var_gpu, handle=handle\n", + ")\n", + "\n", + "torch.testing.assert_close(x_gpu.grad, dx_gpu, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(scale_gpu.grad, dscale_gpu, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(bias_gpu.grad, dbias_gpu, atol=5e-3, rtol=3e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The input `grad` to the backward graph is the gradient of layer norm output computed by PyTorch triggered by `loss.backward()` before. When you run the graph with PyTorch tensors, you use `detach()` to detach the tensors from the PyTorch computational graph. The output from the graph are compared against the corresponding gradients from PyTorch." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Binding APIs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### LayerNorm Training" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them easily when we calculate reference outputs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# input tensor memory in (B, S, H) layout\n", + "x_gpu = torch.randn(\n", + " batch,\n", + " seq_size,\n", + " embedding_dim,\n", + " device=\"cuda\",\n", + " dtype=dtype,\n", + " requires_grad=True,\n", + ")\n", + "scale_gpu = torch.randn(\n", + " 1, 1, embedding_dim, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "bias_gpu = torch.randn(\n", + " 1, 1, embedding_dim, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "# set epsilon to epsilon_value, allocate on cpu.\n", + "eps_cpu = torch.full((1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create reference computation and output tensor GPU buffers using PyTorch" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# create the reference computation outputs here before the cuDNN computation, in order to use .empty_like() to create our output buffers\n", + "out_expected = torch.nn.functional.layer_norm(\n", + " x_gpu,\n", + " [embedding_dim],\n", + " weight=scale_gpu.reshape(-1),\n", + " bias=bias_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "\n", + "mean_expected = x_gpu.to(torch.float32).mean(dim=-1, keepdim=True)\n", + "\n", + "inv_var_expected = torch.rsqrt(\n", + " torch.var(x_gpu.to(torch.float32), dim=-1, keepdim=True) + epsilon_value\n", + ")\n", + "\n", + "# allocate output tensor memory using PyTorch\n", + "out_gpu = torch.empty_like(out_expected)\n", + "mean_gpu = torch.empty_like(mean_expected)\n", + "inv_var_gpu = torch.empty_like(inv_var_expected)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Create cuDNN Foward Graph and tensors" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Create the cuDNN graph\n", + "fwd_graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + ")\n", + "\n", + "# Create tensor handles with the fwd_graph\n", + "x = fwd_graph.tensor_like(x_gpu.detach()).set_name(\"X\")\n", + "scale = fwd_graph.tensor_like(scale_gpu.detach()).set_name(\"scale\")\n", + "bias = fwd_graph.tensor_like(bias_gpu.detach()).set_name(\"bias\")\n", + "epsilon = fwd_graph.tensor_like(eps_cpu).set_name(\"epsilon\")\n", + "\n", + "# Add a layernorm operation\n", + "(out, mean, inv_var) = fwd_graph.layernorm(\n", + " name=\"layernorm\",\n", + " input=x,\n", + " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", + " scale=scale,\n", + " bias=bias,\n", + " epsilon=epsilon,\n", + ")\n", + "\n", + "# Enable all outputs\n", + "out.set_name(\"output\").set_output(True).set_data_type(out_expected.dtype)\n", + "mean.set_name(\"mean\").set_output(True).set_data_type(mean_expected.dtype)\n", + "inv_var.set_name(\"inv_var\").set_output(True).set_data_type(inv_var_expected.dtype);" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Validate and build the forward graph" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Build the fwd_graph\n", + "fwd_graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", + "\n", + "# To run this block more than once, we need to re-run the previous block to get a new fwd_graph.\n", + "# The same instance of a fwd_graph should not be built twice." + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Execute the forward graph" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Instead of mapping UIDs to memory (as in 20_layernorm.ipynb), we can directly map handles to memory. This is simpler but slightly slower to execute." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Mapping of (handles -> memory)\n", + "variant_pack = {\n", + " x: x_gpu.detach(),\n", + " scale: scale_gpu.detach(),\n", + " bias: bias_gpu.detach(),\n", + " epsilon: eps_cpu,\n", + " out: out_gpu,\n", + " mean: mean_gpu,\n", + " inv_var: inv_var_gpu,\n", + "}\n", + "\n", + "workspace = torch.empty(\n", + " fwd_graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8\n", + ")\n", + "fwd_graph.execute(variant_pack, workspace)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Test cuDNN's output against PyTorch's and check correctness" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "torch.cuda.synchronize()\n", + "\n", + "# reference output\n", + "torch.testing.assert_close(out_gpu, out_expected, rtol=5e-3, atol=5e-3)\n", + "torch.testing.assert_close(inv_var_gpu, inv_var_expected, rtol=5e-3, atol=5e-3)\n", + "torch.testing.assert_close(mean_gpu, mean_expected, rtol=5e-3, atol=5e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### LayerNorm Backwards Pass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compute references values for backwards graph" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Reference backwards operation using PyTorch\n", + "target = torch.randn_like(out_expected)\n", + "criterion = torch.nn.MSELoss()\n", + "loss = criterion(out_expected, target)\n", + "\n", + "out_expected.retain_grad()\n", + "x_gpu.retain_grad()\n", + "scale_gpu.retain_grad()\n", + "bias_gpu.retain_grad()\n", + "\n", + "loss.backward()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Build backwards graph" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bwd_graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + ")\n", + "\n", + "# Create tensors associated with the backwards graph. DO NOT reuse tensor handles from the forward graph.\n", + "d_out = bwd_graph.tensor(\n", + " name=\"d_out\", dim=x_gpu.size(), stride=x_gpu.stride(), data_type=x_gpu.dtype\n", + ")\n", + "\n", + "x_bwd = bwd_graph.tensor_like(x, name=\"x\")\n", + "scale_bwd = bwd_graph.tensor_like(scale, name=\"scale\")\n", + "mean_bwd = bwd_graph.tensor_like(mean, name=\"mean\")\n", + "inv_var_bwd = bwd_graph.tensor_like(inv_var, name=\"inv_var\")\n", + "\n", + "# Add the layernorm backwards operation\n", + "(d_x, d_scale, d_bias) = bwd_graph.layernorm_backward(\n", + " name=\"d_layernorm\",\n", + " grad=d_out,\n", + " input=x_bwd,\n", + " scale=scale_bwd,\n", + " mean=mean_bwd,\n", + " inv_variance=inv_var_bwd,\n", + ")\n", + "\n", + "# Enable outputs.\n", + "d_x.set_output(True).set_data_type(x_gpu.dtype)\n", + "d_scale.set_output(True).set_data_type(x_gpu.dtype)\n", + "d_bias.set_output(True).set_data_type(x_gpu.dtype)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Build the bwd_graph\n", + "bwd_graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Execute the graph and check correctness against PyTorch" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Create output buffers for gradients\n", + "d_x_gpu = torch.empty_like(x_gpu)\n", + "d_scale_gpu = torch.empty_like(scale_gpu)\n", + "d_bias_gpu = torch.empty_like(bias_gpu)\n", + "\n", + "workspace = torch.empty(\n", + " bwd_graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8\n", + ")\n", + "\n", + "# For the inputs of the backwards graph (x_bwd, d_out, scale_bwd, mean_bwd, inv_var_bwd), we use the outputs of the forwards graph. For d_out we use pytorches autograd .grad functionality.\n", + "bwd_graph.execute(\n", + " {\n", + " x_bwd: x_gpu.detach(),\n", + " scale_bwd: scale_gpu.detach(),\n", + " d_out: out_expected.grad,\n", + " mean_bwd: mean_gpu.detach(),\n", + " inv_var_bwd: inv_var_gpu.detach(),\n", + " d_x: d_x_gpu,\n", + " d_scale: d_scale_gpu,\n", + " d_bias: d_bias_gpu,\n", + " },\n", + " workspace,\n", + " handle=handle,\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compare results and check correctness" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "torch.cuda.synchronize()\n", + "\n", + "# compare to reference output\n", + "torch.testing.assert_close(x_gpu.grad, d_x_gpu, atol=2e-4, rtol=2e-4)\n", + "torch.testing.assert_close(scale_gpu.grad, d_scale_gpu, atol=2e-4, rtol=2e-4)\n", + "torch.testing.assert_close(bias_gpu.grad, d_bias_gpu, atol=2e-4, rtol=2e-4)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Perform Cleanup" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cudnn.destroy_handle(handle)" + ], + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/samples/python/22_layernorm_inference.ipynb b/samples/python/22_layernorm_inference.ipynb index 50b9fc71d..d876ac906 100644 --- a/samples/python/22_layernorm_inference.ipynb +++ b/samples/python/22_layernorm_inference.ipynb @@ -1,380 +1,373 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# LayerNorm Inference Operation\n", - "\n", - "This notebook shows how a layernorm operation can be done using cuDNN in inference mode. This is different from the forward pass in training mode that the mean and variance are not computed nor stored.\n", - "\n", - "$$\\text{LayerNorm}(x) = \\frac{x-\\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\\cdot\\gamma+\\beta$$" - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# LayerNorm Inference Operation\n", + "\n", + "This notebook shows how a layernorm operation can be done using cuDNN in inference mode. This is different from the forward pass in training mode that the mean and variance are not computed nor stored.\n", + "\n", + "$$\\text{LayerNorm}(x) = \\frac{x-\\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\\cdot\\gamma+\\beta$$" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/22_layernorm_inference.ipynb)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites and Setup\n", + "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('nvidia-smi')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If running on Colab, you will need to install the cudnn python interface." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", + "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", + "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Overview" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the following, we are going to perform the layer norm forward pass with the following problem sizes:\n", + "\n", + "- batch size: 4\n", + "- sequence length: 1024\n", + "- embedding dimension: 768\n", + "\n", + "The tensor will be in 16-bit float format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import cudnn\n", + "import torch\n", + "\n", + "torch.manual_seed(1)\n", + "handle = cudnn.create_handle()\n", + "\n", + "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", + "\n", + "assert torch.cuda.is_available()\n", + "\n", + "batch, seq_size, embedding_dim = 4, 1024, 768\n", + "dtype = torch.float16\n", + "# Epsilon is a small number to prevent division by 0.\n", + "epsilon_value = 1e-3" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Wrapper" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below is how you can use the `Graph` wrapper to perform forward layer norm with the input tensors in PyTorch.\n", + "This is highly similar to the one in the [forward pass notebook](20_layernorm_forward.ipynb)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# allocate random input tensors in (B, S, H) layout\n", + "x_gpu = torch.randn(batch, seq_size, embedding_dim, device=\"cuda\", dtype=dtype)\n", + "scale_gpu = torch.randn(1, 1, embedding_dim, device=\"cuda\", dtype=dtype)\n", + "bias_gpu = torch.randn(1, 1, embedding_dim, device=\"cuda\", dtype=dtype)\n", + "eps_cpu = torch.full((1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", + "\n", + "# create a graph\n", + "with cudnn.Graph(\n", + " io_data_type=cudnn.data_type.HALF,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + " inputs=[\"ln::input\", \"ln::scale\", \"ln::bias\", \"ln::epsilon\"],\n", + " outputs=[\"ln::Y\"],\n", + ") as graph:\n", + " out, mean, inv_var = graph.layernorm(\n", + " name=\"ln\",\n", + " input=x_gpu,\n", + " norm_forward_phase=cudnn.norm_forward_phase.INFERENCE,\n", + " scale=scale_gpu,\n", + " bias=bias_gpu,\n", + " epsilon=eps_cpu,\n", + " )\n", + " out.set_output(True).set_data_type(dtype)\n", + " assert mean is None, \"Expecting mean to be None under inference mode\"\n", + " assert inv_var is None, \"Expecting inv_var to be None under inference mode\"\n", + "\n", + "# execute the graph and retrieve the output tensors\n", + "out_gpu = graph(x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle)\n", + "\n", + "# verify the result with PyTorch operations\n", + "out_ref = torch.nn.functional.layer_norm(\n", + " x_gpu,\n", + " [embedding_dim],\n", + " weight=scale_gpu.reshape(-1),\n", + " bias=bias_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "\n", + "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compared to the one in the [forward pass notebook](20_layernorm_forward.ipynb), the argument `norm_forward_phase` is set to `cudnn.norm_forward_phase.INFERENCE` instead of `cudnn.norm_forward_phase.TRAINING`. The result of this is that the mean and variance are not computed nor stored." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Binding APIs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### LayerNorm Inference" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them to calculate a reference value." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# allocate input tensor memory in (B, S, H) layout\n", + "x_gpu = torch.randn(\n", + " batch,\n", + " seq_size,\n", + " embedding_dim,\n", + " dtype=dtype,\n", + " requires_grad=True,\n", + " device=\"cuda\",\n", + ")\n", + "scale_gpu = torch.randn(\n", + " 1, 1, embedding_dim, dtype=dtype, requires_grad=True, device=\"cuda\"\n", + ")\n", + "bias_gpu = torch.randn(\n", + " 1, 1, embedding_dim, dtype=dtype, requires_grad=True, device=\"cuda\"\n", + ")\n", + "\n", + "# Epsilon must be a scalar value on the cpu.\n", + "epsilon_cpu = torch.full(\n", + " (1, 1, 1), epsilon_value, dtype=torch.float32, requires_grad=False, device=\"cpu\"\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create reference computation and allocate output tensor GPU buffers" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Create the reference computation outputs here so we can use .empty_like() to create our output buffers\n", + "out_expected = torch.nn.functional.layer_norm(\n", + " x_gpu,\n", + " [embedding_dim],\n", + " weight=scale_gpu.reshape(-1),\n", + " bias=bias_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "\n", + "\n", + "# Allocate output tensor memory using PyTorch\n", + "out_gpu = torch.empty_like(out_expected)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Create cuDNN graph and tensors" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Create the cuDNN graph\n", + "graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + ")\n", + "\n", + "# Create tensor handles with the graph API\n", + "x = graph.tensor_like(x_gpu.detach()).set_name(\"X\")\n", + "scale = graph.tensor_like(scale_gpu.detach()).set_name(\"scale\")\n", + "bias = graph.tensor_like(bias_gpu.detach()).set_name(\"bias\")\n", + "epsilon = graph.tensor_like(epsilon_cpu).set_name(\"epsilon\")\n", + "\n", + "(out, mean, inv_var) = graph.layernorm(\n", + " name=\"layernorm\",\n", + " input=x,\n", + " norm_forward_phase=cudnn.norm_forward_phase.INFERENCE, # Note INFERENCE and not TRAINING\n", + " scale=scale,\n", + " bias=bias,\n", + " epsilon=epsilon,\n", + ")\n", + "\n", + "# Enable only the desired output, by default, outputs are disabled\n", + "out.set_name(\"output\").set_output(True).set_data_type(out_expected.dtype)\n", + "\n", + "# Because we have set the norm_forward_phase to INFERENCE, these outputs will be None.\n", + "assert mean is None\n", + "assert inv_var is None" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Build the graph" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Build the graph\n", + "graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", + "\n", + "# To run this block more than once, we need to re-run the previous block to get a new graph.\n", + "# The same instance of a graph should not be built twice." + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Execute the graph" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Mapping of (handles -> memory)\n", + "variant_pack = {\n", + " x: x_gpu.detach(),\n", + " scale: scale_gpu.detach(),\n", + " bias: bias_gpu.detach(),\n", + " epsilon: epsilon_cpu,\n", + " out: out_gpu,\n", + "}\n", + "\n", + "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n", + "graph.execute(variant_pack, workspace)\n", + "torch.cuda.synchronize()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Test cuDNN's output against PyTorch's and check correctness" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# reference output\n", + "torch.testing.assert_close(out_gpu, out_expected, rtol=5e-3, atol=5e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Perform Cleanup" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cudnn.destroy_handle(handle)" + ], + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "language_info": { + "name": "python" + } }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/22_layernorm_inference.ipynb)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites and Setup\n", - "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('nvidia-smi')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If running on Colab, you will need to install the cudnn python interface." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", - "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", - "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Overview" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we are going to perform the layer norm forward pass with the following problem sizes:\n", - "\n", - "- batch size: 4\n", - "- sequence length: 1024\n", - "- embedding dimension: 768\n", - "\n", - "The tensor will be in 16-bit float format." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import cudnn\n", - "import torch\n", - "\n", - "torch.manual_seed(1)\n", - "handle = cudnn.create_handle()\n", - "\n", - "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", - "\n", - "assert torch.cuda.is_available()\n", - "\n", - "batch, seq_size, embedding_dim = 4, 1024, 768\n", - "dtype = torch.float16\n", - "# Epsilon is a small number to prevent division by 0.\n", - "epsilon_value = 1e-3" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Wrapper" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below is how you can use the `Graph` wrapper to perform forward layer norm with the input tensors in PyTorch.\n", - "This is highly similar to the one in the [forward pass notebook](20_layernorm_forward.ipynb)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# allocate random input tensors\n", - "x_gpu = torch.randn(\n", - " batch * seq_size, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype).to(\n", - " memory_format=torch.channels_last\n", - ")\n", - "bias_gpu = torch.randn(1, embedding_dim, 1, 1, device=\"cuda\", dtype=dtype).to(\n", - " memory_format=torch.channels_last\n", - ")\n", - "eps_cpu = torch.full((1, 1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", - "\n", - "# create a graph\n", - "with cudnn.Graph(\n", - " io_data_type=cudnn.data_type.HALF,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - " inputs=[\"ln::input\", \"ln::scale\", \"ln::bias\", \"ln::epsilon\"],\n", - " outputs=[\"ln::Y\"],\n", - ") as graph:\n", - " out, mean, inv_var = graph.layernorm(\n", - " name=\"ln\",\n", - " input=x_gpu,\n", - " norm_forward_phase=cudnn.norm_forward_phase.INFERENCE,\n", - " scale=scale_gpu,\n", - " bias=bias_gpu,\n", - " epsilon=eps_cpu,\n", - " )\n", - " out.set_output(True).set_data_type(dtype)\n", - " assert mean is None, \"Expecting mean to be None under inference mode\"\n", - " assert inv_var is None, \"Expecting inv_var to be None under inference mode\"\n", - "\n", - "# execute the graph and retrieve the output tensors\n", - "out_gpu = graph(x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle)\n", - "\n", - "# verify the result with PyTorch operations\n", - "out_ref = torch.nn.functional.layer_norm(\n", - " x_gpu,\n", - " [embedding_dim, 1, 1],\n", - " weight=scale_gpu.squeeze(0),\n", - " bias=bias_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "\n", - "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Compared to the one in the [forward pass notebook](20_layernorm_forward.ipynb), the argument `norm_forward_phase` is set to `cudnn.norm_forward_phase.INFERENCE` instead of `cudnn.norm_forward_phase.TRAINING`. The result of this is that the mean and variance are not computed nor stored." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Python Binding APIs" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### LayerNorm Inference" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them to calculate a reference value." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# allocate input tensor memory, initialize them to random numbers\n", - "x_gpu = torch.randn(\n", - " batch * seq_size,\n", - " embedding_dim,\n", - " 1,\n", - " 1,\n", - " dtype=dtype,\n", - " requires_grad=True,\n", - " device=\"cuda\",\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, dtype=dtype, requires_grad=True, device=\"cuda\"\n", - ").to(memory_format=torch.channels_last)\n", - "bias_gpu = torch.randn(\n", - " 1, embedding_dim, 1, 1, dtype=dtype, requires_grad=True, device=\"cuda\"\n", - ").to(memory_format=torch.channels_last)\n", - "\n", - "# Epsilon must be a scalar value on the cpu.\n", - "epsilon_cpu = torch.full(\n", - " (1, 1, 1, 1), epsilon_value, dtype=torch.float32, requires_grad=False, device=\"cpu\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create reference computation and allocate output tensor GPU buffers" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create the reference computation outputs here so we can use .empty_like() to create our output buffers\n", - "out_expected = torch.nn.functional.layer_norm(\n", - " x_gpu,\n", - " [embedding_dim, 1, 1],\n", - " weight=scale_gpu.squeeze(0),\n", - " bias=bias_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "\n", - "\n", - "# Allocate output tensor memory using PyTorch\n", - "out_gpu = torch.empty_like(out_expected)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Create cuDNN graph and tensors" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create the cuDNN graph\n", - "graph = cudnn.pygraph(\n", - " handle=handle,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - ")\n", - "\n", - "# Create tensor handles with the graph API\n", - "x = graph.tensor_like(x_gpu.detach()).set_name(\"X\")\n", - "scale = graph.tensor_like(scale_gpu.detach()).set_name(\"scale\")\n", - "bias = graph.tensor_like(bias_gpu.detach()).set_name(\"bias\")\n", - "epsilon = graph.tensor_like(epsilon_cpu).set_name(\"epsilon\")\n", - "\n", - "(out, mean, inv_var) = graph.layernorm(\n", - " name=\"layernorm\",\n", - " input=x,\n", - " norm_forward_phase=cudnn.norm_forward_phase.INFERENCE, # Note INFERENCE and not TRAINING\n", - " scale=scale,\n", - " bias=bias,\n", - " epsilon=epsilon,\n", - ")\n", - "\n", - "# Enable only the desired output, by default, outputs are disabled\n", - "out.set_name(\"output\").set_output(True).set_data_type(out_expected.dtype)\n", - "\n", - "# Because we have set the norm_forward_phase to INFERENCE, these outputs will be None.\n", - "assert mean is None\n", - "assert inv_var is None" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Build the graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Build the graph\n", - "graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", - "\n", - "# To run this block more than once, we need to re-run the previous block to get a new graph.\n", - "# The same instance of a graph should not be built twice." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Execute the graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Mapping of (handles -> memory)\n", - "variant_pack = {\n", - " x: x_gpu.detach(),\n", - " scale: scale_gpu.detach(),\n", - " bias: bias_gpu.detach(),\n", - " epsilon: epsilon_cpu,\n", - " out: out_gpu,\n", - "}\n", - "\n", - "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n", - "graph.execute(variant_pack, workspace)\n", - "torch.cuda.synchronize()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Test cuDNN's output against PyTorch's and check correctness" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# reference output\n", - "torch.testing.assert_close(out_gpu, out_expected, rtol=5e-3, atol=5e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Perform Cleanup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "cudnn.destroy_handle(handle)" - ] - } - ], - "metadata": { - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/samples/python/29_rmsnorm.ipynb b/samples/python/29_rmsnorm.ipynb index 6170ac687..eb2690a7b 100644 --- a/samples/python/29_rmsnorm.ipynb +++ b/samples/python/29_rmsnorm.ipynb @@ -1,520 +1,518 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# RMSNorm Operation\n", - "\n", - "This notebook shows how to compute an RMS norm using the cuDNN python frontend.\n", - "\n", - "$$\\text{RMSNorm}(x) = \\frac{x}{\\sqrt{\\mathbb{E}(x^2) + \\epsilon}}\\cdot\\gamma+\\beta$$" - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# RMSNorm Operation\n", + "\n", + "This notebook shows how to compute an RMS norm using the cuDNN python frontend.\n", + "\n", + "$$\\text{RMSNorm}(x) = \\frac{x}{\\sqrt{\\mathbb{E}(x^2) + \\epsilon}}\\cdot\\gamma+\\beta$$" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/28_rmsnorm.ipynb)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites and Setup\n", + "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('nvidia-smi')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If running on Colab, you will need to install the cudnn python interface." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", + "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", + "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Overview" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the following, we are going to perform the RMS norm forward pass with the following problem sizes:\n", + "\n", + "- batch size: 4\n", + "- sequence length: 1024\n", + "- hidden dimension: 768\n", + "\n", + "The tensor will be in 16-bit float format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import cudnn\n", + "import torch\n", + "\n", + "torch.manual_seed(1)\n", + "\n", + "handle = cudnn.create_handle()\n", + "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", + "\n", + "assert torch.cuda.is_available()\n", + "assert (\n", + " cudnn.backend_version_string() >= \"9.1.0\"\n", + "), \"RMSNorm requires cuDNN 9.1.0 or higher\"\n", + "\n", + "batch, seq_length, hidden_size = 4, 1024, 128\n", + "dtype = torch.float16\n", + "# Epsilon is a small number to prevent division by 0.\n", + "epsilon_value = 1e-3" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Wrapper" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below is how you can use the `Graph` wrapper to perform RMS norm with the input tensors in PyTorch format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# forward pass of rmsnorm using PyTorch in (B, S, H) layout\n", + "x_gpu = torch.randn(\n", + " batch,\n", + " seq_length,\n", + " hidden_size,\n", + " device=\"cuda\",\n", + " dtype=dtype,\n", + " requires_grad=True,\n", + ")\n", + "scale_gpu = torch.randn(\n", + " 1, 1, hidden_size, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "bias_gpu = torch.randn(\n", + " 1, 1, hidden_size, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "eps_cpu = torch.full((1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", + "\n", + "# forward pass of rmsnorm using cuDNN graph\n", + "with cudnn.Graph(\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + " inputs=[\"rms_fwd::input\", \"rms_fwd::scale\", \"rms_fwd::bias\", \"rms_fwd::epsilon\"],\n", + " outputs=[\"rms_fwd::Y\", \"rms_fwd::INV_VARIANCE\"],\n", + ") as rmsnorm_graph:\n", + " out, inv_var = rmsnorm_graph.rmsnorm(\n", + " name=\"rms_fwd\",\n", + " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", + " input=x_gpu,\n", + " scale=scale_gpu,\n", + " bias=bias_gpu,\n", + " epsilon=eps_cpu,\n", + " )\n", + " # set output, inv_var must be float32 tensor\n", + " out.set_output(True).set_data_type(dtype)\n", + " inv_var.set_output(True).set_data_type(cudnn.data_type.FLOAT)\n", + "\n", + "# execute the graph and retrieve the output tensors\n", + "out_gpu, inv_var_gpu = rmsnorm_graph(x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle)\n", + "\n", + "# verify the result with PyTorch API\n", + "out_ref = torch.nn.functional.rms_norm(\n", + " x_gpu,\n", + " [hidden_size],\n", + " weight=scale_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "out_ref = out_ref + bias_gpu\n", + "inv_var_ref = torch.rsqrt(\n", + " torch.mean(x_gpu.float().pow(2), dim=-1, keepdim=True) + epsilon_value\n", + ")\n", + "\n", + "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(inv_var_gpu, inv_var_ref, atol=5e-3, rtol=3e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The above performs forward pass of RMSNorm. Next, we will perform the backward pass:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Compute gradients: Ask PyTorch not to discard the grads after use so that we can read it twice\n", + "# out_ref.grad will be used in the cudnn graph, x_gpu.grad, scale_gpu.grad, and bias_gpu.grad will\n", + "# be used to compare with the cudnn graph output.\n", + "target = torch.randn_like(out_ref)\n", + "criterion = torch.nn.MSELoss()\n", + "loss = criterion(out_ref, target)\n", + "\n", + "out_ref.retain_grad()\n", + "x_gpu.retain_grad()\n", + "scale_gpu.retain_grad()\n", + "bias_gpu.retain_grad()\n", + "\n", + "loss.backward()\n", + "\n", + "# Backward pass\n", + "with cudnn.Graph(\n", + " io_data_type=cudnn.data_type.HALF,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + " inputs=[\n", + " \"rms_bwd::grad\",\n", + " \"rms_bwd::input\",\n", + " \"rms_bwd::scale\",\n", + " \"rms_bwd::inv_variance\",\n", + " ],\n", + " outputs=[\"rms_bwd::DX\", \"rms_bwd::Dscale\", \"rms_bwd::Dbias\"],\n", + ") as bwd_graph:\n", + " dx, dscale, dbias = bwd_graph.rmsnorm_backward(\n", + " name=\"rms_bwd\",\n", + " grad=out_ref,\n", + " input=x_gpu,\n", + " scale=scale_gpu,\n", + " inv_variance=inv_var_gpu,\n", + " has_dbias=True,\n", + " )\n", + " dx.set_output(True).set_data_type(x_gpu.dtype)\n", + " dscale.set_output(True).set_data_type(x_gpu.dtype)\n", + " dbias.set_output(True).set_data_type(x_gpu.dtype)\n", + "\n", + "# execute the graph and retrieve the output tensors\n", + "dx_gpu, dscale_gpu, dbias_gpu = bwd_graph(\n", + " out_ref.grad, x_gpu, scale_gpu, inv_var_gpu, handle=handle\n", + ")\n", + "\n", + "torch.testing.assert_close(x_gpu.grad, dx_gpu, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(scale_gpu.grad, dscale_gpu, atol=5e-3, rtol=3e-3)\n", + "torch.testing.assert_close(bias_gpu.grad, dbias_gpu, atol=5e-3, rtol=3e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Binding APIs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### RMS norm forward pass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them easily when we calculate reference outputs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# input tensor memory in (B, S, H) layout\n", + "x_gpu = torch.randn(\n", + " batch,\n", + " seq_length,\n", + " hidden_size,\n", + " device=\"cuda\",\n", + " dtype=dtype,\n", + " requires_grad=True,\n", + ")\n", + "scale_gpu = torch.randn(\n", + " 1, 1, hidden_size, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "bias_gpu = torch.randn(\n", + " 1, 1, hidden_size, device=\"cuda\", dtype=dtype, requires_grad=True\n", + ")\n", + "# set epsilon to epsilon_value, allocate on cpu.\n", + "eps_cpu = torch.full((1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create cuDNN graph and cuDNN tensors" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + ")\n", + "\n", + "# create tensor handles with the graph API\n", + "x = graph.tensor_like(x_gpu.detach()).set_name(\"X\")\n", + "scale = graph.tensor_like(scale_gpu.detach()).set_name(\"scale\")\n", + "bias = graph.tensor_like(bias_gpu.detach()).set_name(\"bias\")\n", + "epsilon = graph.tensor_like(eps_cpu).set_name(\"epsilon\")\n", + "\n", + "out, inv_var = graph.rmsnorm(\n", + " name=\"rmsnorm\",\n", + " input=x,\n", + " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", + " scale=scale,\n", + " bias=bias,\n", + " epsilon=epsilon,\n", + ")\n", + "\n", + "# enable all outputs\n", + "out.set_name(\"output\").set_output(True).set_data_type(dtype)\n", + "inv_var.set_name(\"inv_var\").set_output(True).set_data_type(cudnn.data_type.FLOAT);" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Build and execute the graph" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Build the graph\n", + "graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", + "\n", + "# Mapping of (handles -> memory)\n", + "variant_pack = {\n", + " x: x_gpu.detach(),\n", + " scale: scale_gpu.detach(),\n", + " bias: bias_gpu.detach(),\n", + " epsilon: eps_cpu,\n", + " out: out_gpu,\n", + " inv_var: inv_var_gpu,\n", + "}\n", + "\n", + "# Execute the graph\n", + "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n", + "graph.execute(variant_pack, workspace)\n", + "torch.cuda.synchronize()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compute reference ouputs and verify" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# reference output using PyTorch\n", + "out_ref = torch.nn.functional.rms_norm(\n", + " x_gpu,\n", + " [hidden_size],\n", + " weight=scale_gpu.reshape(-1),\n", + " eps=epsilon_value,\n", + ")\n", + "out_ref = out_ref + bias_gpu\n", + "inv_var_ref = torch.rsqrt(\n", + " torch.mean(x_gpu.float().pow(2), dim=-1, keepdim=True) + epsilon_value\n", + ")\n", + "\n", + "# Test cuDNN's output against PyTorch's and check correctness\n", + "torch.testing.assert_close(out_gpu, out_ref, rtol=5e-3, atol=5e-3)\n", + "torch.testing.assert_close(inv_var_gpu, inv_var_ref, rtol=5e-3, atol=5e-3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### RMSNorm Backwards Pass" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "target = torch.randn_like(out_ref)\n", + "criterion = torch.nn.MSELoss()\n", + "loss = criterion(out_ref, target)\n", + "\n", + "out_ref.retain_grad()\n", + "x_gpu.retain_grad()\n", + "scale_gpu.retain_grad()\n", + "bias_gpu.retain_grad()\n", + "\n", + "loss.backward()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bwd_graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + ")\n", + "\n", + "d_out = bwd_graph.tensor_like(out_ref.grad)\n", + "x_bwd = bwd_graph.tensor_like(x, name=\"x\")\n", + "scale_bwd = bwd_graph.tensor_like(scale, name=\"scale\")\n", + "inv_var_bwd = bwd_graph.tensor_like(inv_var, name=\"inv_var\")\n", + "\n", + "d_x, d_scale, d_bias = bwd_graph.rmsnorm_backward(\n", + " name=\"d_rmsnorm\",\n", + " grad=d_out,\n", + " input=x_bwd,\n", + " scale=scale_bwd,\n", + " inv_variance=inv_var_bwd,\n", + " has_dbias=True,\n", + ")\n", + "\n", + "d_x.set_output(True).set_data_type(dtype)\n", + "d_scale.set_output(True).set_data_type(dtype)\n", + "d_bias.set_output(True).set_data_type(dtype);" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Build the bwd_graph\n", + "bwd_graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", + "\n", + "# Create output buffers for gradients\n", + "d_x_gpu = torch.empty_like(x_gpu)\n", + "d_scale_gpu = torch.empty_like(scale_gpu)\n", + "d_bias_gpu = torch.empty_like(bias_gpu)\n", + "\n", + "# Execute the graph\n", + "workspace = torch.empty(\n", + " bwd_graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8\n", + ")\n", + "bwd_graph.execute(\n", + " {\n", + " x_bwd: x_gpu.detach(),\n", + " scale_bwd: scale_gpu.detach(),\n", + " d_out: out_ref.grad,\n", + " inv_var_bwd: inv_var_gpu.detach(),\n", + " d_x: d_x_gpu,\n", + " d_scale: d_scale_gpu,\n", + " d_bias: d_bias_gpu,\n", + " },\n", + " workspace,\n", + " handle=handle,\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compare results and check correctness" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "torch.cuda.synchronize()\n", + "\n", + "torch.testing.assert_close(x_gpu.grad, d_x_gpu, atol=2e-4, rtol=2e-4)\n", + "torch.testing.assert_close(scale_gpu.grad, d_scale_gpu, atol=2e-4, rtol=2e-4)\n", + "torch.testing.assert_close(bias_gpu.grad, d_bias_gpu, atol=2e-4, rtol=2e-4)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Perform Cleanup" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cudnn.destroy_handle(handle)" + ], + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "language_info": { + "name": "python" + } }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/28_rmsnorm.ipynb)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites and Setup\n", - "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('nvidia-smi')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If running on Colab, you will need to install the cudnn python interface." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# get_ipython().system('pip install nvidia-cudnn-cu12')\n", - "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", - "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Overview" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we are going to perform the RMS norm forward pass with the following problem sizes:\n", - "\n", - "- batch size: 4\n", - "- sequence length: 1024\n", - "- hidden dimension: 768\n", - "\n", - "The tensor will be in 16-bit float format." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import cudnn\n", - "import torch\n", - "\n", - "torch.manual_seed(1)\n", - "\n", - "handle = cudnn.create_handle()\n", - "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n", - "\n", - "assert torch.cuda.is_available()\n", - "assert (\n", - " cudnn.backend_version_string() >= \"9.1.0\"\n", - "), \"RMSNorm requires cuDNN 9.1.0 or higher\"\n", - "\n", - "batch, seq_length, hidden_size = 4, 1024, 128\n", - "dtype = torch.float16\n", - "# Epsilon is a small number to prevent division by 0.\n", - "epsilon_value = 1e-3" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Wrapper" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below is how you can use the `Graph` wrapper to perform RMS norm with the input tensors in PyTorch format." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# forward pass of rmsnorm using PyTorch\n", - "x_gpu = torch.randn(\n", - " batch * seq_length,\n", - " hidden_size,\n", - " 1,\n", - " 1,\n", - " device=\"cuda\",\n", - " dtype=dtype,\n", - " requires_grad=True,\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(\n", - " 1, hidden_size, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "bias_gpu = torch.randn(\n", - " 1, hidden_size, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "eps_cpu = torch.full((1, 1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")\n", - "\n", - "# forward pass of rmsnorm using cuDNN graph\n", - "with cudnn.Graph(\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - " inputs=[\"rms_fwd::input\", \"rms_fwd::scale\", \"rms_fwd::bias\", \"rms_fwd::epsilon\"],\n", - " outputs=[\"rms_fwd::Y\", \"rms_fwd::INV_VARIANCE\"],\n", - ") as rmsnorm_graph:\n", - " out, inv_var = rmsnorm_graph.rmsnorm(\n", - " name=\"rms_fwd\",\n", - " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", - " input=x_gpu,\n", - " scale=scale_gpu,\n", - " bias=bias_gpu,\n", - " epsilon=eps_cpu,\n", - " )\n", - " # set output, inv_var must be float32 tensor\n", - " out.set_output(True).set_data_type(dtype)\n", - " inv_var.set_output(True).set_data_type(cudnn.data_type.FLOAT)\n", - "\n", - "# execute the graph and retrieve the output tensors\n", - "out_gpu, inv_var_gpu = rmsnorm_graph(x_gpu, scale_gpu, bias_gpu, eps_cpu, handle=handle)\n", - "\n", - "# verify the result with PyTorch API\n", - "out_ref = torch.nn.functional.rms_norm(\n", - " x_gpu,\n", - " [hidden_size, 1, 1], # RMS norm over last 3 dimensions\n", - " weight=scale_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "out_ref = out_ref + bias_gpu\n", - "inv_var_ref = torch.rsqrt(\n", - " torch.mean(x_gpu.float().pow(2), dim=(1, 2, 3), keepdim=True) + epsilon_value\n", - ")\n", - "\n", - "torch.testing.assert_close(out_gpu, out_ref, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(inv_var_gpu, inv_var_ref, atol=5e-3, rtol=3e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The above performs forward pass of RMSNorm. Next, we will perform the backward pass:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Compute gradients: Ask PyTorch not to discard the grads after use so that we can read it twice\n", - "# out_ref.grad will be used in the cudnn graph, x_gpu.grad, scale_gpu.grad, and bias_gpu.grad will\n", - "# be used to compare with the cudnn graph output.\n", - "target = torch.randn_like(out_ref)\n", - "criterion = torch.nn.MSELoss()\n", - "loss = criterion(out_ref, target)\n", - "\n", - "out_ref.retain_grad()\n", - "x_gpu.retain_grad()\n", - "scale_gpu.retain_grad()\n", - "bias_gpu.retain_grad()\n", - "\n", - "loss.backward()\n", - "\n", - "# Backward pass\n", - "with cudnn.Graph(\n", - " io_data_type=cudnn.data_type.HALF,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - " inputs=[\n", - " \"rms_bwd::grad\",\n", - " \"rms_bwd::input\",\n", - " \"rms_bwd::scale\",\n", - " \"rms_bwd::inv_variance\",\n", - " ],\n", - " outputs=[\"rms_bwd::DX\", \"rms_bwd::Dscale\", \"rms_bwd::Dbias\"],\n", - ") as bwd_graph:\n", - " dx, dscale, dbias = bwd_graph.rmsnorm_backward(\n", - " name=\"rms_bwd\",\n", - " grad=out_ref,\n", - " input=x_gpu,\n", - " scale=scale_gpu,\n", - " inv_variance=inv_var_gpu,\n", - " has_dbias=True,\n", - " )\n", - " dx.set_output(True).set_data_type(x_gpu.dtype)\n", - " dscale.set_output(True).set_data_type(x_gpu.dtype)\n", - " dbias.set_output(True).set_data_type(x_gpu.dtype)\n", - "\n", - "# execute the graph and retrieve the output tensors\n", - "dx_gpu, dscale_gpu, dbias_gpu = bwd_graph(\n", - " out_ref.grad, x_gpu, scale_gpu, inv_var_gpu, handle=handle\n", - ")\n", - "\n", - "torch.testing.assert_close(x_gpu.grad, dx_gpu, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(scale_gpu.grad, dscale_gpu, atol=5e-3, rtol=3e-3)\n", - "torch.testing.assert_close(bias_gpu.grad, dbias_gpu, atol=5e-3, rtol=3e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Using Python Binding APIs" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### RMS norm forward pass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them easily when we calculate reference outputs." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# input tensor memory, initialize them to random numbers\n", - "x_gpu = torch.randn(\n", - " batch * seq_length,\n", - " hidden_size,\n", - " 1,\n", - " 1,\n", - " device=\"cuda\",\n", - " dtype=dtype,\n", - " requires_grad=True,\n", - ").to(memory_format=torch.channels_last)\n", - "scale_gpu = torch.randn(\n", - " 1, hidden_size, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "bias_gpu = torch.randn(\n", - " 1, hidden_size, 1, 1, device=\"cuda\", dtype=dtype, requires_grad=True\n", - ").to(memory_format=torch.channels_last)\n", - "# set epsilon to epsilon_value, allocate on cpu.\n", - "eps_cpu = torch.full((1, 1, 1, 1), epsilon_value, dtype=torch.float32, device=\"cpu\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create cuDNN graph and cuDNN tensors" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph = cudnn.pygraph(\n", - " handle=handle,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - ")\n", - "\n", - "# create tensor handles with the graph API\n", - "x = graph.tensor_like(x_gpu.detach()).set_name(\"X\")\n", - "scale = graph.tensor_like(scale_gpu.detach()).set_name(\"scale\")\n", - "bias = graph.tensor_like(bias_gpu.detach()).set_name(\"bias\")\n", - "epsilon = graph.tensor_like(eps_cpu).set_name(\"epsilon\")\n", - "\n", - "out, inv_var = graph.rmsnorm(\n", - " name=\"rmsnorm\",\n", - " input=x,\n", - " norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n", - " scale=scale,\n", - " bias=bias,\n", - " epsilon=epsilon,\n", - ")\n", - "\n", - "# enable all outputs\n", - "out.set_name(\"output\").set_output(True).set_data_type(dtype)\n", - "inv_var.set_name(\"inv_var\").set_output(True).set_data_type(cudnn.data_type.FLOAT);" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Build and execute the graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Build the graph\n", - "graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", - "\n", - "# Mapping of (handles -> memory)\n", - "variant_pack = {\n", - " x: x_gpu.detach(),\n", - " scale: scale_gpu.detach(),\n", - " bias: bias_gpu.detach(),\n", - " epsilon: eps_cpu,\n", - " out: out_gpu,\n", - " inv_var: inv_var_gpu,\n", - "}\n", - "\n", - "# Execute the graph\n", - "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n", - "graph.execute(variant_pack, workspace)\n", - "torch.cuda.synchronize()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Compute reference ouputs and verify" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# reference output using PyTorch\n", - "out_ref = torch.nn.functional.rms_norm(\n", - " x_gpu,\n", - " [hidden_size, 1, 1], # RMS norm over last 3 dimensions\n", - " weight=scale_gpu.squeeze(0),\n", - " eps=epsilon_value,\n", - ")\n", - "out_ref = out_ref + bias_gpu\n", - "inv_var_ref = torch.rsqrt(\n", - " torch.mean(x_gpu.float().pow(2), dim=(1, 2, 3), keepdim=True) + epsilon_value\n", - ")\n", - "\n", - "# Test cuDNN's output against PyTorch's and check correctness\n", - "torch.testing.assert_close(out_gpu, out_ref, rtol=5e-3, atol=5e-3)\n", - "torch.testing.assert_close(inv_var_gpu, inv_var_ref, rtol=5e-3, atol=5e-3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### RMSNorm Backwards Pass" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "target = torch.randn_like(out_ref)\n", - "criterion = torch.nn.MSELoss()\n", - "loss = criterion(out_ref, target)\n", - "\n", - "out_ref.retain_grad()\n", - "x_gpu.retain_grad()\n", - "scale_gpu.retain_grad()\n", - "bias_gpu.retain_grad()\n", - "\n", - "loss.backward()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "bwd_graph = cudnn.pygraph(\n", - " handle=handle,\n", - " intermediate_data_type=cudnn.data_type.FLOAT,\n", - " compute_data_type=cudnn.data_type.FLOAT,\n", - ")\n", - "\n", - "d_out = bwd_graph.tensor_like(out_ref.grad)\n", - "x_bwd = bwd_graph.tensor_like(x, name=\"x\")\n", - "scale_bwd = bwd_graph.tensor_like(scale, name=\"scale\")\n", - "inv_var_bwd = bwd_graph.tensor_like(inv_var, name=\"inv_var\")\n", - "\n", - "d_x, d_scale, d_bias = bwd_graph.rmsnorm_backward(\n", - " name=\"d_rmsnorm\",\n", - " grad=d_out,\n", - " input=x_bwd,\n", - " scale=scale_bwd,\n", - " inv_variance=inv_var_bwd,\n", - " has_dbias=True,\n", - ")\n", - "\n", - "d_x.set_output(True).set_data_type(dtype)\n", - "d_scale.set_output(True).set_data_type(dtype)\n", - "d_bias.set_output(True).set_data_type(dtype);" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Build the bwd_graph\n", - "bwd_graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", - "\n", - "# Create output buffers for gradients\n", - "d_x_gpu = torch.empty_like(x_gpu)\n", - "d_scale_gpu = torch.empty_like(scale_gpu)\n", - "d_bias_gpu = torch.empty_like(bias_gpu)\n", - "\n", - "# Execute the graph\n", - "workspace = torch.empty(\n", - " bwd_graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8\n", - ")\n", - "bwd_graph.execute(\n", - " {\n", - " x_bwd: x_gpu.detach(),\n", - " scale_bwd: scale_gpu.detach(),\n", - " d_out: out_ref.grad,\n", - " inv_var_bwd: inv_var_gpu.detach(),\n", - " d_x: d_x_gpu,\n", - " d_scale: d_scale_gpu,\n", - " d_bias: d_bias_gpu,\n", - " },\n", - " workspace,\n", - " handle=handle,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Compare results and check correctness" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "torch.cuda.synchronize()\n", - "\n", - "torch.testing.assert_close(x_gpu.grad, d_x_gpu, atol=2e-4, rtol=2e-4)\n", - "torch.testing.assert_close(scale_gpu.grad, d_scale_gpu, atol=2e-4, rtol=2e-4)\n", - "torch.testing.assert_close(bias_gpu.grad, d_bias_gpu, atol=2e-4, rtol=2e-4)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Perform Cleanup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "cudnn.destroy_handle(handle)" - ] - } - ], - "metadata": { - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/test/python/test_layernorm.py b/test/python/test_layernorm.py index f5ee920cc..be3809ae8 100644 --- a/test/python/test_layernorm.py +++ b/test/python/test_layernorm.py @@ -33,15 +33,14 @@ def test_layernorm(param_extract, cudnn_handle): atol, rtol = 1e-2, 1e-2 batch_size, seq_size = 16, 128 - N, C, H, W = batch_size * seq_size, embedding_dim, 1, 1 epsilon_value = 1e-3 - x_gpu = 3 * torch.randn(N, C, H, W, requires_grad=True, device="cuda", dtype=input_type).to(memory_format=torch.channels_last) - 0.5 - scale_gpu = 5 * torch.randn(1, C, H, W, requires_grad=True, device="cuda", dtype=input_type).to(memory_format=torch.channels_last) - 1 - bias_gpu = 7 * torch.randn(1, C, H, W, requires_grad=True, device="cuda", dtype=input_type).to(memory_format=torch.channels_last) - 2 + x_gpu = 3 * torch.randn(batch_size, seq_size, embedding_dim, requires_grad=True, device="cuda", dtype=input_type) - 0.5 + scale_gpu = 5 * torch.randn(1, 1, embedding_dim, requires_grad=True, device="cuda", dtype=input_type) - 1 + bias_gpu = 7 * torch.randn(1, 1, embedding_dim, requires_grad=True, device="cuda", dtype=input_type) - 2 epsilon_cpu = torch.full( - (1, 1, 1, 1), + (1, 1, 1), epsilon_value, requires_grad=False, device="cpu", @@ -50,13 +49,13 @@ def test_layernorm(param_extract, cudnn_handle): Y_expected = torch.nn.functional.layer_norm( x_gpu, - [C, H, W], - weight=scale_gpu.squeeze(0), - bias=bias_gpu.squeeze(0), + [embedding_dim], + weight=scale_gpu.reshape(-1), + bias=bias_gpu.reshape(-1), eps=epsilon_value, ) - mean_expected = x_gpu.to(torch.float32).mean(dim=(1, 2, 3), keepdim=True) - inv_var_expected = torch.rsqrt(torch.var(x_gpu.to(torch.float32), dim=(1, 2, 3), keepdim=True) + epsilon_value) + mean_expected = x_gpu.to(torch.float32).mean(dim=-1, keepdim=True) + inv_var_expected = torch.rsqrt(torch.var(x_gpu.to(torch.float32), dim=-1, keepdim=True) + epsilon_value) stream = torch.cuda.current_stream().cuda_stream cudnn.set_stream(handle=cudnn_handle, stream=stream) diff --git a/test/python/test_rmsnorm.py b/test/python/test_rmsnorm.py index 43d1997f7..a2f32b44c 100644 --- a/test/python/test_rmsnorm.py +++ b/test/python/test_rmsnorm.py @@ -55,15 +55,14 @@ def test_rmsnorm(param_extract, cudnn_handle): embedding_dim, input_type, has_bias = param_extract batch_size, seq_size = 16, 128 - N, C, H, W = batch_size * seq_size, embedding_dim, 1, 1 epsilon_value = 1e-3 - x_gpu = 2 * torch.randn(N, C, H, W, requires_grad=True, device="cuda", dtype=input_type) - 1.25 - scale_gpu = 3 * torch.randn(1, C, H, W, requires_grad=True, device="cuda", dtype=input_type) - 2.75 - bias_gpu = torch.randn(1, C, H, W, requires_grad=True, device="cuda", dtype=input_type) + x_gpu = 2 * torch.randn(batch_size, seq_size, embedding_dim, requires_grad=True, device="cuda", dtype=input_type) - 1.25 + scale_gpu = 3 * torch.randn(1, 1, embedding_dim, requires_grad=True, device="cuda", dtype=input_type) - 2.75 + bias_gpu = torch.randn(1, 1, embedding_dim, requires_grad=True, device="cuda", dtype=input_type) epsilon_cpu = torch.full( - (1, 1, 1, 1), + (1, 1, 1), epsilon_value, requires_grad=False, device="cpu", @@ -72,7 +71,7 @@ def test_rmsnorm(param_extract, cudnn_handle): print("Running reference") - model = RMSNorm(eps=epsilon_value, dim=(1, 2, 3)).float() + model = RMSNorm(eps=epsilon_value, dim=-1).float() Y_expected, inv_var_expected = model(x_gpu, scale_gpu, bias_gpu if has_bias else None) print("Building cudnn graph")