---
title: "Profiler-Driven Kernel Optimization for Fine-Tuning"
publisher: "nvidia"
type: "playbook"
updated: "2026-05-26T18:38:42.972Z"
description: "Use torch.profiler to find training bottlenecks, then write custom Triton kernels to optimize LLaMA 8B fine-tuning"
canonical: "https://build.nvidia.com/station/kernel-dev-ft.md"
---

# Basic idea

DGX Station puts a full Blackwell GPU on your desk, which makes it an ideal environment for profiling and optimizing GPU kernels used during model training. This playbook walks through a real optimization workflow: **profiling a LLaMA 3.1 8B fine-tuning run to identify bottlenecks, then writing custom Triton kernels that eliminate those bottlenecks** — specifically a fused RMSNorm and a fused cross-entropy loss using online softmax.

For inference workloads, tools like `torch.compile` and serving frameworks (vLLM, TensorRT-LLM) already ship highly optimized fused kernels. But training workloads are different. Backward passes double the kernel count, large vocabularies create massive intermediate tensors during loss computation, and `torch.compile` does not restructure algorithms to avoid these allocations. Projects like [Liger-Kernel](https://github.com/linkedin/Liger-Kernel) and [Unsloth](https://github.com/unslothai/unsloth) demonstrate that custom training kernels deliver real results: 20-60% memory reduction and 10-30% throughput improvement.

This playbook uses **Triton** instead of raw CUDA C++. Triton is a Python-native GPU programming language that JIT-compiles to optimized GPU code — no `nvcc` compiler, no C++ build systems, no manual thread indexing. It is the standard for custom training kernels: Liger-Kernel, Unsloth, and FlashAttention are all written in Triton.

**No prior Triton, CUDA, or GPU programming experience is required.** The instructions explain each concept as it comes up.

# What you'll accomplish

You will profile a LLaMA 3.1 8B fine-tuning workload, identify the key performance bottlenecks, and write custom Triton kernels that address them.

- **Profile** a baseline fine-tuning step using `torch.profiler` and interpret the results to identify two targets: RMSNorm (memory-bandwidth-bound) and cross-entropy loss (memory-capacity-bound).
- **Write a fused RMSNorm kernel** in Triton that processes normalization in a single GPU pass instead of multiple separate operations, improving memory bandwidth utilization from ~11% to ~80-90% of peak.
- **Write a fused cross-entropy kernel** using the online softmax algorithm (Milakov-Gimelshein) that computes loss without materializing intermediate softmax tensors, achieving ~6x memory reduction and up to 4x latency improvement at realistic batch sizes.
- **Verify correctness** of both kernels (forward and backward passes) against PyTorch reference implementations.
- **Benchmark** the kernels to measure latency, throughput, and memory savings.
- **Integrate** both kernels into an end-to-end LLaMA 3.1 8B fine-tuning loop and measure real training throughput and memory improvements.

# What to know before starting

- Comfortable with Linux command line and shell scripting.
- Basic familiarity with Python and PyTorch (tensors, autograd, training loops).
- Understanding of what fine-tuning is (training a pre-trained model on new data).
- No Triton, CUDA, or GPU programming experience required — all code is explained.

# Prerequisites

**Hardware:**
- NVIDIA DGX Station with GB300 Ultra Superchip.
- At least 150 GB available storage for the container image, model weights (~16 GB for LLaMA 3.1 8B in BF16), profiler traces, and optimizer states.

**Software:**
- Docker with NVIDIA Container Toolkit: `docker run --rm --gpus all nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu24.04 nvidia-smi`
- On a DGX Station, immediately confirm which device index belongs to the GB300 so later steps can target it explicitly. Run `nvidia-smi --query-gpu=index,name --format=csv,noheader` and note the index for the row showing `NVIDIA GB300`. Subsequent steps recommend `--gpus '"device=N"'` (with `N` = that index) instead of `--gpus all` so profiling and benchmark numbers stay on a single, known GPU.
- Network access to pull container images from NGC and download model weights from Hugging Face.
- A Hugging Face account with access to [meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B) and a [Hugging Face access token](https://huggingface.co/settings/tokens).

# Ancillary files

All required assets are in the playbook directory `nvidia/station-kernel-dev-ft/assets` (see the [dgx-spark-playbooks](https://github.com/NVIDIA/dgx-spark-playbooks) repository).

- `assets/Dockerfile` — Development container based on NVIDIA's PyTorch NGC image with Triton, transformers, and profiling dependencies.
- `assets/requirements.txt` — Python dependencies installed inside the container.
- `assets/profile_baseline.py` — Profiling script that captures a `torch.profiler` trace of a LLaMA 3.1 8B training step and prints a breakdown of GPU time by operation. Supports flags to enable custom kernels for re-profiling.
- `assets/rmsnorm_kernel.py` — Fused RMSNorm Triton kernel with forward and backward passes, wrapped as a drop-in `torch.nn.Module` replacement. Heavily commented with explanations of each Triton concept.
- `assets/rmsnorm_test.py` — Correctness tests comparing the custom RMSNorm against PyTorch's reference implementation (forward and backward, FP32 and BF16).
- `assets/cross_entropy_kernel.py` — Fused cross-entropy Triton kernel using online softmax, with forward and backward passes. Processes the vocabulary in chunks to avoid materializing the full logit tensor.
- `assets/cross_entropy_test.py` — Correctness tests and memory usage comparison against `torch.nn.CrossEntropyLoss`.
- `assets/benchmark_kernels.py` — Benchmarking script that measures latency, throughput, bandwidth utilization, and peak memory for both custom kernels.
- `assets/finetune_baseline.py` — Minimal LLaMA 3.1 8B fine-tuning script using vanilla PyTorch, reporting tokens/sec and peak memory.
- `assets/finetune_optimized.py` — Identical fine-tuning script with both custom kernels monkey-patched in for direct comparison.

# Time & risk

* **Estimated time:** About 2 hours. Steps 1-4 (setup through baseline profiling) take about 30 minutes. Steps 5-7 (RMSNorm kernel) take about 30 minutes. Steps 8-10 (cross-entropy kernel) take about 40 minutes. Step 11 (end-to-end integration) takes about 20 minutes. Steps 12-13 (cleanup and next steps) are a few minutes.
* **Risk level:** Low
* All work runs inside a Docker container — no host system modifications.
* LLaMA 3.1 8B model weights (~16 GB in BF16) are downloaded from Hugging Face on first run and cached locally.
* Requires a Hugging Face token with access to the LLaMA 3.1 model.
* **Rollback:** Exit the container. Your source files are preserved in the mounted `assets/` directory; everything else is discarded.
* **Last Updated:** 05/26/2026
* First publication

## More

- [Instructions](/station/kernel-dev-ft/instructions.md)
- [Troubleshooting](/station/kernel-dev-ft/troubleshooting.md)