🧠 Complete LLM Pretraining Pipeline

Build a ~100M parameter general-text language model from scratch on a single GPU

100M Parameters Single GPU General Text Causal LM (GPT-style) Hugging Face Ecosystem Open Datasets

📋 Table of Contents

1. Project Overview & Architecture 2. Environment Setup 3. Data Sourcing & Preparation 4. Train Custom Tokenizer 5. Model Configuration 6. Training Script 7. Memory Optimizations 8. Evaluation & Benchmarking 9. Inference & Text Generation 10. Optional: SFT/DPO Alignment

1. Project Overview & Architecture

Goal: Train a small but capable causal language model from scratch using only publicly available data and a single GPU (8-24GB VRAM).

Model Architecture (100M params)

Component	Setting	Notes
Architecture	GPT-Neo style (decoder-only)	Proven, efficient, well-supported
Hidden Size	768	Embedding dimension
Layers	12	Transformer blocks
Attention Heads	12	64 dim per head
FFN Dim	3072	4× hidden size
Vocab Size	32,000	Trained BPE tokenizer
Context Length	512	Trainable on single GPU
Parameters	~100M	Perfect for experimentation

Recommended GPU Specs

💻 Minimum: RTX 3060 12GB

Batch size 4-8, gradient accumulation

🎯 Recommended: RTX 4090 24GB

Batch size 16-32, faster training

☁️ Cloud: A100 40GB / L4 24GB

Maximum throughput, bf16 support

Training Time Estimate: On RTX 4090, ~2-4 hours for 1 epoch over 10B tokens. Full convergence may take 3-10 epochs depending on data quality.

2. Environment Setup

1 Create project structure:

llm-project/ ├── data/ # Raw and processed datasets ├── tokenizer/ # Trained tokenizer files ├── models/ # Saved model checkpoints ├── scripts/ # Training scripts │ ├── train_tokenizer.py │ ├── pretrain.py │ ├── evaluate.py │ └── generate.py ├── requirements.txt └── README.md

2 Install dependencies:

bash

# Create virtual environment
python -m venv venv
source venv/bin/activate  # Linux/Mac
# venv\Scripts\activate   # Windows

# Install core dependencies
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers datasets tokenizers accelerate wandb
pip install huggingface_hub sentencepiece protobuf

# Optional: for optimization
pip install flash-attn --no-build-isolation  # Requires CUDA 11.6+
pip install bitsandbytes deepspeed

# For evaluation
pip install lm-eval

# For alignment (SFT/DPO)
pip install trl peft

txt # requirements.txt (save this)

torch>=2.1.0
transformers>=4.40.0
datasets>=2.18.0
tokenizers>=0.19.0
accelerate>=0.29.0
huggingface-hub>=0.22.0
wandb>=0.16.0
sentencepiece>=0.2.0
protobuf>=4.25.0
flash-attn>=2.5.0
bitsandbytes>=0.43.0
deepspeed>=0.14.0
trl>=0.8.0
peft>=0.10.0
lm-eval>=0.4.0

Login to Hugging Face: You'll need a token to upload models and access some datasets.

huggingface-cli login
# Or in Python:
from huggingface_hub import login
login(token="your_hf_token")

3. Data Sourcing & Preparation

We'll use a mix of high-quality open datasets that together provide ~15-20B tokens — perfect for training a 100M model.

Recommended Dataset Mix

Dataset	Size	Type	HF Hub ID
Fineweb-Edu	~1.3B tokens	Educational web text	`HuggingFaceFW/fineweb-edu`
The Pile (subset)	~2B tokens	Diverse academic/text	`EleutherAI/pile`
Project Gutenberg (PG19)	~3B tokens	Classic literature	`pg19`
OpenWebText	~8B tokens	Web text (Reddit links)	`Skylion007/openwebtext`
C4 (en)	~15B tokens	Cleaned Common Crawl	`allenai/c4`

3 Create data preparation script:

python

# scripts/prepare_data.py
import os
from datasets import load_dataset, concatenate_datasets, DatasetDict
from transformers import AutoTokenizer
import random

SEED = 42
random.seed(SEED)

