Lazy loaded image8.1 nanoGPT: Training a GPT from Scratch

8.1 nanoGPT: Training a GPT from Scratch

Full code provided in the end for downloading

8.1.1 Preparations

  • Prepare the Environment: We provide two methods for setting up your environment: using Anaconda or Google Colab.
      1. Using Anaconda: Create a new Conda environment and install the required packages. (If you're not familiar with managing Anaconda environments, it's recommended to learn that first.)
        1. Ensure your environment has CUDA version 11.8 or higher
          Run nvidia-smi in the command prompt. The output will include your CUDA version.
          notion image
          • Create the environment:
            • Activate the environment:
              • Install the necessary libraries:
            1. Using CoLab: If you prefer to run the experiment online, click this link: nanoGPT example.
                • You don’t need to manually download the data when using Colab. The first code block fetches the dataset online and loads it into memory:
                  • By default, Colab uses the CPU. To switch to a GPU:
                    • Click Runtime → Change runtime type:
                      • notion image
                    • Select T4 GPU. (Note that TPU v2-8 work only for TensorFlow, but we are using PyTorch.)
                      • notion image
          • Prepare the Data: Download tinyshakespeare dataset used by nanoGPT. You can use one of the following methods:
            • On Linux, use wget:
              • On Windows (Windows 10 or 11), use curl:
                • Or download via your browser. Click the dataset link: tinyshakespeare, then press Ctrl + S to save the file manually. Or, you can copy the text and save it to a .txt file.
              <ins/>

              8.1.2 Code Explanation

              • Import related packages and set hyperparameters
                • Read the dataset, doing a simple tokenization
                  • The output is
                • Dataset splitting and batching
                  • Model Evaluation
                    • Model Definition
                      • A single head in the multi-head self-attention mechanism
                        • Multi-head self-attention, FFN, and block definition:
                          • Define a Bigram Language Model
                          • Model Training
                            <ins/>

                            8.1.3 Experimental Results

                            • Experimental Results: First, let's examine the model parameters. From the implementation above, we can see that this small GPT-2 model has approximately 0.209729 million parameters. Using the rough estimation method introduced earlier:
                              • This estimate is roughly consistent with the actual value.
                            • Generate new text from a blank context:
                              • Training loss
                                • notion image
                              • GPU Memory Usage During Training
                                • Training requires only 348MB of GPU memory. The code uses FP32 precision, meaning each element occupies 4 bytes.
                                  • notion image
                                • When CUDA is initialized, it consumes 298MB for context initialization. You can observe this by running torch.tensor(1).to('cuda'), which triggers CUDA initialization in PyTorch.
                                  • notion image
                                • The remaining memory is used for model parameters, gradients, optimizer state, intermediate activations, and the PyTorch memory cache.
                                  • To understand how much memory is consumed by intermediate activations, it's helpful to print torch.cuda.memory_allocated() after each layer during a forward pass.
                                    Monitor GPU usage.
                                    Instead of using nvidia-smi for a one-time snapshot of GPU usage, you can use these Python-based tools:
                                    1. gpustat: Install with pip install gpustat. Run it in an activated environment to get a detailed overview of GPU usage.
                                    1. nvitop: Install with pip install nvitop. Run it to get a dynamic, real-time view of GPU processes and memory usage.
                                    ⚠️ On Windows, neither gpustat nor nvitop can provide detailed per-process GPU memory usage.
                                Prev
                                7.3 Knowledge Distillation
                                Next
                                8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example
                                Loading...
                                Article List
                                LLM Learning Roadmap
                                ✨ Awesome-Anything
                                🖼️ Digital Image Processing
                                🍃 LLM Components
                                🌱 LLM Pre-training
                                ☘️ LLM Post-Training
                                🍀 LLM Popular Models
                                🪴 LLM Applications
                                🌿 LLM Optimization
                                🌾 LLM Compression
                                🌵 LLM Hands-on Practice
                                🌴 LLM Must-read Papers
                                🌳 LLM Q&A
                                🐝 VLM Image Encoders
                                📝 MISC.