Gemma语言模型

Setup

1. Clone the repository:

    git clone https://github.com/lmassaron/gemma_from_scratch.git
    cd gemma_from_scratch
    

2. Create a virtual environment and install dependencies:

    python -m venv venv
    source venv/bin/activate  # On Windows, use `venv\Scripts\activate`
    pip install -r requirements.txt
    

How to Use

The process is divided into two main steps: data preparation and model training.

Usage

This repository provides two main scripts for running inference:

• inference_google_gemma.py: This script uses the pre-trained Gemma 3 270M model from the Hugging Face Hub. It downloads the weights, loads them into the model, and generates text from a list of predefined sentences. This is the recommended script for most users.

    python inference_google_gemma.py
    

• inference_custom.py: This script is for running inference with a custom model. It requires a .pth file with the model weights. You can use this script to test your own trained models.

    python inference_custom.py --model-path /path/to/your/model.pth
    

Both scripts will output the generated text to the console.

References

• Original Notebook: LLMs-from-scratch/ch05/12_gemma3/standalone-gemma3.ipynb by Sebastian Raschka.
• Hugging Face Model Card: google/gemma-3-270m

Key Components & Logic

• prepare_dataset.py: A flexible data processing script. It parallelizes the tokenization step across all available CPU cores for maximum efficiency and uses memory-mapped NumPy arrays to handle datasets larger than RAM.

• train.py: The main training loop. It implements modern training best practices:
• Mixed-Precision Training: Uses torch.amp.autocast with bfloat16 for faster training and reduced memory usage on supported hardware.
• Optimizer: Employs the AdamW optimizer, which adds weight decay for better regularization.
• Learning Rate Scheduler: Uses a SequentialLR scheduler that combines a linear warmup phase with a cosine decay, helping to stabilize training.
• Gradient Accumulation: Allows for training with large effective batch sizes even on memory-constrained GPUs.
• Gradient Clipping: Prevents exploding gradients by clipping the norm of the gradients before the optimizer step.

• gemma_scratch/ (The Core Package):
• model.py: Defines the Gemma3Model class, a PyTorch nn.Module that assembles the complete transformer architecture.
• layers.py: Contains the TransformerBlock, the core repeating unit of the model, which includes multi-head attention and the MLP layers.
• rope.py: Implements Rotary Positional Embeddings (RoPE), a modern technique for injecting positional information into the self-attention mechanism.
• normalization.py: Provides an efficient RMSNorm (Root Mean Square Normalization) layer, which is used throughout the Gemma architecture instead of traditional LayerNorm.
• config.py: A simple file to store the model's hyperparameters (e.g., number of layers, heads, embedding dimensions).
• tokenizer.py: A wrapper for the GPT-2 tokenizer used for encoding the text data.

Two Workflows: Training vs. Inference

This repository supports two primary use cases: 1. Inference with Official Gemma Weights: The inference_google_gemma.py script uses the official Gemma tokenizer and loads the pre-trained 270M model from Hugging Face. This demonstrates the architectural compatibility of our implementation. 2. Training a Model From Scratch: The prepare_dataset.py and train.py scripts allow you to train a model on your own data (e.g., TinyStories). For this workflow, we use the simpler and faster GPT-2 tokenizer (tiktoken). This is a self-contained training pipeline, and the resulting model should be used with inference_custom.py, which also uses the GPT-2 tokenizer.