Repo before implementing concepts of the paper memorizing transformer

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+data/hellaswag/hellaswag_val.jsonl filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+data/edu_fineweb10B
+log/model*
+!log/model_final.pt

configs/config.json

@@ -0,0 +1,21 @@
+{
+  "model": {
+    "block_size": 1024,
+    "vocab_size": 50304,
+    "n_layer": 12,
+    "n_head": 12,
+    "n_embd": 768
+  },
+  "training": {
+    "max_steps": 19073,
+    "log_dir": "log",
+    "total_batch_size": 524288,
+    "B": 64,
+    "T": 1024,
+    "max_lr": 0.0006,
+    "min_lr": 0.00006,
+    "warmup_steps": 715,
+    "weight_decay": 0.1,
+    "learning_rate": 0.0006
+  }
+}

data/fineweb.py

@@ -0,0 +1,79 @@
+"""
+FineWeb-Edu dataset (for srs pretraining)
+https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu
+Downloads and tokenizes the data and saves data shards to disk.
+Will save shards to the local directory "edu_fineweb10B".
+"""
+import os
+import multiprocessing as mp
+import numpy as np
+import tiktoken
+from datasets import load_dataset 
+from tqdm import tqdm 
+
+
+local_dir = "edu_fineweb10B"
+remote_name = "sample-10BT"
+shard_size = int(1e8) # 100M tokens per shard, total of 100 shards
+
+DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), local_dir)
+os.makedirs(DATA_CACHE_DIR, exist_ok=True)
+print("Shards will be saved to:",DATA_CACHE_DIR)
+
+#dataset download
+fw = load_dataset("HuggingFaceFW/fineweb-edu", name=remote_name, split="train")
+
+#tokenizer
+enc = tiktoken.get_encoding("gpt2")
+eot = enc._special_tokens['<|endoftext|>'] # end of text token
+
+def tokenize(doc):
+    # tokenizes a single document and returns a numpy array of uint16 tokens
+    tokens = [eot]
+    tokens.extend(enc.encode_ordinary(doc["text"]))
+    tokens_np = np.array(tokens)
+    assert (0 <= tokens_np).all() and (tokens_np < 2**16).all(), "token dictionary too large for uint16"
+    tokens_np_uint16 = tokens_np.astype(np.uint16)
+    return tokens_np_uint16
+
+def write_datafile(filename, tokens_np):
+    np.save(filename, tokens_np)
+
+nprocs = max(1, os.cpu_count()//2)
+with mp.Pool(nprocs) as pool:
+    shard_index = 0
+    # preallocate buffer to hold current shard
+    all_tokens_np = np.empty((shard_size,), dtype=np.uint16)
+    token_count = 0
+    progress_bar = None
+    for tokens in pool.imap(tokenize, fw, chunksize=16):
+
+        # is there enough space in the current shard for the new tokens?
+        if token_count + len(tokens) < shard_size:
+            # simply append tokens to current shard
+            all_tokens_np[token_count:token_count+len(tokens)] = tokens
+            token_count += len(tokens)
+            # update progress bar
+            if progress_bar is None:
+                progress_bar = tqdm(total=shard_size, unit="tokens", desc=f"Shard {shard_index}")
+            progress_bar.update(len(tokens))
+        else:
+            # write the current shard and start a new one
+            split = "val" if shard_index == 0 else "train"
+            filename = os.path.join(DATA_CACHE_DIR, f"edufineweb_{split}_{shard_index:06d}")
+            # split the document into whatever fits in this shard; the remainder goes to next one
+            remainder = shard_size - token_count
+            progress_bar.update(remainder)
+            all_tokens_np[token_count:token_count+remainder] = tokens[:remainder]
+            write_datafile(filename, all_tokens_np)
+            shard_index += 1
+            progress_bar = None
+            # populate the next shard with the leftovers of the current doc
+            all_tokens_np[0:len(tokens)-remainder] = tokens[remainder:]
+            token_count = len(tokens)-remainder
+
+    # write any remaining tokens as the last shard
+    if token_count != 0:
+        split = "val" if shard_index == 0 else "train"
+        filename = os.path.join(DATA_CACHE_DIR, f"edufineweb_{split}_{shard_index:06d}")
+        write_datafile(filename, all_tokens_np[:token_count])

evaluation/hellaswag.py

@@ -0,0 +1,113 @@
+"""
+Downloads and evaluates HellaSwag in Python.
+https://github.com/rowanz/hellaswag
+
+"""
+import os
+import json
+import requests
+import tiktoken
+from tqdm import tqdm
+import torch
+from torch.nn import functional as F
+
+DATA_DOWNLOADED_PATH = '"data/hellaswag"'
+
+def download_file(url:str, fname:str, chunk_size=1024):
+    resp = requests.get(url, stream=True)
+    total = int(resp.headers.get("content-length", 0 ))
+    with open(fname, "wb") as file, tqdm(
+        desc = fname,
+        total=total,
+        unit="iB",
+        unit_scale=True,
+        unit_divisor=1024
+    )as bar:
+        for data in resp.iter_content(chunk_size=chunk_size):
+            size = file.write(data)
+            bar.update(size)
+
+hellaswags = {
+    "train": "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_train.jsonl",
+    "val": "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl",
+    "test": "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_test.jsonl",
+}
+
+enc = tiktoken.get_encoding("gpt2")
+
+def download(split):
+    """Downloads HellaSwag DATA_DOWNLOADED_PATH"""
+    os.makedirs(DATA_DOWNLOADED_PATH, exist_ok=True)
+    data_url = hellaswags[split]
+    data_filename = os.path.join(DATA_DOWNLOADED_PATH, f"hellaswag_{split}.jsonl")
+    if not os.path.exists(data_filename):
+        print(f"Downloading {data_url} to {data_filename}...")
+        download_file(data_url, data_filename)
+
+def render_example(example):
+    """
+    Given the example as a dictionary, render it as three torch tensors:
+    - tokens (the tokens of context + completion, of size 4xN, as there are always 4 candidates)
+    - mask (is 1 in the region of the candidate completion, where we evaluate likelihoods)
+    - label (the index of the correct completion, which we hope has the highest likelihood)
+    """
+    ctx = example["ctx"]
+    label = example["label"]
+    endings = example["endings"]
+
+    # data needed to reproduce this eval on the C size
+    data = {
+        "label": label,
+        "ctx_tokens": None,
+        "ending_tokens": [],
+    }
+
+    # gather up all the tokens
+    ctx_tokens = enc.encode(ctx)
+    data["ctx_tokens"] = ctx_tokens
+    tok_rows = []
+    mask_rows = []
+    for end in endings:
+        end_tokens = enc.encode(" " + end) # note: prepending " " because GPT-2 tokenizer
+        tok_rows.append(ctx_tokens + end_tokens)
+        mask_rows.append([0]*len(ctx_tokens) + [1]*len(end_tokens))
+        data["ending_tokens"].append(end_tokens)
+
+    # have to be careful during the collation because the number of tokens in each row can differ
+    max_len = max(len(row) for row in tok_rows)
+    tokens = torch.zeros((4, max_len), dtype=torch.long)
+    mask = torch.zeros((4, max_len), dtype=torch.long)
+    for i, (tok_row, mask_row) in enumerate(zip(tok_rows, mask_rows)):
+        tokens[i, :len(tok_row)] = torch.tensor(tok_row)
+        mask[i, :len(mask_row)] = torch.tensor(mask_row)
+
+    return data, tokens, mask, label
+
+def iterate_examples(split):
+    # there are 10,042 examples in total in val
+    download(split)
+    with open(os.path.join(DATA_DOWNLOADED_PATH, f"hellaswag_{split}.jsonl"), "r") as f:
+        for line in f:
+            example = json.loads(line)
+            yield example
+
+
+def get_most_likely_row(tokens, mask, logits):
+        shift_logits = (logits[..., :-1, :]).contiguous() #this will be x for loss calculation
+        shift_tokens = (tokens[..., 1:]).contiguous() #this will be y for loss calculation
+        shift_mask = (mask[..., 1:]).contiguous() #shifting same as tokens shifted
+        flat_shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+        flat_shift_tokens = shift_tokens.view(-1)
+        shift_losses = F.cross_entropy(flat_shift_logits, flat_shift_tokens, reduction='none')
+        shift_losses = shift_losses.view(tokens.size(0), -1)
+        masked_shift_losses = shift_losses * shift_mask
+        sum_loss = masked_shift_losses.sum(dim=1)
+        avg_loss = sum_loss / shift_mask.sum(dim=1)
+        pred_norm = avg_loss.argmin().item() #taking the index of minimum loss
+        return pred_norm   
+
+
+
+
+
+

evaluation/val_hellaswag.py

@@ -0,0 +1,94 @@
+import torch
+from ..hellaswag import render_example, iterate_examples, get_most_likely_row
+import torch.distributed as dist 
+from torch.distributed import init_process_group, destroy_process_group
+from torch.nn.parallel import DistributedDataParallel as DDP
+import os
+from ..ModelGPT2 import GPT,log_file
+
+ddp = int(os.environ.get('RANK', -1)) != -1 #will be True if ddp run
+if ddp:
+    assert torch.cuda.is_available()
+    init_process_group(backend='nccl')
+    ddp_rank = int(os.environ['RANK'])
+    ddp_local_rank = int(os.environ['LOCAL_RANK'])
+    ddp_world_size = int(os.environ['WORLD_SIZE'])
+    device = f"cuda:{ddp_local_rank}"
+    torch.cuda.set_device(device)
+    master_process = ddp_rank == 0 #this is the process doing checkpoint,logging,etc
+else:
+    ddp_rank = 0
+    ddp_local_rank = 0
+    ddp_world_size = 1
+    master_process = True
+    #attempt to autodetect the device
+    device = 'cpu'
+    if torch.cuda.is_available():
+        device = 'cuda'
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"  #for mac users use apple silicon cpu which allready have gpu.mps is backend for apple silicon
+    print(f"Using device: {device}")
+# device = "cpu" #OVERRIDE
+
+device_type = "cuda" if device.startswith("cuda") else "cpu"
+
+torch.manual_seed(1337)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed(1337)
+
+
+#Creating model by loading the model weights
+checkpoint_path = '../log/model_final.pt' 
+if master_process:
+        print(f"Loading checkpoint from {checkpoint_path}")
+
+checkpoint = torch.load(checkpoint_path, map_location=device)
+
+# Extract config and create model
+model_config = checkpoint['config']
+model_config.vocab_size = 50304 #for computational effciency(power of 2)
+model = GPT(model_config)
+# Load model state dict
+model.load_state_dict(checkpoint['model'])
+model = DDP(model, device_ids=[ddp_local_rank])
+model.to(device)
+
+
+def evaluate_hellaswag(model, device, device_type, ddp, ddp_rank, ddp_world_size, log_file, master_process):
+         
+    num_correct_norm = 0
+    num_total = 0
+
+    for i, example in enumerate(iterate_examples("val")):
+        # only process example where i % ddp_world_size ==ddp_rank#this is for proper managemnt of which part is deal by which gpu
+        if ddp:
+            if i % ddp_world_size != ddp_rank:
+                continue
+        #rendering example into tokens and labels
+        _, tokens, mask, label = render_example(example)
+        tokens = tokens.to(device)
+        mask = mask.to(device)
+        #get the logits
+        with torch.no_grad():
+            with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+                logits, loss = model(tokens)
+            pred_norm = get_most_likely_row(tokens, mask, logits)
+        num_total += 1
+        num_correct_norm += int(pred_norm == label)
+    #reduce the stats accross all process
+    if ddp:
+        num_total = torch.tensor(num_total, dtype=torch.long, device=device)
+        num_correct_norm = torch.tensor(num_correct_norm, dtype=torch.long, device=device)
+        dist.all_reduce(num_total, op=dist.ReduceOp.SUM)
+        dist.all_reduce(num_correct_norm, op=dist.ReduceOp.SUM)
+        num_total = num_total.item()
+        num_correct_norm = num_correct_norm.item()
+    acc_norm = num_correct_norm / num_total #accuracy of hellaswag
+    if master_process:
+        print(f"HellaSwag accuracy: {num_correct_norm}/{num_total}={acc_norm:.4f}")
+        with open(log_file, "a") as f:
+            f.write(f"Final Hellaswag accuracy: {acc_norm:.4f}\n")   
+
+evaluate_hellaswag(model, device, device_type, ddp, ddp_rank, ddp_world_size, log_file, master_process)
+if ddp:
+    destroy_process_group()

log/log.txt

@@ -0,0 +1 @@
+0 val 10.9528

model_core/attention.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class CasualSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        self.c_proj.NANOGPT_SCALE_INIT = 1
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+
+    def forward(self, x):
+        B, T, C = x.size()
+        qkv = self.c_attn(x)
+        q, k, v = qkv.split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1,2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1,2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1,2) # (B, nh, T, hs)
+
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True) #flash attention
+
+        y = y.transpose(1,2).contiguous().view(B, T, C) # (B, T, C) 
+        y = self.c_proj(y)
+        return y