def load_and_mix_datasets(target_tokens=5_000_000_000, streaming=True):
    """
    Load and mix datasets for pretraining.
    Target: ~5B tokens (good starting point for 100M model)
    """
    
    datasets_config = [
        {"name": "HuggingFaceFW/fineweb-edu", "split": "train", "weight": 0.30, "text_col": "text"},
        {"name": "Skylion007/openwebtext", "split": "train", "weight": 0.25, "text_col": "text"},
        {"name": "allenai/c4", "split": "train", "weight": 0.25, "text_col": "text"},
        {"name": "bookcorpus", "split": "train", "weight": 0.20, "text_col": "text"},
    ]
    
    all_datasets = []
    
    for config in datasets_config:
        print(f"Loading {config['name']}...")
        try:
            ds = load_dataset(
                config["name"], 
                split=config["split"], 
                streaming=streaming,
                trust_remote_code=True
            )
            # Take weighted subset
            ds = ds.shuffle(seed=SEED)
            all_datasets.append((ds, config["text_col"], config["weight"]))
        except Exception as e:
            print(f"Warning: Could not load {config['name']}: {e}")
    
    return all_datasets

def tokenize_function(examples, tokenizer, text_column="text"):
    """Tokenize texts with truncation and padding."""
    return tokenizer(
        examples[text_column],
        truncation=True,
        max_length=512,
        return_overflowing_tokens=False,
        return_length=True,
    )

def prepare_pretraining_data(
    tokenizer_name_or_path="gpt2",
    output_dir="./data/pretrain",
    context_length=512,
    num_proc=4
):
    """
    Prepare concatenated pretraining data.
    Groups texts into context-length chunks for efficient training.
    """
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Load tokenizer (we'll train our own later, use GPT-2 for now)
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    print("Loading datasets...")
    # For initial prep, use a single reliable dataset
    ds = load_dataset("openwebtext", streaming=True, split="train")
    ds = ds.take(5_000_000)  # ~2.5B tokens worth
    
    # Convert streaming to in-memory for processing
    print("Converting to in-memory dataset...")
    texts = []
    for i, example in enumerate(ds):
        texts.append(example["text"])
        if i % 100000 == 0:
            print(f"  Processed {i} examples...")
    
    from datasets import Dataset
    ds = Dataset.from_dict({"text": texts})
    
    # Tokenize
    print("Tokenizing...")
    def tokenize_batch(batch):
        return tokenizer(batch["text"], truncation=True, max_length=context_length)
    
    tokenized = ds.map(tokenize_batch, batched=True, num_proc=num_proc, remove_columns=ds.column_names)
    
    # Concatenate into chunks of exactly context_length
    print("Grouping into chunks...")
    
    def group_texts(examples):
        concatenated = {k: sum(examples[k], []) for k in examples.keys()}
        total_length = len(concatenated[list(examples.keys())[0]])
        
        # Drop remainder
        total_length = (total_length // context_length) * context_length
        
        result = {}
        for k, t in concatenated.items():
            result[k] = [t[i:i+context_length] for i in range(0, total_length, context_length)]
        
        result["labels"] = result["input_ids"].copy()
        return result
    
    lm_dataset = tokenized.map(group_texts, batched=True, num_proc=num_proc)
    
    # Split train/val
    lm_dataset = lm_dataset.train_test_split(test_size=0.001, seed=SEED)
    
    print(f"Train samples: {len(lm_dataset['train'])}")
    print(f"Validation samples: {len(lm_dataset['test'])}")
    
    lm_dataset.save_to_disk(output_dir)
    print(f"Saved to {output_dir}")
    
    return lm_dataset

if __name__ == "__main__":
    prepare_pretraining_data()

4. Train Custom Tokenizer

A custom tokenizer trained on your domain data will be more efficient than using GPT-2's. We use Byte-Pair Encoding (BPE) with the tokenizers library.

4 Train tokenizer on your corpus:

python

# scripts/train_tokenizer.py
import os
from datasets import load_dataset
from tokenizers import Tokenizer, models, pre_tokenizers, trainers, processors
from transformers import PreTrainedTokenizerFast

VOCAB_SIZE = 32000
CONTEXT_LENGTH = 512
OUTPUT_DIR = "./tokenizer"

def train_tokenizer():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    print("Loading training corpus for tokenizer...")
    # Use a sample of web text for training
    ds = load_dataset("openwebtext", streaming=True, split="train")
    
    # Collect ~100MB of text for tokenizer training
    texts = []
    total_chars = 0
    target_chars = 100_000_000  # 100MB
    
    for example in ds:
        text = example["text"]
        texts.append(text)
        total_chars += len(text)
        if total_chars >= target_chars:
            break
        if len(texts) % 10000 == 0:
            print(f"Collected {len(texts)} texts, {total_chars/1e6:.1f}MB...")
    
    print(f"Training tokenizer on {len(texts)} documents ({total_chars/1e6:.1f}MB)...")
    
    # Initialize BPE tokenizer
    tokenizer = Tokenizer(models.BPE())
    
    # Use ByteLevel pre-tokenizer (like GPT-2) - handles all Unicode
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    
    # Trainer with special tokens
    trainer = trainers.BpeTrainer(
        vocab_size=VOCAB_SIZE,
        special_tokens=[
            "<|endoftext|>",   # EOS / PAD
            "<|bos|>",         # Beginning of sequence
            "<|unk|>",         # Unknown token
            "<|pad|>",         # Padding
            "<|mask|>",        # Mask (for potential MLM use)
        ],
        min_frequency=2,
        show_progress=True,
    )
    
    # Train
    tokenizer.train_from_iterator(texts, trainer=trainer, length=len(texts))
    
    # Add post-processor for template processing
    tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)
    
    # Enable padding and truncation
    tokenizer.enable_padding(
        pad_id=tokenizer.token_to_id("<|pad|>"),
        pad_token="<|pad|>",
        length=CONTEXT_LENGTH
    )
    tokenizer.enable_truncation(max_length=CONTEXT_LENGTH)
    
    # Save raw tokenizer
    tokenizer_path = os.path.join(OUTPUT_DIR, "tokenizer.json")
    tokenizer.save(tokenizer_path)
    
    # Wrap as HuggingFace tokenizer
    hf_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        bos_token="<|bos|>",
        eos_token="<|endoftext|>",
        unk_token="<|unk|>",
        pad_token="<|pad|>",
        mask_token="<|mask|>",
    )
    
    # Save in HF format
    hf_tokenizer.save_pretrained(OUTPUT_DIR)
    
    print(f"\n✅ Tokenizer saved to {OUTPUT_DIR}")
    print(f"Vocab size: {len(hf_tokenizer)}")
    print(f"Model max length: {hf_tokenizer.model_max_length}")
    
    # Test
    test_text = "Hello world! This is a test of the custom tokenizer."
    encoded = hf_tokenizer.encode(test_text)
    decoded = hf_tokenizer.decode(encoded)
    print(f"\nTest encode/decode:")
    print(f"  Original: {test_text}")
    print(f"  Decoded:  {decoded}")
    print(f"  Tokens:   {len(encoded)}")

