train.py

import torch
import torch.nn as nn

from torch.utils.data import Dataset, DataLoader , random_split
from datasets import load_dataset , concatenate_datasets
from tokenizers import Tokenizer
from tokenizers.models import BPE,WordLevel
from tokenizers.trainers import BpeTrainer,WordLevelTrainer
from tokenizers.pre_tokenizers import ByteLevel,Whitespace
from tokenizers.processors import TemplateProcessing
from  tokenizers import  decoders
from torch.cuda.amp import autocast, GradScaler

import time 

from torch.utils.tensorboard import SummaryWriter
from itertools import islice

from config import get_weights_file_path, get_config

from tqdm import tqdm

from pathlib import Path

import warnings
from dataset import BilingualDataset
from model import build_gpt

g = torch.Generator()
g.manual_seed(23)

def get_all_sentences(ds):
    for item in ds:
        yield item['text']


def get_or_build_tokenizer(config, ds):
    tokenizer_path = Path(config['tokenizer_file'])

    if not tokenizer_path.exists():
        # Define tokenizer with BPE model
        tokenizer = Tokenizer(BPE(unk_token="<unk>"))

        # ByteLevel pre-tokenizer and decoder
        tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=True)
        tokenizer.decoder = decoders.ByteLevel()

        # Optional: Add post-processing for special tokens
        tokenizer.post_processor = TemplateProcessing(
            single="<s> $A </s>",
            pair="<s> $A </s> <s> $B </s>",
            special_tokens=[
                ("<s>", 0),
                ("</s>", 1),
            ],
        )

        # Trainer
        trainer = BpeTrainer(
            vocab_size = 30000,
            min_frequency=2,
            special_tokens=["<s>", "</s>", "<pad>", "<unk>", "<mask>","<user>","<ai>","<search_start>","<search_end>","<think>","</think>"]
        )

        # Train from dataset
        tokenizer.train_from_iterator(get_all_sentences(ds), trainer=trainer)

        # Save as single .json file
        tokenizer.save(str(tokenizer_path))

    else:
        tokenizer = Tokenizer.from_file(str(tokenizer_path))

    return tokenizer

def get_ds(config):
    # ds_raw  = load_dataset("json",data_files={'train':config['train'],'test':config['test']})
    ds_raw  = load_dataset("json",data_files='./dataset/openwebtext_500k_docs.jsonl',split="train",streaming=True)
    ds_test  = load_dataset("json",data_files='./dataset/openwebtext_test.jsonl',split="train",streaming=True)
    # ds_raw = ds_raw[:1]
    # ds_raw  = load_dataset("stas/openwebtext-10k")
    tokenizer = get_or_build_tokenizer(config,ds_raw)
    # tokenizer = get_or_build_tokenizer(config,ds_raw)
    train_ds = BilingualDataset(ds_raw, tokenizer, config['seq_len'])
    val_ds = BilingualDataset(ds_test, tokenizer, config['seq_len'])
    train_dataloader = DataLoader(train_ds, num_workers=6,prefetch_factor=2,pin_memory=True,batch_size=config['batch_size'])
    val_dataloader = DataLoader(val_ds, batch_size=1)
    
    return train_dataloader, val_dataloader, tokenizer

def get_model(config, vocab_size):
    # model = build_transformer(vocab_src_len,vocab_tgt_len,config['seq_len'],config['seq_len'],config['d_model'], config['N'] , config['h'], config['d_ff'])
    model = build_gpt( vocab_size, config['seq_len'], config['d_model'], config['N']  , config['q_head'],config['kv_head'], config['d_ff'],config['dropout'])
    return model

def validate_model(val_dataloader, model,device,loss_fn,vocab_size):
    total_loss = 0
    model.eval()
    i = 0
    with torch.no_grad():
        for batch in val_dataloader:
            input_tokens = batch['input'].to(device,non_blocking=True)
            label = batch['label'].to(device,non_blocking=True)
            decoder_output = model.decode(input_tokens)
            project_output = model.project(decoder_output)
            total_loss += loss_fn(
                        project_output.view(-1,vocab_size), 
                        label.view(-1)
                    )
            i+=1
        print(f"Validation loss : {total_loss/i:4f}")
            
            


