train.py

import argparse
import os
import sys
import shutil
import random
import numpy as np
import time
import copy
import math
import matplotlib.pyplot as plt

import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.autograd import Variable
import transformers
from transformers import GPT2TokenizerFast

# --- Mask for causal (autoregressive) attention ---
def subsequent_mask(size):
    """Mask out subsequent positions."""
    attn_shape = (1, size, size)
    mask = torch.triu(torch.ones(attn_shape), diagonal=1).bool()
    return mask

# --- Corpus reading ---
def read_corpus(filename, tokenizer):
    print(f"Reading corpus from {filename}...")
    seq = []
    with open(filename, 'rt') as f:
        for line in f:
            line = line.rstrip('\n')
            tokens = tokenizer(line)
            seq.extend(tokens['input_ids'])
    print(f"Read {len(seq)} tokens from {filename}")
    return seq

class Embedder(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.d_model = d_model
        self.embed = nn.Embedding(vocab_size, d_model)
    def forward(self, x):
        return self.embed(x.long())

class PositionalEncoder(nn.Module):
    def __init__(self, d_model, max_seq_len = 4096, dropout = 0.1):
        super().__init__()
        self.d_model = d_model
        self.dropout = nn.Dropout(dropout)
        # create constant 'pe' matrix with values dependent on pos and i
        pe = torch.zeros(max_seq_len, d_model)
        for pos in range(max_seq_len):
            for i in range(0, d_model, 2):
                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i)/d_model)))
                if i + 1 < d_model:
                    pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1))/d_model)))
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        # make embeddings relatively larger
        x = x * math.sqrt(self.d_model)
        # add constant to embedding
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len]
        return self.dropout(x)

class Norm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.size = d_model
        # create two learnable parameters to calibrate normalisation
        self.alpha = nn.Parameter(torch.ones(self.size))
        self.bias = nn.Parameter(torch.zeros(self.size))
        self.eps = eps
    
    def forward(self, x):
        norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) / \
               (x.std(dim=-1, keepdim=True) + self.eps) + self.bias
        return norm

def attention(q, k, v, d_k, mask=None, dropout=None):
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        mask = mask.unsqueeze(1)
        scores = scores.masked_fill(mask == 0, -1e9)
    scores = F.softmax(scores, dim=-1)
    if dropout is not None:
        scores = dropout(scores)
    output = torch.matmul(scores, v)
    return output

class MultiHeadAttention(nn.Module):
    def __init__(self, heads, d_model, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.d_k = d_model // heads
        self.h = heads
        self.q_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        self.out = nn.Linear(d_model, d_model)
    
    def forward(self, q, k, v, mask=None):
        bs = q.size(0)
        # perform linear operation and split into N heads
        k = self.k_linear(k).view(bs, -1, self.h, self.d_k)
        q = self.q_linear(q).view(bs, -1, self.h, self.d_k)
        v = self.v_linear(v).view(bs, -1, self.h, self.d_k)
        # transpose to get dimensions bs * N * sl * d_model
        k = k.transpose(1,2)
        q = q.transpose(1,2)
        v = v.transpose(1,2)
        # calculate attention
        scores = attention(q, k, v, self.d_k, mask, self.dropout)
        # concatenate heads and put through final linear layer
        concat = scores.transpose(1,2).contiguous().view(bs, -1, self.d_model)
        output = self.out(concat)
        return output

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff=2048, dropout=0.1):
        super().__init__()
        # We set d_ff as a default to 2048
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)
    
    def forward(self, x):
        x = self.dropout(F.relu(self.linear_1(x)))
        x = self.linear_2(x)
        return x

def get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

class DecoderLayer(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super().__init__()
        self.norm_1 = Norm(d_model)
        self.norm_2 = Norm(d_model)
        self.attn = MultiHeadAttention(heads, d_model, dropout)
        self.ff = FeedForward(d_model, dropout=dropout)
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)
    
    def forward(self, x, trg_mask):
        x2 = self.norm_1(x)
        x = x + self.dropout_1(self.attn(x2, x2, x2, trg_mask))
        x2 = self.norm_2(x)
        x = x + self.dropout_2(self.ff(x2))
        return x

class Decoder(nn.Module):
    def __init__(self, vocab_size, d_model, N, heads, dropout):
        super().__init__()
        self.N = N
        self.embed = Embedder(vocab_size, d_model)
        self.pe = PositionalEncoder(d_model, dropout=dropout)
        self.layers = get_clones(DecoderLayer(d_model, heads, dropout), N)
        self.norm = Norm(d_model)
    
    def forward(self, x, trg_mask):
        x = self.embed(x)
        x = self.pe(x)
        for layer in self.layers:
            x = layer(x, trg_mask)
        return self.norm(x)

class GPT2LM(nn.Module):
    def __init__(self, vocab_size, d_model, N, heads, dropout, tie_weights=False):
        super().__init__()
        self.decoder = Decoder(vocab_size, d_model, N, heads, dropout)
        self.out = nn.Linear(d_model, vocab_size)
        if tie_weights:
            self.out.weight = self.decoder.embed.embed.weight
            print("✅ Tied embeddings enabled.")
    
    def forward(self, x, mask):
        d_output = self.decoder(x, mask)
        return self.out(d_output)