model_core/dataloader.py

@@ -0,0 +1,53 @@
+import os
+import numpy as np
+import torch
+
+#Data loader
+class DataLoader_1:
+    def __init__(self, B, T, process_rank, num_processes, split, master_process):
+        self.B = B
+        self.T = T
+        self.process_rank = process_rank
+        self.num_processes = num_processes
+        assert split in {'train', 'val'}
+        
+        
+        data_root = "data/edu_fineweb10B"
+        shards = os.listdir(data_root)
+        shards = [s for s in shards if split in s]
+        shards = sorted(shards)
+        shards = [os.path.join(data_root, s) for s in shards]
+        self.shards = shards
+        assert len(shards)> 0, f"no shards found for split {split}"   
+        if master_process:
+            print(f"found {len(shards)} shards for split {split}")  
+        self.reset()
+
+    def load_tokens(self, filename):
+        npt = np.load(filename)
+        npt = npt.astype(np.int32)
+        ptt = torch.tensor(npt, dtype=torch.long)
+        return ptt
+    
+
+    def reset(self):
+    #state, init at shard 0
+        self.current_shard = 0
+        self.tokens = self.load_tokens(self.shards[self.current_shard])
+        self.current_position = self.B * self.T * self.process_rank 
+
+    def next_batch(self):
+        B, T = self.B, self.T
+        buf = self.tokens[self.current_position:self.current_position + B*T+1]
+        x = (buf[:-1]).view(B,T) 
+        y = (buf[1:]).view(B,T) 
+        
+        self.current_position += B * T * self.num_processes
+
+        if self.current_position + (B * T * self.num_processes + 1) > len(self.tokens):
+            self.current_shard = (self.current_shard + 1) % len(self.shards)
+            self.tokens = self.load_tokens(self.shards[self.current_shard])
+            self.current_position = B * T * self.process_rank
+        return x, y
+
+