if __name__ == "__main__":
    train_tokenizer()

After training, your tokenizer will be saved in ./tokenizer/ with vocab size 32,000. This is more efficient than GPT-2's 50,257 vocab for general English text.

5. Model Configuration

5 Define model architecture (~100M parameters):

python

# scripts/create_model.py
from transformers import GPT2Config, GPT2LMHeadModel, AutoTokenizer
import os

OUTPUT_DIR = "./models/llm-100m"

def create_model_config():
    """
    Create a ~100M parameter GPT-style model.
    Architecture: 12 layers, 768 hidden, 12 heads, 3072 FFN
    """
    
    config = GPT2Config(
        # Core architecture
        vocab_size=32000,          # Match your tokenizer
        n_positions=512,           # Max context length
        n_embd=768,                # Hidden dimension
        n_layer=12,                # Transformer layers
        n_head=12,                 # Attention heads
        
        # FFN intermediate size (4× hidden = 3072)
        n_inner=3072,
        
        # Activation and normalization
        activation_function="gelu_new",
        layer_norm_eps=1e-5,
        
        # Dropout (0.1 for training, 0 for inference)
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
        
        # Initialization
        initializer_range=0.02,
        
        # Tie embeddings to output (saves params, often better)
        tie_word_embeddings=True,
        
        # For gradient checkpointing compatibility
        use_cache=False,
    )
    
    return config