# --- Data batching for arbitrary sizes ---
def batchify(data, batch_size, seq_len):
    nbatch = len(data) // batch_size
    data = torch.tensor(data[:nbatch * batch_size], dtype=torch.long)
    data = data.view(batch_size, -1)
    for i in range(0, data.size(1) - 1, seq_len):
        seq_len_i = min(seq_len, data.size(1) - 1 - i)
        src = data[:, i:i + seq_len_i]
        tgt = data[:, i + 1:i + 1 + seq_len_i]
        yield src, tgt

# --- Training and evaluation loops with tokens per second reporting ---
def train_model(model, opt):
    print("Starting training...")
    model.train()
    train_ppls = []
    valid_ppls = []
    total_tokens = 0
    total_time = 0
    
    for epoch in range(opt.epochs):
        total_loss = 0
        batches = 0
        epoch_tokens = 0
        epoch_start_time = time.time()
        
        for src, tgt in batchify(opt.train, opt.batchsize, opt.seqlen):
            batch_start_time = time.time()
            src, tgt = src.to(opt.device), tgt.to(opt.device)
            mask = subsequent_mask(src.size(1)).to(opt.device)
            output = model(src, mask)
            loss = F.cross_entropy(output.view(-1, opt.vocab_size), tgt.reshape(-1), ignore_index=opt.src_pad)
            opt.optimizer.zero_grad()
            loss.backward()
            opt.optimizer.step()
            
            batch_time = time.time() - batch_start_time
            total_loss += loss.item()
            batches += 1
            
            # Count tokens in this batch: batch_size * sequence_length
            tokens_in_batch = src.numel()
            epoch_tokens += tokens_in_batch
            total_tokens += tokens_in_batch
            total_time += batch_time
            
            tokens_per_sec = tokens_in_batch / batch_time
            if batches % opt.printevery == 0:
                print(f"Epoch {epoch+1}, Batch {batches}, Loss: {loss.item():.4f}, Speed: {tokens_per_sec:.2f} tokens/sec")
        
        epoch_time = time.time() - epoch_start_time
        epoch_tokens_per_sec = epoch_tokens / epoch_time
        
        avg_loss = total_loss / batches
        train_ppl = math.exp(avg_loss)
        train_ppls.append(train_ppl)
        
        print(f"Epoch {epoch+1}/{opt.epochs}, Loss: {avg_loss:.4f}, Perplexity: {train_ppl:.2f}")
        print(f"Epoch training speed: {epoch_tokens_per_sec:.2f} tokens/sec")
        
        valid_ppl = test_model(model, opt.valid, opt, tag=f"valid-epoch{epoch+1}")
        valid_ppls.append(valid_ppl)

    # Report final training speed
    avg_tokens_per_sec = total_tokens / total_time
    print(f"\nTraining completed.")
    print(f"Average training speed: {avg_tokens_per_sec:.2f} tokens/sec")
    
    # Save training speed to file
    with open(os.path.join("saved", opt.dir_name, "training_speed.txt"), "w") as f:
        f.write(f"Total tokens processed: {total_tokens}\n")
        f.write(f"Total training time: {total_time:.2f} seconds\n")
        f.write(f"Average training speed: {avg_tokens_per_sec:.2f} tokens/sec\n")

    # Ensure directory exists
    dir_name = os.path.join("saved", opt.dir_name)
    if not os.path.exists(dir_name):
        os.makedirs(dir_name)
        print(f"Created directory: {dir_name}")

    # Save the model
    save_path = os.path.join(dir_name, "gpt2lm_wiki103.pth")
    print(f"Saving model to: {save_path}")
    print(f"Directory exists: {os.path.exists(dir_name)}")
    print(f"Write permissions: {os.access(dir_name, os.W_OK)}")
    torch.save(model.state_dict(), save_path)
    print(f"Model saved successfully to {save_path}")

    # Plot learning curve
    plt.plot(range(1, opt.epochs+1), train_ppls, label='Train PPL')
    plt.plot(range(1, opt.epochs+1), valid_ppls, label='Valid PPL')
    plt.xlabel("Epoch")
    plt.ylabel("Perplexity")
    plt.legend()
    plt.title("Training & Validation Perplexity")
    plt.savefig(os.path.join(dir_name, "learning_curve.png"))
    print(f"Saved learning curve to {dir_name}/learning_curve.png")

    # Save perplexity log
    with open(os.path.join(dir_name, "perplexity_log.txt"), "w") as f:
        for i in range(opt.epochs):
            f.write(f"Epoch {i+1}: Train PPL = {train_ppls[i]:.2f}, Valid PPL = {valid_ppls[i]:.2f}\n")
    
    return avg_tokens_per_sec