model_core/model.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+import inspect
+from .attention import CasualSelfAttention
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
+        self.gelu = nn.GELU(approximate='tanh')
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
+        self.c_proj.NANOGPT_SCALE_INIT = 1
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        return x
+    
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = CasualSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024 #max sequence length
+    vocab_size: int = 50257 #number of tokens: 50000 BPE merges + 256 byte tokens +1 special token which is endoftext
+    n_layer: int = 12 #number of layers
+    n_head: int = 12 #number of heads
+    n_embd: int = 768 #embedding dimensions
+
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = nn.LayerNorm(config.n_embd),
+        ))
+
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        #Weight sharing scheme
+        self.transformer.wte.weight = self.lm_head.weight
+
+        # init params
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            std = 0.02
+            if hasattr(module, 'NANOGPT_SCALE_INIT'):
+                std *= (2 * self.config.n_layer) ** -0.5
+            torch.nn.init.normal_(module.weight, mean = 0.0, std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.size()
+        assert T <=self.config.block_size, f"Cannot forward sequence of length {T} ,block size is only {self.config.block_size}"
+
+        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
+        pos_emb = self.transformer.wpe(pos)
+        tok_emb = self.transformer.wte(idx)
+        x = tok_emb + pos_emb
+
+        for block in self.transformer.h:
+            x = block(x)
+
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x) #(B, T, vocab_size)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        
+        return logits, loss
+    
+    def configure_optimizers(self, weight_decay, learning_rate, device_type, master_process):
+        param_dict = {pn:p for pn, p in self.named_parameters()}
+        param_dict = {pn:p for pn, p in param_dict.items() if p.requires_grad}
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [{'params':decay_params, ' weight_decay': weight_decay},
+                       {'params':nodecay_params, 'weight_decay': 0.0}
+                       ]
+        num_decay_params = sum(p.numel() for p in decay_params)  
+        num_nodecay_params = sum(p.numel() for p in nodecay_params) 
+        if master_process:
+            print(f"num decayed parameters tensors: {len(decay_params)}, with{num_decay_params}:parameters")
+            print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == "cuda"    
+        if master_process:
+            print(f"using fused AdamW: {use_fused}")
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=(0.9,0.95), eps=1e-8, fused=use_fused)
+        return optimizer 
+    