def create_model():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained("./tokenizer")
    
    # Create config
    config = create_model_config()
    
    # Initialize model from config
    model = GPT2LMHeadModel(config)
    
    # Resize embeddings to match tokenizer (in case of mismatch)
    model.resize_token_embeddings(len(tokenizer))
    
    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"✅ Model created!")
    print(f"   Total parameters: {total_params:,} ({total_params/1e6:.1f}M)")
    print(f"   Trainable: {trainable_params:,}")
    print(f"   Layers: {config.n_layer}")
    print(f"   Hidden: {config.n_embd}")
    print(f"   Heads: {config.n_head}")
    print(f"   Context: {config.n_positions}")
    print(f"   Vocab: {config.vocab_size}")
    
    # Save
    model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)
    
    print(f"\n💾 Saved to {OUTPUT_DIR}")
    
    return model, tokenizer

if __name__ == "__main__":
    create_model()

Expected output: ~100M parameters. With tied embeddings, this is efficient for a 12-layer model.

6. Training Script

6 Main pretraining script with optimizations:

python

# scripts/pretrain.py
import os
import torch
import wandb
from transformers import (
    GPT2LMHeadModel,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    get_cosine_schedule_with_warmup,
)
from datasets import load_from_disk
from accelerate import Accelerator

# ==================== CONFIGURATION ====================

CONFIG = {
    "model_path": "./models/llm-100m",
    "data_path": "./data/pretrain",
    "output_dir": "./models/llm-100m-trained",
    
    # Training hyperparameters
    "batch_size": 16,           # Per device (adjust based on GPU memory)
    "gradient_accumulation": 4,  # Effective batch = 16 × 4 = 64
    "learning_rate": 5e-4,
    "weight_decay": 0.01,
    "warmup_steps": 1000,
    "max_steps": 50_000,        # ~3-4 epochs on 5B tokens
    "max_grad_norm": 1.0,
    
    # Optimization
    "fp16": False,            # Use bf16 if available (RTX 4090, A100)
    "bf16": True,             # Better than fp16 on Ampere+ GPUs
    "gradient_checkpointing": True,  # Trade compute for memory
    
    # Logging & Checkpointing
    "logging_steps": 100,
    "eval_steps": 1000,
    "save_steps": 5000,
    "save_total_limit": 3,
    
    # Other
    "seed": 42,
    "dataloader_num_workers": 4,
}

# ==================== TRAINING SETUP ====================

def setup_training():
    # Initialize wandb (optional)
    wandb.init(
        project="llm-pretraining",
        name="llm-100m-general",
        config=CONFIG,
    )
    
    # Set seed
    torch.manual_seed(CONFIG["seed"])
    
    # Load tokenizer and model
    print("Loading model and tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_path"])
    model = GPT2LMHeadModel.from_pretrained(CONFIG["model_path"])
    
    # Enable gradient checkpointing for memory savings
    if CONFIG["gradient_checkpointing"]:
        model.gradient_checkpointing_enable()
        model.config.use_cache = False  # Required when using gradient checkpointing
    
    # Load dataset
    print("Loading dataset...")
    if os.path.exists(CONFIG["data_path"]):
        dataset = load_from_disk(CONFIG["data_path"])
    else:
        print("Dataset not found! Run prepare_data.py first.")
        return
    
    # Data collator for causal LM (creates labels by shifting inputs)
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,  # Causal LM, not masked
    )
    
    # Training arguments
    training_args = TrainingArguments(
        output_dir=CONFIG["output_dir"],
        overwrite_output_dir=True,
        
        # Batch sizes
        per_device_train_batch_size=CONFIG["batch_size"],
        per_device_eval_batch_size=CONFIG["batch_size"],
        gradient_accumulation_steps=CONFIG["gradient_accumulation"],
        
        # Learning rate
        learning_rate=CONFIG["learning_rate"],
        weight_decay=CONFIG["weight_decay"],
        max_grad_norm=CONFIG["max_grad_norm"],
        warmup_steps=CONFIG["warmup_steps"],
        
        # Schedule
        lr_scheduler_type="cosine",
        max_steps=CONFIG["max_steps"],
        
        # Precision
        fp16=CONFIG["fp16"],
        bf16=CONFIG["bf16"],
        
        # Evaluation
        evaluation_strategy="steps",
        eval_steps=CONFIG["eval_steps"],
        
        # Logging
        logging_strategy="steps",
        logging_steps=CONFIG["logging_steps"],
        report_to=["wandb"],
        
        # Checkpointing
        save_strategy="steps",
        save_steps=CONFIG["save_steps"],
        save_total_limit=CONFIG["save_total_limit"],
        
        # Performance
        dataloader_num_workers=CONFIG["dataloader_num_workers"],
        dataloader_pin_memory=True,
        remove_unused_columns=False,
        
        # Hub
        push_to_hub=False,  # Set True to upload checkpoints
        hub_model_id=None,
        
        seed=CONFIG["seed"],
    )
    
    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        data_collator=data_collator,
    )
    
    # Train
    print("\n🚀 Starting training...\n")
    trainer.train()
    
    # Save final model
    trainer.save_model(CONFIG["output_dir"])
    tokenizer.save_pretrained(CONFIG["output_dir"])
    
    # Final evaluation
    eval_results = trainer.evaluate()
    print(f"\n📊 Final eval loss: {eval_results['eval_loss']:.4f}")
    print(f"📊 Final perplexity: {torch.exp(torch.tensor(eval_results['eval_loss'])).item():.2f}")
    
    wandb.finish()