def test_model(model, data, opt, tag="valid"):
    print(f"Running {tag} set...")
    model.eval()
    total_loss, batches = 0, 0
    with torch.no_grad():
        for src, tgt in batchify(data, opt.batchsize, opt.seqlen):
            src, tgt = src.to(opt.device), tgt.to(opt.device)
            mask = subsequent_mask(src.size(1)).to(opt.device)
            output = model(src, mask)
            loss = F.cross_entropy(output.view(-1, opt.vocab_size), tgt.reshape(-1), ignore_index=opt.src_pad)
            total_loss += loss.item()
            batches += 1
    avg_loss = total_loss / batches
    ppl = math.exp(avg_loss)
    print(f"{tag.capitalize()} PPL: {ppl:.2f}")
    model.train()
    return ppl

def main():
    random.seed(10)
    parser = argparse.ArgumentParser()
    parser.add_argument('-no_cuda', action='store_true')
    parser.add_argument('-SGDR', action='store_true')
    parser.add_argument('-epochs', type=int, default=1)  # Reduced for faster iteration
    parser.add_argument('-d_model', type=int, default=512)
    parser.add_argument('-n_layers', type=int, default=6)
    parser.add_argument('-heads', type=int, default=8)
    parser.add_argument('-dropout', type=float, default=0.1)
    parser.add_argument('-batchsize', type=int, default=8)  # Increased batch size
    parser.add_argument('-printevery', type=int, default=100)
    parser.add_argument('-lr', type=float, default=0.0001)  # Slightly higher learning rate
    parser.add_argument('-seqlen', type=int, default=512)
    parser.add_argument('-threshold', type=int, default=3)
    parser.add_argument('-savename', type=str)
    parser.add_argument('-loadname', type=str)
    parser.add_argument('-tied', type=int, default=1)
    parser.add_argument('-dir_name', type=str, default='wiki103_model')
    parser.add_argument('-norm', type=float, default=2.0)
    opt = parser.parse_args()
    opt.verbose = False
    
    # Use GPU if available and not explicitly disabled
    if not opt.no_cuda and torch.cuda.is_available():
        opt.device = torch.device("cuda:0")
        print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
    else:
        opt.device = torch.device("cpu")
        print("Using CPU for training")

    time_name = time.strftime("%y%m%d_%H%M%S")
    opt.time_name = time_name
    dir_name = "saved/%s" % (opt.dir_name)
    if not os.path.exists(dir_name):
        os.makedirs(dir_name)
    source_name = sys.argv[0]
    shutil.copy(source_name, dir_name + "/" + os.path.basename(source_name))
    opt.log_file = dir_name + "/log_file.txt"
    print(str(opt))

    # Load and tokenize Wikitext-103 dataset
    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    # Changed file paths for Wikitext-103
    opt.train = read_corpus('wiki103.train.txt', tokenizer)
    opt.valid = read_corpus('wiki103.valid.txt', tokenizer)
    opt.test = read_corpus('wiki103.test.txt', tokenizer)
    
    obs = len(opt.train)
    print(f"Training set: {obs} tokens")
    
    opt.vocab_size = 50257  # GPT-2 tokenizer vocabulary size
    temp = []
    for i in range(opt.vocab_size):
        temp.append(i)
    opt.indices = torch.tensor(temp).to(opt.device)

    # Initialize model
    model = GPT2LM(opt.vocab_size, opt.d_model, opt.n_layers, opt.heads, opt.dropout, tie_weights=(opt.tied == 1)).to(opt.device)
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    text = 'total params: %d' % (params)
    print(text)

    # Choose optimizer
    opt.optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr, betas=(0.9, 0.98), eps=1e-9)
    
    if opt.savename is not None:
        try:
            os.mkdir(opt.savename)
        except:
            nothing = 1
    opt.src_pad = 0
    opt.trg_pad = 0

    # Train and evaluate
    avg_tokens_per_sec = train_model(model, opt)
    test_model(model, opt.valid, opt, tag="valid")
    test_model(model, opt.test, opt, tag="test")
    
    print(f"\nFinal training speed: {avg_tokens_per_sec:.2f} tokens/sec")
    
    # Write a summary of speed improvement suggestions
    with open(os.path.join(dir_name, "speed_improvement_suggestions.txt"), "w") as f:
        f.write("Suggestions for improving training speed:\n\n")
        f.write("1. Use mixed precision training (FP16/BF16)\n")
        f.write("2. Increase batch size and use gradient accumulation\n")
        f.write("3. Enable tensor core operations on compatible GPUs\n")
        f.write("4. Optimize data loading with prefetching and parallel workers\n")
        f.write("5. Use model parallelism or distributed training\n")
        f.write("6. Consider using optimized implementations like FlashAttention\n")
        f.write("7. Experiment with smaller model sizes or pruning\n")
        f.write("8. Profile and optimize bottlenecks\n")
        f.write("9. Use memory-efficient optimizers\n")
        f.write("10. Consider efficient implementations like xformers or rotary embeddings\n")

if __name__ == "__main__":
    main()