model_core/training.py

@@ -0,0 +1,171 @@
+from torch.distributed import init_process_group, destroy_process_group
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
+import os
+import torch
+import time
+import json
+import math
+from .model import GPT,GPTConfig
+
+
+def train_memgpt(config_path,dataloader_class=None):    
+
+    with open(config_path,'r') as f:
+        cfg = json.load(f)
+    
+    model_cfg_params = cfg['model']
+    train_cfg_params = cfg['training']
+
+    ddp = int(os.environ.get('RANK', -1)) != -1
+    if ddp:
+        assert torch.cuda.is_available()
+        init_process_group(backend='nccl')
+        ddp_rank = int(os.environ['RANK'])
+        ddp_local_rank = int(os.environ['LOCAL_RANK'])
+        ddp_world_size = int(os.environ['WORLD_SIZE'])
+        device = f"cuda:{ddp_local_rank}"
+        torch.cuda.set_device(device)
+        master_process = ddp_rank == 0
+    else:
+        ddp_rank = 0
+        ddp_local_rank = 0
+        ddp_world_size = 1
+        master_process = True
+        device = 'cpu'
+        if torch.cuda.is_available():
+            device = 'cuda'
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            device = "mps"
+        if master_process:
+            print(f"Using device: {device}")
+
+    device_type = "cuda" if device.startswith("cuda") else "cpu"
+
+    torch.manual_seed(1337)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(1337)
+    
+
+       
+    total_batch_size = train_cfg_params['total_batch_size']
+    B = train_cfg_params['B']
+    T = train_cfg_params['T']
+    max_steps = train_cfg_params['max_steps']
+    log_dir = train_cfg_params['log_dir']
+    max_lr = train_cfg_params['max_lr']
+    min_lr = train_cfg_params['min_lr']
+    warmup_steps = train_cfg_params['warmup_steps']
+    weight_decay = train_cfg_params['weight_decay']
+    base_learning_rate = train_cfg_params['learning_rate'] 
+
+    assert total_batch_size % (B * T * ddp_world_size) == 0
+    grad_accum_steps = total_batch_size // (B * T * ddp_world_size)
+    if master_process:
+        print(f"Total desired batch size: {total_batch_size}")
+        print(f"Calculated gradient accumulation steps: {grad_accum_steps}")
+
+    train_loader = dataloader_class(B=B, T=T, process_rank=ddp_rank, num_processes=ddp_world_size, split="train",master_process=master_process)
+    val_loader = dataloader_class(B=B, T=T, process_rank=ddp_rank, num_processes=ddp_world_size, split="val",master_process=master_process)
+
+    torch.set_float32_matmul_precision('high')
+
+    # Create Model
+    model = GPT(GPTConfig(**model_cfg_params))
+    model.to(device)
+    use_compile = True
+    if use_compile:
+        model = torch.compile(model)
+    if ddp:
+        model = DDP(model, device_ids=[ddp_local_rank])
+    raw_model = model.module if ddp else model
+
+    def get_lr(it):
+        if it < warmup_steps:
+            return max_lr * (it + 1) / warmup_steps
+        if it > max_steps:
+            return min_lr
+        decay_ratio = (it - warmup_steps) / (max_steps - warmup_steps)
+        assert 0 <= decay_ratio <= 1
+        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+        return min_lr + coeff * (max_lr - min_lr)
+
+    optimizer = raw_model.configure_optimizers(weight_decay=weight_decay, learning_rate=base_learning_rate, device_type=device_type, master_process=master_process)
+
+    os.makedirs(log_dir, exist_ok=True)
+    log_file = os.path.join(log_dir, "log.txt")
+    with open(log_file, "w") as f:
+        pass
+
+    for step in range(max_steps):
+        t0 = time.time()
+        last_step = (step == max_steps - 1)
+
+        if step % 350 == 0 or last_step:
+            model.eval()
+            val_loader.reset()
+            with torch.no_grad():
+                val_loss_accum = 0.0
+                val_loss_steps = 20
+                for _ in range(val_loss_steps):
+                    x, y = val_loader.next_batch()
+                    x, y = x.to(device), y.to(device)
+                    with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+                        logits, loss = model(x, y)
+                    loss = loss / val_loss_steps
+                    val_loss_accum += loss.detach()
+            if ddp:
+                dist.all_reduce(val_loss_accum, op=dist.ReduceOp.AVG)
+            if master_process:
+                print(f"Validation loss: {val_loss_accum.item():.4f}")
+                with open(log_file, "a") as f:
+                    f.write(f"{step} val {val_loss_accum.item():.4f}\n")
+
+                checkpoint_name = f"model_final.pt" if last_step else f"model_{step:05d}.pt"
+                checkpoint_path = os.path.join(log_dir, checkpoint_name)
+
+                checkpoint = {
+                    'model': raw_model.state_dict(),
+                    'optimizer': optimizer.state_dict(),
+                    'step': step,
+                    'val_loss': val_loss_accum.item(),
+                    'config': raw_model.config
+                }
+                torch.save(checkpoint, checkpoint_path)
+
+               
+        model.train()
+        optimizer.zero_grad()
+        loss_accum = 0.0
+        for micro_step in range(grad_accum_steps):
+            x, y = train_loader.next_batch()
+            x, y = x.to(device), y.to(device)
+            if ddp:
+                model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1)
+            with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+                logits, loss = model(x, y)
+            loss = loss / grad_accum_steps
+            loss_accum += loss.detach()
+            loss.backward()
+
+        if ddp:
+            dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG)
+
+        norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        lr = get_lr(step)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+        optimizer.step()
+        if device_type == 'cuda':
+            torch.cuda.synchronize()
+        t1 = time.time()
+        dt = (t1 - t0) * 1000
+        tokens_processed = train_loader.B * train_loader.T * grad_accum_steps * ddp_world_size
+        tokens_per_sec = tokens_processed / dt
+        if master_process:
+            print(f"Step:{step:5d} | Loss: {loss_accum.item():.6f} | lr: {lr:.4e} | Norm:{norm:.4f} | dt: {dt:.2f}ms | Tok/sec: {tokens_per_sec:.2f}")
+            with open(log_file, 'a') as f:
+                f.write(f"{step} train {loss_accum.item():.6f}\n")
+
+    if ddp:
+        destroy_process_group()