if __name__ == "__main__":
    setup_training()

Run training: python scripts/pretrain.py

7. Memory Optimizations for Single GPU

If you run out of memory, use these techniques in order of impact:

Technique	Memory Saved	How to Enable
Gradient Checkpointing	~60%	`model.gradient_checkpointing_enable()`
BF16 / FP16	~50%	`bf16=True` in TrainingArguments
Gradient Accumulation	Scales effective batch	`gradient_accumulation_steps=4+`
Smaller Batch Size	Linear	`per_device_train_batch_size=4`
DeepSpeed ZeRO-2	~20%	Use DeepSpeed config (see below)
8-bit AdamW	~15%	`pip install bitsandbytes`, use `bnb.optim.AdamW8bit`

DeepSpeed Config (for aggressive memory optimization)

json

# ds_config.json
{
    "bf16": { "enabled": true },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8
    },
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 5e-4,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 0.01
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": { "warmup_min_lr": 0, "warmup_max_lr": 5e-4, "warmup_num_steps": 1000 }
    }
}

Run with DeepSpeed: deepspeed scripts/pretrain.py --deepspeed ds_config.json

8. Evaluation & Benchmarking

7 Evaluate your trained model:

python

# scripts/evaluate.py
import torch
from transformers import GPT2LMHeadModel, AutoTokenizer
from datasets import load_dataset
import math

def evaluate_perplexity(model_path="./models/llm-100m-trained"):
    """Calculate perplexity on validation set."""
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = GPT2LMHeadModel.from_pretrained(model_path)
    model.eval()
    
    # Load a small validation set
    ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation")
    
    encodings = tokenizer("\n\n".join(ds["text"]), return_tensors="pt")
    
    max_length = 512
    stride = 512
    seq_len = encodings.input_ids.size(1)
    
    nlls = []
    prev_end_loc = 0
    
    for begin_loc in range(0, seq_len, stride):
        end_loc = min(begin_loc + max_length, seq_len)
        trg_len = end_loc - prev_end_loc
        input_ids = encodings.input_ids[:, begin_loc:end_loc]
        target_ids = input_ids.clone()
        target_ids[:, :-trg_len] = -100
        
        with torch.no_grad():
            outputs = model(input_ids, labels=target_ids)
            neg_log_likelihood = outputs.loss * trg_len
        
        nlls.append(neg_log_likelihood)
        prev_end_loc = end_loc
        if end_loc == seq_len:
            break
    
    ppl = torch.exp(torch.stack(nlls).sum() / end_loc)
    print(f"Perplexity: {ppl.item():.2f}")
    return ppl.item()

def generate_samples(model_path="./models/llm-100m-trained", prompts=None):
    """Generate text samples to qualitatively evaluate."""
    
    if prompts is None:
        prompts = [
            "The future of artificial intelligence is",
            "Once upon a time in a distant galaxy",
            "The key to happiness lies in",
            "In the field of machine learning,",
            "The quick brown fox",
        ]
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = GPT2LMHeadModel.from_pretrained(model_path)
    model.eval()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)
    
    print("="*60)
    for prompt in prompts:
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=100,
                do_sample=True,
                temperature=0.8,
                top_k=50,
                top_p=0.95,
                num_return_sequences=1,
                pad_token_id=tokenizer.eos_token_id,
            )
        
        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
        print(f"\n📝 Prompt: {prompt}")
        print(f"🤖 Output: {generated}")
        print("-"*60)