def train_model(config):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device : {device}")

    # Enable TF32 (optional, speeds up matmul)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    Path(config['model_folder']).mkdir(parents=True, exist_ok=True)

    train_dataloader , val_dataloader, tokenizer = get_ds(config)
    print(tokenizer.get_vocab_size())
    model = get_model(config, tokenizer.get_vocab_size()).to(device)
    # TensorBoard
    writer = SummaryWriter(config['experiment_name'])

    optimizer = torch.optim.AdamW(model.parameters(), lr=config['lr'], betas=(0.9, 0.95), weight_decay=0.1)
    scaler = GradScaler()  # <- added scaler for mixed precision

    initial_epoch = 0
    global_step = 0
    tqdm_state = {'n':0}
    
    model_filename = None
    if config['preload']:
        model_filename = get_weights_file_path(config, config['preload'])
        print(f"Preloading Model {model_filename}")
        state = torch.load(model_filename)
        model.load_state_dict(state['model_state_dict'])
        optimizer.load_state_dict(state['optimizer_state_dict'])
        initial_epoch = state['epoch'] if 'mid-' in model_filename else state['epoch'] + 1 
        global_step = state['global_step']
        tqdm_state = state['tqdm_state']  if 'mid-' in model_filename else {'n':0}
    else:
        print("No Model to preload. Setting from scratch.")

    loss_fn = nn.CrossEntropyLoss(
        ignore_index=tokenizer.token_to_id('<pad>'), 
        label_smoothing=0.05
    ).to(device)
    e = 0
    
    try:
        for epoch in range(initial_epoch, config['num_epochs']):
            model.train()
            batch_iterator = tqdm(islice(train_dataloader,tqdm_state['n'],None), desc=f'Processing epoch {epoch:02d}',initial=tqdm_state['n'] ,total=140000)#total=217013)
            e = epoch
            if 'elapsed' in tqdm_state and "mid-" in model_filename :
                batch_iterator.start_t = time.time() - tqdm_state['elapsed']
            # total_len = len(batch_iterator)
            for batch in batch_iterator:
                input_tokens = batch['input'].to(device,non_blocking=True)
                label = batch['label'].to(device,non_blocking=True)
                optimizer.zero_grad(set_to_none=True)
                with autocast(dtype=torch.float16):
                    decoder_output = model.decode(input_tokens)
                    project_output = model.project(decoder_output)  # (B, Seq_len, tgt_vocab_size)
                    loss = loss_fn(
                        project_output.view(-1, tokenizer.get_vocab_size()), 
                        label.view(-1)
                    )
                if global_step%10 ==0:
                    batch_iterator.set_postfix({f"loss": f"{loss.item():6.3f}"})
                    writer.add_scalar("train loss", loss.item(), global_step)
                    writer.flush()
                if global_step % 10000 == 0 and global_step != 0:
                    validate_model(val_dataloader,model,device,loss_fn,tokenizer.get_vocab_size())
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                global_step += 1
                tqdm_state = {'n': batch_iterator.n,'elapsed':batch_iterator.format_dict["elapsed"]}
                # if()
            tqdm_state['n'] = 0
            del tqdm_state['elapsed']

            model_filename = get_weights_file_path(config, f'{epoch:02d}')
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'global_step': global_step,
                'tqdm_state':tqdm_state
            }, model_filename)
            validate_model(validate_model,model,device,loss_fn,tokenizer.get_vocab_size())
    except KeyboardInterrupt:
        print("You are stoping training : ... ")
        model_filename = get_weights_file_path(config, f'mid-{e:02d}{input("Checkpoint Name: ")}')
        torch.save({
            'epoch': e,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'global_step': global_step,
            'tqdm_state':tqdm_state
        }, model_filename)
        
if __name__ == "__main__":
    warnings.filterwarnings('ignore')
    config = get_config("./openweb.config.json")
    train_model(config)