train_slm.py

# train_slm.py
import os, math, torch
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import LambdaLR
from datasets import load_from_disk
from transformers import AutoTokenizer
from config_slm import SLMConfig
from embedding_module import TokenPositionalEmbedding, EmbeddingConfig
from transformer_block import TransformerBlock, TransformerBlockConfig
import json


# Build model
class SmallLanguageModel(torch.nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        mcfg = cfg.model

        self.embed = TokenPositionalEmbedding(
            EmbeddingConfig(
                vocab_size=mcfg.vocab_size,
                d_model=mcfg.d_model,
                block_size=mcfg.block_size,
                dropout=mcfg.dropout
            )
        )

        # stack of transformer blocks
        self.blocks = torch.nn.ModuleList([
            TransformerBlock(
                TransformerBlockConfig(
                    d_model=mcfg.d_model,
                    n_heads=mcfg.n_heads_per_layer[i],
                    block_size=mcfg.block_size,
                    ffn_hidden_mult=mcfg.ffn_hidden_mult,
                    ffn_activation=mcfg.ffn_activation,
                    attn_dropout=mcfg.dropout,
                    resid_dropout=mcfg.dropout
                )
            )
            for i in range(mcfg.n_layers)
        ])

        self.ln_f = torch.nn.LayerNorm(mcfg.d_model)
        self.lm_head = torch.nn.Linear(mcfg.d_model, mcfg.vocab_size, bias=False)

        # tie weights (important)
        self.lm_head.weight = self.embed.token_emb.weight

    def forward(self, idx, targets=None):
        x = self.embed(idx)  # (B, T, d_model)

        for block in self.blocks:
            x, _ = block(x)

        x = self.ln_f(x)
        logits = self.lm_head(x)  # (B, T, vocab_size)

        loss = None
        if targets is not None:
            loss = torch.nn.functional.cross_entropy(
                logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1
            )
        return logits, loss



# Load data
def get_dataloaders(cfg):
    ds_train = load_from_disk(os.path.join(cfg.dataset.dataset_path, "train"))
    ds_val = load_from_disk(os.path.join(cfg.dataset.dataset_path, "validation"))

    def collate_fn(batch):
        input_ids = [torch.tensor(item["input_ids"], dtype=torch.long) for item in batch]
        labels = [torch.tensor(item["labels"], dtype=torch.long) for item in batch]
        input_ids = torch.stack(input_ids)
        labels = torch.stack(labels)
        return {"input_ids": input_ids, "labels": labels}

    train_loader = DataLoader(ds_train, batch_size=cfg.dataset.batch_size,
                              shuffle=True, num_workers=cfg.dataset.num_workers,
                              collate_fn=collate_fn)
    val_loader = DataLoader(ds_val, batch_size=cfg.dataset.batch_size,
                            shuffle=False, num_workers=cfg.dataset.num_workers,
                            collate_fn=collate_fn)
    return train_loader, val_loader



def get_lr_scheduler(opt, cfg, total_steps):
    """Builds a warmup + cosine/linear scheduler."""
    warmup_steps = int(cfg.training.warmup_ratio * total_steps)
    min_lr = cfg.training.lr * cfg.training.min_lr_ratio
    max_lr = cfg.training.lr

    def lr_lambda(step):
        if step < warmup_steps:
            return float(step) / float(max(1, warmup_steps))
        progress = (step - warmup_steps) / float(max(1, total_steps - warmup_steps))
        progress = min(1.0, max(0.0, progress))

        if cfg.training.scheduler_type == "cosine":
            return min_lr / max_lr + 0.5 * (1 + math.cos(math.pi * progress)) * (1 - min_lr / max_lr)
        elif cfg.training.scheduler_type == "linear":
            return 1.0 - progress * (1.0 - min_lr / max_lr)
        elif cfg.training.scheduler_type == "constant":
            return 1.0
        else:
            raise ValueError(f"Unknown scheduler type {cfg.training.scheduler_type}")

    return LambdaLR(opt, lr_lambda)


def train(cfg):
    device = cfg.training.device if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    model = SmallLanguageModel(cfg).to(device)
    train_loader, val_loader = get_dataloaders(cfg)

    opt = torch.optim.AdamW(
        model.parameters(),
        lr=cfg.training.lr,
        weight_decay=cfg.training.weight_decay
    )

    total_steps = len(train_loader) * cfg.training.max_epochs
    scheduler = get_lr_scheduler(opt, cfg, total_steps)

    global_step = 0
    for epoch in range(cfg.training.max_epochs):
        model.train()
        total_loss = 0
        for step, batch in enumerate(train_loader):
            opt.zero_grad()
            input_ids = batch["input_ids"].to(device)
            labels = batch["labels"].to(device)
            _, loss = model(input_ids, labels)
            loss.backward()

            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.training.grad_clip)
            opt.step()
            scheduler.step()
            global_step += 1

            total_loss += loss.item()
            if step % 100 == 0:
                lr = scheduler.get_last_lr()[0]
                print(f"Epoch {epoch+1} | Step {step} | LR {lr:.6f} | Loss {loss.item():.4f}")

        avg_loss = total_loss / len(train_loader)
        print(f"Epoch {epoch+1} complete. Train Loss: {avg_loss:.4f}")

        # validation
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for batch in val_loader:
                input_ids = batch["input_ids"].to(device)
                labels = batch["labels"].to(device)
                _, loss = model(input_ids, labels)
                val_loss += loss.item()
        val_loss /= len(val_loader)
        print(f"Validation Loss: {val_loss:.4f}")

    os.makedirs(cfg.training.save_dir, exist_ok=True)
    print("💾 Saving model manually to", cfg.training.save_dir)

    # Save model weights
    torch.save(model.state_dict(), os.path.join(cfg.training.save_dir, "pytorch_model.bin"))

    # Save model config for reproducibility
    model_config = {
        "vocab_size": cfg.model.vocab_size,
        "d_model": cfg.model.d_model,
        "n_layers": cfg.model.n_layers,
        "n_heads_per_layer": cfg.model.n_heads_per_layer,
        "block_size": cfg.model.block_size,
        "dropout": cfg.model.dropout,
        "ffn_hidden_mult": cfg.model.ffn_hidden_mult,
        "ffn_activation": cfg.model.ffn_activation,
    }
    with open(os.path.join(cfg.training.save_dir, "config.json"), "w") as f:
        json.dump(model_config, f, indent=2)

    # Save tokenizer
    tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
    tokenizer.save_pretrained(cfg.training.save_dir)

    print(f"✅ Model + tokenizer saved at {cfg.training.save_dir}")



# Run
if __name__ == "__main__":
    cfg = SLMConfig()
    train(cfg)