requirement.txt

@@ -0,0 +1,10 @@
+--extra-index-url https://download.pytorch.org/whl/cu121
+
+safetensors==0.5.3
+tiktoken==0.9.0
+tokenizers==0.21.1
+transformers==4.50.1
+tqdm==4.67.1
+requests==2.32.3
+numpy<1.27,>=1.22
+torch==2.3.1+cu121

scripts/evaluate.py

@@ -0,0 +1,2 @@
+#To run all evaluation at once
+#Code yet to be added

scripts/generate.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn.functional as F
+import tiktoken
+from model import GPT
+
+def generate_text(model, prompt, num_return_sequences=4, max_length=32, device='cuda'):
+        model.eval()
+        enc = tiktoken.get_encoding('gpt2')
+        tokens = enc.encode(prompt)
+        tokens = torch.tensor(tokens, dtype=torch.long)
+        tokens = tokens.unsqueeze(0).repeat(num_return_sequences, 1)
+        xgen = tokens.to(device)
+        sample_rng = torch.Generator(device=device)
+        sample_rng.manual_seed(42)
+        
+        while xgen.size(1) < max_length:
+            with torch.no_grad():
+                logits, loss = model(xgen)  # (B, T, vocab_size)
+                logits = logits[:, -1, :]  # (B, vocab_size)
+                probs = F.softmax(logits, dim=-1)  
+                topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)  
+                ix = torch.multinomial(topk_probs, 1, generator=sample_rng)  
+                xcol = torch.gather(topk_indices, -1, ix)  
+                xgen = torch.cat((xgen, xcol), dim=1)  
+        
+        generated_texts = []
+        for i in range(num_return_sequences):
+            tokens = xgen[i, :max_length].tolist()
+            decoded = enc.decode(tokens)
+            generated_texts.append(decoded)
+            print(f"Sample {i + 1}: {decoded}")
+
+        
+        return generated_texts
+    
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"running with {device}")
+
+
+checkpoint_path = 'log/model_final.pt'
+
+print(f"Loading checkpoint from {checkpoint_path}")
+checkpoint = torch.load(checkpoint_path,map_location=device)
+model_config = checkpoint['config']
+model_config.vocab_size = 50304 
+model = GPT(model_config)
+
+
+model.load_state_dict(checkpoint['model'])
+model.to(device)
+
+
+
+prompt = "Hello, I'm a language model,"
+
+generated_texts = generate_text(
+        model=model,
+        prompt=prompt,
+        num_return_sequences=4,
+        max_length=32,
+        device=device
+    )

scripts/train.py

@@ -0,0 +1,11 @@
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from model_core.training import train_memgpt
+from model_core.dataloader import DataLoader_1
+
+if __name__ == "__main__":
+    config_path = "configs/config.json"
+    print("Training starter")
+    train_memgpt(config_path=config_path,dataloader_class=DataLoader_1)