if __name__ == "__main__":
    print("Evaluating perplexity...")
    evaluate_perplexity()
    
    print("\nGenerating samples...")
    generate_samples()

9. Inference & Text Generation

8 Simple inference script:

python

# scripts/generate.py
import torch
from transformers import pipeline

def create_generator(model_path="./models/llm-100m-trained"):
    """Create a text generation pipeline."""
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    generator = pipeline(
        "text-generation",
        model=model_path,
        tokenizer=model_path,
        device=device,
        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    )
    
    return generator

def chat_loop():
    """Interactive text generation."""
    
    print("Loading model...")
    generator = create_generator()
    
    print("\n🤖 LLM Ready! Type your prompt (or 'quit' to exit)\n")
    
    while True:
        prompt = input("You: ").strip()
        if prompt.lower() in ["quit", "exit", "q"]:
            break
        
        if not prompt:
            continue
        
        # Generate
        outputs = generator(
            prompt,
            max_new_tokens=150,
            do_sample=True,
            temperature=0.8,
            top_p=0.92,
            top_k=50,
            repetition_penalty=1.1,
            pad_token_id=generator.tokenizer.eos_token_id,
        )
        
        generated_text = outputs[0]["generated_text"]
        response = generated_text[len(prompt):].strip()
        
        print(f"\n🤖 Bot: {response}\n")

if __name__ == "__main__":
    chat_loop()

10. Optional: SFT / DPO Alignment

After pretraining, align your model with instruction-following data using TRL.

Supervised Fine-Tuning (SFT)

python

# scripts/sft_train.py
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
from datasets import load_dataset

# Load pretrained model
model = AutoModelForCausalLM.from_pretrained("./models/llm-100m-trained")
tokenizer = AutoTokenizer.from_pretrained("./models/llm-100m-trained")

# Load instruction dataset (e.g., Alpaca, Dolly, OpenAssistant)
dataset = load_dataset("tatsu-lab/alpaca", split="train")

# Format: "### Instruction:\n{instruction}\n\n### Response:\n{output}"
def format_prompt(example):
    if example["input"]:
        prompt = f"### Instruction:\n{example['instruction']}\n\n### Input:\n{example['input']}\n\n### Response:\n{example['output']}"
    else:
        prompt = f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"
    return {"text": prompt}

dataset = dataset.map(format_prompt)

# Train with SFTTrainer
training_args = TrainingArguments(
    output_dir="./models/llm-100m-sft",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    learning_rate=2e-5,
    fp16=True,
    logging_steps=10,
    save_strategy="epoch",
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=training_args,
    dataset_text_field="text",
    max_seq_length=512,
)

trainer.train()
trainer.save_model("./models/llm-100m-sft")

Direct Preference Optimization (DPO)

python

# scripts/dpo_train.py
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOTrainer, DPOConfig
from datasets import load_dataset

# Load SFT model as base
model = AutoModelForCausalLM.from_pretrained("./models/llm-100m-sft")
tokenizer = AutoTokenizer.from_pretrained("./models/llm-100m-sft")

# Load preference dataset (e.g., Anthropic HH, SHP)
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

# DPO training
training_args = DPOConfig(
    output_dir="./models/llm-100m-dpo",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    learning_rate=1e-6,
    beta=0.1,  # DPO temperature parameter
    logging_steps=10,
)

trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

trainer.train()
trainer.save_model("./models/llm-100m-dpo")

🚀 Quick Start Commands

bash

# 1. Setup
mkdir llm-project && cd llm-project
python -m venv venv
source venv/bin/activate
pip install -r requirements.txt

# 2. Train tokenizer
python scripts/train_tokenizer.py

# 3. Prepare data
python scripts/prepare_data.py

# 4. Create model
python scripts/create_model.py

# 5. Pretrain!
python scripts/pretrain.py

# 6. Evaluate
python scripts/evaluate.py

# 7. Chat with your model
python scripts/generate.py

# 8. (Optional) Align with SFT
python scripts/sft_train.py

Expected Results: After ~50K steps on 5B tokens, you should see perplexity around 15-25 on WikiText-2. The model will generate coherent sentences and understand basic context.