removed

Naive_gpt/config.py

@@ -1,16 +0,0 @@
-# Model/config.py
-import torch
-
-# Define hyperparameters and constants
-BATCH_SIZE = 16
-BLOCK_SIZE = 1024
-MAX_ITERS = 5
-EVAL_INTERVAL = 500
-LEARNING_RATE = 6e-4
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-EVAL_ITERS = 200
-N_EMBD = 768
-N_HEAD = 12
-N_LAYER = 12
-DROPOUT = 0.2
-MODEL_PATH = "Naive_gpt\model_weights_llama"  # Where to save weights

Naive_gpt/config_1.5B.py

@@ -1,16 +0,0 @@
-# Model/config.py
-import torch
-
-# Define hyperparameters and constants
-BATCH_SIZE = 16
-BLOCK_SIZE = 1024
-MAX_ITERS = 5
-EVAL_INTERVAL = 500
-LEARNING_RATE = 6e-4
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-EVAL_ITERS = 200
-N_EMBD = 1600
-N_HEAD = 25
-N_LAYER = 48
-DROPOUT = 0.2
-MODEL_PATH = "Naive_gpt\model_weights_llama"  # Where to save weights

Naive_gpt/config_355.py

@@ -1,16 +0,0 @@
-# Model/config.py
-import torch
-
-# Define hyperparameters and constants
-BATCH_SIZE = 16
-BLOCK_SIZE = 1024
-MAX_ITERS = 5
-EVAL_INTERVAL = 500
-LEARNING_RATE = 6e-4
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-EVAL_ITERS = 200
-N_EMBD = 1024
-N_HEAD = 16
-N_LAYER = 24
-DROPOUT = 0.2
-MODEL_PATH = "Naive_gpt\model_weights_llama"  # Where to save weights

Naive_gpt/config_770.py

@@ -1,16 +0,0 @@
-# Model/config.py
-import torch
-
-# Define hyperparameters and constants
-BATCH_SIZE = 16
-BLOCK_SIZE = 1024
-MAX_ITERS = 5
-EVAL_INTERVAL = 500
-LEARNING_RATE = 6e-4
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-EVAL_ITERS = 200
-N_EMBD = 1280
-N_HEAD = 20
-N_LAYER = 36
-DROPOUT = 0.2
-MODEL_PATH = "Naive_gpt\model_weights_llama"  # Where to save weights

Naive_gpt/data_loader.py

@@ -1,252 +0,0 @@
-# Model/data_loader.py
-import torch
-import os
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-class TextDataLoader:
-    def __init__(self, file_path, batch_size, block_size, tokenizer, chunk_size=10**4):
-        self.file_path = file_path
-        self.batch_size = batch_size
-        self.block_size = block_size
-        self.tokenizer = tokenizer
-        self.chunk_size = chunk_size
-        self.file = open(self.file_path, 'r', encoding='utf-8')
-        self.data = None
-        self.end_of_file = False
-
-        # Load the initial chunk of data
-        self.load_chunk()
-
-    def load_chunk(self):
-        """Load a chunk from the file, encode it, and handle end-of-file conditions."""
-        text = self.file.read()
-        if not text:
-            self.end_of_file = True
-            logging.info("End of file reached.")
-        else:
-            try:
-                # Encode the text using the tokenizer
-                encoded = self.tokenizer.encode(text)
-                if len(encoded) > 0:
-                    self.data = torch.tensor(encoded, dtype=torch.long)
-                    logging.info(f"Loaded new data chunk of size: {len(self.data)} tokens.")
-                    # save the encoded data to a file
-                    torch.save(self.data, "encoded_data.pth")
-            except Exception as e:
-                logging.error(f"Error encoding text chunk: {e}")
-                self.end_of_file = True
-
-    def num_batches(self):
-        """Calculate the total number of batches in the current chunk."""
-        if self.data is not None:
-            return (len(self.data) - 1) // self.block_size  # Total batches in the current chunk
-        return 0
-
-    def get_batch(self):
-        """Retrieve a batch of data from the current chunk or load a new chunk if needed."""
-        if self.end_of_file:
-            return None, None  # Return None when no data is left
-
-        # Generate a batch of data
-        ix = torch.randint(len(self.data) - self.block_size, (self.batch_size,))
-        x = torch.stack([self.data[i:i+self.block_size] for i in ix])
-        y = torch.stack([self.data[i+1:i+self.block_size+1] for i in ix])
-        return x, y
-
-    def reset(self):
-        """Reset the file and flags for a new epoch."""
-        self.file.seek(0)
-        self.end_of_file = False
-        logging.info("Resetting file for a new epoch.")
-        self.load_chunk()
-
-    def close(self):
-        """Clean up file resources when done."""
-        self.file.close()
-        logging.info("File closed.")
-
-    def __iter__(self):
-        """Make the data loader iterable so it can be used in a loop."""
-        while not self.end_of_file:
-            x, y = self.get_batch()
-            if x is None or y is None:
-                break  # Stop iteration if there's no more data
-
-            yield x, y  # Yield a batch of data
-
-        # Once iteration is done, close the file
-        self.close()
-
-#before parallelizing
-# Set up logging
-# logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-# class TextDataLoader:
-#     def __init__(self, file_path, batch_size, block_size, tokenizer, device='cpu', chunk_size=10**4):
-#         self.file_path = file_path
-#         self.batch_size = batch_size
-#         self.block_size = block_size
-#         self.tokenizer = tokenizer
-#         self.device = device
-#         self.chunk_size = chunk_size
-#         self.file = open(self.file_path, 'r', encoding='utf-8')
-#         self.data = None
-#         self.end_of_file = False
-
-#         # Load the initial chunk of data
-#         self.load_chunk()
-
-#     def load_chunk(self):
-#         """Load a chunk from the file, encode it, and handle end-of-file conditions."""
-#         text = self.file.read()
-#         if not text:
-#             self.end_of_file = True
-#             logging.info("End of file reached.")
-#         else:
-#             try:
-#                 # Encode the text using the tokenizer
-#                 encoded = self.tokenizer.encode(text)
-#                 if len(encoded) > 0:
-#                     self.data = torch.tensor(encoded, dtype=torch.long).to(self.device)
-#                     logging.info(f"Loaded new data chunk of size: {len(self.data)} tokens.")
-#             except Exception as e:
-#                 logging.error(f"Error encoding text chunk: {e}")
-#                 self.end_of_file = True
-
-#     def num_batches(self):
-#         """Calculate the total number of batches in the current chunk."""
-#         if self.data is not None:
-#             return (len(self.data) - 1) // self.block_size  # Total batches in the current chunk
-#         return 0
-
-#     def get_batch(self):
-#         """Retrieve a batch of data from the current chunk or load a new chunk if needed."""
-#         if self.end_of_file:
-#             return None, None  # Return None when no data is left
-        
-#         # Generate a batch of data
-#         ix = torch.randint(len(self.data) - self.block_size, (self.batch_size,))
-#         x = torch.stack([self.data[i:i+self.block_size] for i in ix])
-#         y = torch.stack([self.data[i+1:i+self.block_size+1] for i in ix])
-#         return x, y
-
-#     def reset(self):
-#         """Reset the file and flags for a new epoch."""
-#         self.file.seek(0)
-#         self.end_of_file = False
-#         logging.info("Resetting file for a new epoch.")
-#         self.load_chunk()
-
-#     def close(self):
-#         """Clean up file resources when done."""
-#         self.file.close()
-#         logging.info("File closed.")
-
-#     def __iter__(self):
-#         """Make the data loader iterable so it can be used in a loop."""
-#         while not self.end_of_file:
-#             x, y = self.get_batch()
-#             if x is None or y is None:
-#                 break  # Stop iteration if there's no more data
-            
-#             yield x, y  # Yield a batch of data
-
-#         # Once iteration is done, close the file
-#         self.close()
-
-
-# # Set up logging
-# logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-# class TextDataLoader:
-#     def __init__(self, file_path, batch_size, block_size, tokenizer, device='cpu', chunk_size=10**4):
-#         self.file_path = file_path
-#         self.batch_size = batch_size
-#         self.block_size = block_size
-#         self.tokenizer = tokenizer
-#         self.device = device
-#         self.chunk_size = chunk_size
-#         self.file = open(self.file_path, 'r', encoding='utf-8')
-#         self.data = None
-#         self.end_of_file = False
-
-#         # Print a preview of the file
-#         # self.print_file_preview()
-        
-#         # Initial chunk loading
-#         self.load_chunk()
-
-#     def print_file_preview(self):
-#         """Prints the first few lines of the text file for preview"""
-#         self.file.seek(0)  # Go to the beginning of the file
-#         lines = [self.file.readline() for _ in range(5)]
-#         preview_text = ''.join(lines)
-#         print("File preview:\n", preview_text)
-#         self.file.seek(0)  # Reset to the start of the file for chunk reading
-
-#     def load_chunk(self):
-#         """Load a chunk from the file, encode it, and handle end-of-file conditions."""
-#         text = self.file.read()
-#         if not text:
-#             self.end_of_file = True
-#             logging.info("End of file reached.")
-#         else:
-#             try:
-#                 # Log the first 100 characters of the text chunk to verify Urdu content
-#                 # logging.info(f"First 100 characters of the chunk: {text[:100]}")
-#                 # print("This is the chunk:", text)
-
-#                 # Encode the text using the tokenizer
-#                 # print("Tokenizer:", self.tokenizer)
-#                 encoded = self.tokenizer.encode(text)
-#                 print(len(encoded))
-#                 print("encoded data: ")
-
-#                 # Log the encoded output length to confirm successful encoding
-#                 logging.info(f"Encoded data length: {len(encoded)} tokens")
-
-#                 # if len(encoded) < self.block_size:
-#                 #     # Only stop if there's absolutely no usable data left
-#                 #     self.end_of_file = len(encoded) == 0
-#                 #     if self.end_of_file:
-#                 #         logging.warning("Insufficient data in chunk; stopping further loading.")
-#                 #     else:
-#                 #         logging.warning("Data chunk smaller than block size loaded; may limit training batch size.")
-                
-#                 if len(encoded) > 0:
-#                     self.data = torch.tensor(encoded, dtype=torch.long).to(self.device)
-#                     logging.info(f"Loaded new data chunk of size: {len(self.data)} tokens.")
-#             except Exception as e:
-#                 logging.error(f"Error encoding text chunk: {e}")
-#                 self.end_of_file = True
-
-#     def get_batch(self):
-#         """Retrieve a batch of data from the current chunk or load a new chunk if needed."""
-#         # if self.end_of_file:
-#         #     return None, None  # Return None when no data is left
-        
-#         # if self.data is None or len(self.data) <= self.block_size:
-#         #     self.load_chunk()
-#         #     if self.end_of_file or self.data is None or len(self.data) < self.block_size:
-#         #         return None, None  # Stop if there’s insufficient data
-
-#         # Generate a batch of data
-#         ix = torch.randint(len(self.data) - self.block_size, (self.batch_size,))
-#         x = torch.stack([self.data[i:i+self.block_size] for i in ix])
-#         y = torch.stack([self.data[i+1:i+self.block_size+1] for i in ix])
-#         return x, y
-
-#     def reset(self):
-#         """Reset the file and flags for a new epoch."""
-#         self.file.seek(0)
-#         self.end_of_file = False
-#         logging.info("Resetting file for a new epoch.")
-#         self.load_chunk()
-
-#     def close(self):
-#         """Clean up file resources when done."""
-#         self.file.close()
-#         logging.info("File closed.")

Naive_gpt/model.py

@@ -1,258 +0,0 @@
-# Model/model.py
-import torch
-import torch.nn as nn
-from .config import *
-import inspect
-
-
-class CausalSelfAttention(nn.Module):
-
-    def __init__(self):
-        super().__init__()
-        assert N_EMBD % N_HEAD == 0
-        # key, query, value projections for all heads, but in a batch
-        self.c_attn = nn.Linear(N_EMBD, 3 * N_EMBD)
-        # output projection
-        self.c_proj = nn.Linear(N_EMBD, N_EMBD)
-        self.c_proj.NANOGPT_SCALE_INIT = 1
-        # regularization
-        self.n_head = N_HEAD
-        self.n_embd = N_EMBD
-
-    def forward(self, x):
-        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
-        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
-        # nh is "number of heads", hs is "head size", and C (number of channels) = nh * hs
-        # e.g. in GPT-2 (124M), n_head=12, hs=64, so nh*hs=C=768 channels in the Transformer
-        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 = nn.functional.scaled_dot_product_attention(
-            q, k, v, is_causal=True)  # flash attention
-        # re-assemble all head outputs side by side
-        y = y.transpose(1, 2).contiguous().view(B, T, C)
-        # output projection
-        y = self.c_proj(y)
-        return y
-    
-# class Head(nn.Module):  #this is sebastian's causal attention
-#     """ one head of self-attention """
-
-#     def __init__(self, head_size):
-#         super().__init__()
-#         self.key = nn.Linear(N_EMBD, head_size, bias=False)
-#         self.query = nn.Linear(N_EMBD, head_size, bias=False)
-#         self.value = nn.Linear(N_EMBD, head_size, bias=False)
-#         self.register_buffer('tril', torch.tril(torch.ones(BLOCK_SIZE, BLOCK_SIZE)))
-#         self.dropout = nn.Dropout(DROPOUT)
-
-#     def forward(self, x):
-#         B, T, C = x.shape
-#         k = self.key(x)   # (B, T, head_size)
-#         q = self.query(x) # (B, T, head_size)
-#         wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5
-#         wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
-#         wei = nn.functional.softmax(wei, dim=-1)
-#         wei = self.dropout(wei)
-#         v = self.value(x)
-#         out = wei @ v
-#         return out
-
-# class MultiHeadAttention(nn.Module):
-#     """ multiple heads of self-attention in parallel """
-
-#     def __init__(self, num_heads, head_size):
-#         super().__init__()
-#         self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
-#         self.proj = nn.Linear(head_size * num_heads, N_EMBD)
-#         self.dropout = nn.Dropout(DROPOUT)
-
-#     def forward(self, x):
-#         out = torch.cat([h(x) for h in self.heads], dim=-1)
-#         out = self.dropout(self.proj(out))
-#         return out
-
-class FeedFoward(nn.Module):   #yeh MLP hai karpathy wala -> Feed forward hai sebastian wala
-    def __init__(self):
-        super().__init__()
-        self.c_fc = nn.Linear(N_EMBD, 4 * N_EMBD)
-        self.gelu = nn.GELU(approximate='tanh')
-        self.c_proj = nn.Linear(4 * N_EMBD, 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
-    """ a simple linear layer followed by a non-linearity """
-
-    # def __init__(self, n_embd):
-    #     super().__init__()
-    #     self.net = nn.Sequential(
-    #         nn.Linear(N_EMBD, 4 * N_EMBD),
-    #         nn.ReLU(),
-    #         nn.Linear(4 * N_EMBD, N_EMBD),
-    #         nn.Dropout(DROPOUT),
-    #     )
-
-    # def forward(self, x):
-    #     return self.net(x)
-
-class Block(nn.Module):
-    """ Transformer block: communication followed by computation """
-
-    def __init__(self, n_embd, n_head):
-        super().__init__()
-        head_size = N_EMBD // n_head
-        self.sa = CausalSelfAttention()
-        self.ffwd = FeedFoward()
-        self.ln1 = nn.LayerNorm(N_EMBD)
-        self.ln2 = nn.LayerNorm(N_EMBD)
-
-    def forward(self, x):
-        x = x + self.sa(self.ln1(x))
-        x = x + self.ffwd(self.ln2(x))
-        return x
-
-class GPTLanguageModel(nn.Module):
-
-    def __init__(self, vocab_size, config):
-        super().__init__()
-        print("This is vocab size:", vocab_size)
-        self.token_embedding_table = nn.Embedding(vocab_size, config.n_embd)
-        self.position_embedding_table = nn.Embedding(config.block_size, config.n_embd)
-        self.blocks = nn.Sequential(
-            *[Block(config.n_embd, n_head=config.n_head) for _ in range(config.n_layer)]
-        )
-        self.ln_f = nn.LayerNorm(config.n_embd)
-        self.lm_head = nn.Linear(config.n_embd, vocab_size)
-
-        self.token_embedding_table.weight = self.lm_head.weight
-
-        self.apply(self._init_weights)
-        self.config = {"BLOCK_SIZE": config.block_size, "N_EMBD": config.n_embd, "N_HEAD":config.n_head, "N_LAYER": config.n_layer}
-
-
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-                std = 0.02
-                if hasattr(module, 'NANOGPT_SCALE_INIT'):
-                    std *= (2 * 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 _init_weights(self, module):
-    #     if isinstance(module, nn.Linear):
-    #         torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
-    #         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.shape
-        assert T <= BLOCK_SIZE, f"Cannot forward sequence of length {T}, block size is only {BLOCK_SIZE}"
-
-
-        tok_emb = self.token_embedding_table(idx)
-        pos_emb = self.position_embedding_table(torch.arange(0, T, dtype=torch.long, device=idx.device))
-        x = tok_emb + pos_emb
-        x = self.blocks(x)
-        x = self.ln_f(x)
-        logits = self.lm_head(x)
-
-        if targets is None:
-            loss = None
-        else:
-            # B, T, C = logits.shape
-            # logits = logits.view(B*T, C)
-            # targets = targets.view(B*T)
-            # loss = nn.functional.cross_entropy(logits, targets)
-            loss = nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
-
-        return logits, loss
-
-    def generate(self, idx, max_new_tokens, temperature=1.0):
-        """
-        Generate tokens using the language model.
-        Args:
-            idx: Input token indices
-            max_new_tokens: Number of tokens to generate
-            temperature: Controls randomness in generation
-                        - temperature > 1.0 increases randomness
-                        - temperature < 1.0 decreases randomness
-                        - temperature = 0 makes it deterministic (always picks highest probability)
-        """
-        for _ in range(max_new_tokens):
-            # Truncate the sequence to the last BLOCK_SIZE tokens
-            idx_cond = idx[:, -BLOCK_SIZE:]
-            # Get logits from the model
-            logits, _ = self(idx_cond)
-            # Focus only on the last time step
-            logits = logits[:, -1, :]
-
-            if temperature == 0.0:
-                # For temperature = 0, simply take the argmax
-                idx_next = torch.argmax(logits, dim=-1, keepdim=True)
-            else:
-                # Apply temperature scaling
-                logits = logits / temperature
-                # Convert to probabilities
-                probs = torch.softmax(logits, dim=-1)
-                # Sample from the distribution
-                idx_next = torch.multinomial(probs, num_samples=1)
-
-            # Append the new token to the sequence
-            idx = torch.cat((idx, idx_next), dim=1)
-        return idx
-
-    # def save(self, path=MODEL_PATH):
-    #     torch.save(self.state_dict(), path)
-
-    # def load(self, path=MODEL_PATH):
-    #     self.load_state_dict(torch.load(path,map_location=torch.device('cpu')))
-
-    def save(self, path=MODEL_PATH):
-        torch.save(self.state_dict(), path)
-
-    # def load(self, path=MODEL_PATH):
-    #     self.load_state_dict(torch.load(path))
-
-    def load(self, path=MODEL_PATH):
-    # Load the state dict
-        state_dict = torch.load(path)["model"]
-
-        # Rename the keys to match the expected ones (remove "orig_mod." prefix)
-        new_state_dict = {}
-        for key, value in state_dict.items():
-            new_key = key.replace('_orig_mod.', '')  # Remove 'orig_mod.' prefix
-            new_state_dict[new_key] = value
-
-        # Load the renamed state dict into the model
-        self.load_state_dict(new_state_dict)
-
-
-    def configure_optimizers(self, weight_decay=0.1, learning_rate=LEARNING_RATE, device=DEVICE):
-        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_parameters = [p for n, p in param_dict.items() if p.dim() >= 2]
-        nodecay_parameters = [p for n, p in param_dict.items() if p.dim() < 2]
-        optim_groups = [
-            {"params": decay_parameters, "weight_decay": weight_decay},
-            {"params": nodecay_parameters, "weight_decay": 0.0},
-        ]
-        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
-        use_fused = fused_available and device == "cuda"
-        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=(0.9, 0.95), eps=1e-8, fused = use_fused)
-        return optimizer

Naive_gpt/train.py

@@ -1,303 +0,0 @@
-import torch
-from tqdm import tqdm
-from .config import *
-from .data_loader import TextDataLoader
-from .model import GPTLanguageModel
-import math
-
-max_lr = 6e-4
-min_lr = max_lr * 0.1
-warmup_steps = 715
-max_steps = 19073
-
-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)
-
-total_batch_size = 524288
-assert total_batch_size % (BATCH_SIZE * BLOCK_SIZE) == 0, "make sure total_batch_size is divisible by BATCH_SIZE * BLOCK_SIZE"
-grad_accumulation_steps = total_batch_size // (BATCH_SIZE * BLOCK_SIZE)
-print(f"grad_accumulation_steps: {grad_accumulation_steps}")
-print(f"total_batch_size: {total_batch_size}")
-
-import sys
-import os
-sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..")
-
-from DataLoader import create_dataloader
-
-
-
-def train(folder_path, tokenizer, model=None, optimizer=None, vocab_size=10000, platform='none', checkpoint=None, is_tokenized_data = False):
-
-    torch.set_float32_matmul_precision('high')  #hammad added this line (need to check if it is necessary)
-    if model is None:
-        model = GPTLanguageModel(vocab_size=vocab_size)
-        print("Model Initialised")
-        if checkpoint != None:
-            print("loading checkpoint........")
-            model.load(checkpoint)
-            print("Model loaded from checkpoint", checkpoint)
-
-        if platform == 'kaggle':
-            model = torch.nn.DataParallel(model, device_ids=[0, 1])
-            model = model.to(DEVICE)
-            optimizer = model.module.configure_optimizers(weight_decay=0.1, learning_rate=LEARNING_RATE, device=DEVICE) #hammad added this line
-        else:
-            model = model.to(DEVICE)
-            model = torch.compile(model) #hammad added this line
-            optimizer = model.configure_optimizers(weight_decay=0.1, learning_rate=LEARNING_RATE, device=DEVICE)
-        # optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, betas = (0.9, 0.95), eps = 1e-8)
-        
-
-    # # Initialize the data loader
-    # loader = TextDataLoader(file_path, BATCH_SIZE, BLOCK_SIZE, tokenizer)
-    
-    
-    # Set up a tqdm progress bar for the epoch
-    for epoch in range(MAX_ITERS):
-        print(f"Epoch {epoch}")
-        epoch_loss = None  # Track loss for the epoch
-
-
-        for i in range(len(os.listdir(folder_path))):
-            file_path = os.path.join(folder_path, os.listdir(folder_path)[i])
-            print(f"loading file: {file_path}")
-            loader = create_dataloader(tokenizer, file_path, BATCH_SIZE, BLOCK_SIZE, BLOCK_SIZE, tokenized_data = is_tokenized_data, filename = os.listdir(folder_path)[i]) #hammad added this line
-
-        
-            # Create a progress bar for batch processing
-            batch_progress_bar = tqdm(loader, desc=f"Epoch {epoch+1} Batch Progress", unit="batch", ncols=100)
-            count = 0
-            loss_accum = 0
-            for xb, yb in batch_progress_bar:
-                if xb is None:
-                    break  # No more batches, stop the epoch
-                optimizer.zero_grad()
-
-                    # Forward pass and loss computation
-                xb = xb.to(DEVICE)
-                yb = yb.to(DEVICE)
-                #with torch.autocast(DEVICE, dtype=torch.bfloat16): #hammad added this line
-                logits, loss = model(xb, yb)
-                loss = loss / grad_accumulation_steps
-                if platform == 'kaggle':
-                    loss_accum += loss.mean().detach()
-                    loss.mean().backward()
-                else:
-                    loss_accum += loss.detach()
-                    loss.backward()  # Backpropagate the loss
-                # for micro_batch in range(grad_accumulation_steps):
-                if count % grad_accumulation_steps == 0:
-                    print("one batch completed at (xb,yb):", count)
-                    loss_accum = 0
-                    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) #hammad added this line
-                    lr = get_lr(count) #need to check if this is correct
-                    for param_group in optimizer.param_groups:
-                        param_group['lr'] = lr
-                    optimizer.step()  # Update model parameters
-                    torch.cuda.synchronize() #wait for the computation to finish before moving to the next iteration
-                
-                # Update epoch_loss to the most recent loss value
-                if platform == 'kaggle':
-                    epoch_loss = loss.mean().item()
-                else:
-                    epoch_loss = loss.item()
-                
-                # Update tqdm with the latest loss value
-                batch_progress_bar.set_postfix(loss=epoch_loss)
-
-                count+=1
-                if count%5000 == 0:
-                    if platform == 'kaggle':
-                        torch.save(model.module.state_dict(), f"model_weights_checkpoint_{count}.pth")
-                    else:
-                        torch.save(model.state_dict(), f"model_weights_checkpoint_{count}.pth")
-                    print(f"Model weights saved at checkpoint {count}")
-            
-            # Save model weights after each chunk or epoch
-            if platform == 'kaggle':
-                torch.save(model.module.state_dict(),
-                        f"model_weights_epoch_{epoch}_{file_path[-6:-4]}.pth")
-            else:
-                torch.save(model.state_dict(),
-                        f"model_weights_epoch_{epoch}_{file_path[-6:-4]}.pth")
-            print(f"Model weights saved at epoch {epoch}")
-            
-            # Print the loss at the end of the epoch
-            if epoch_loss is not None:
-                print(f"Epoch {epoch}, Loss: {epoch_loss}")
-            else:
-                print(f"Epoch {epoch}, No data available for loss calculation.")
-        
-        # Reset the loader for a new epoch
-    #     loader.reset()
-    
-    # loader.close()  # Ensure the file is properly closed at the end
-    torch.cuda.empty_cache()
-
-    return model, optimizer
-
-
-#before parallelizing the model
-# def train(file_path, tokenizer, model=None, optimizer=None, vocab_size=10000, platform='none'):
-#     if model is None:
-#         model = GPTLanguageModel(vocab_size=vocab_size)
-#         if platform == 'kaggle':
-#             model = torch.nn.DataParallel(model, device_ids=[0, 1]).to(DEVICE)
-#         else:
-#             model = model.to(DEVICE)
-#         optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
-
-#     # Initialize the data loader
-#     loader = TextDataLoader(file_path, BATCH_SIZE, BLOCK_SIZE, tokenizer, DEVICE)
-    
-#     # Set up a tqdm progress bar for the epoch
-#     for epoch in range(MAX_ITERS):
-#         print(f"Epoch {epoch}")
-#         epoch_loss = None  # Track loss for the epoch
-        
-#         # Create a progress bar for batch processing
-#         batch_progress_bar = tqdm(loader, total=loader.num_batches(), desc=f"Epoch {epoch+1} Batch Progress", unit="batch", ncols=100)
-        
-#         for xb, yb in batch_progress_bar:
-#             if xb is None:
-#                 break  # No more batches, stop the epoch
-            
-#             # Forward pass and loss computation
-#             logits, loss = model(xb, yb)
-#             optimizer.zero_grad()
-#             loss.backward()  # Backpropagate the loss
-#             optimizer.step()  # Update model parameters
-            
-#             # Update epoch_loss to the most recent loss value
-#             epoch_loss = loss.item()
-            
-#             # Update tqdm with the latest loss value
-#             batch_progress_bar.set_postfix(loss=epoch_loss)
-        
-#         # Save model weights after each chunk or epoch
-#         model.save(f"model_weights_epoch_{epoch}.pth")
-#         print(f"Model weights saved at epoch {epoch}")
-        
-#         # Print the loss at the end of the epoch
-#         if epoch_loss is not None:
-#             print(f"Epoch {epoch}, Loss: {epoch_loss}")
-#         else:
-#             print(f"Epoch {epoch}, No data available for loss calculation.")
-        
-#         # Reset the loader for a new epoch
-#         loader.reset()
-    
-#     loader.close()  # Ensure the file is properly closed at the end
-
-#     return model, optimizer
-
-# def train(file_path, tokenizer, model = None, optimizer = None, vocab_size=10000, platform='none'):
-#     if model is None:
-#         model = GPTLanguageModel(vocab_size=vocab_size)
-#         if platform == 'kaggle':
-#             model = torch.nn.DataParallel(model, device_ids=[0, 1]).to(DEVICE)
-#         else:
-#             model = model.to(DEVICE)
-#         optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
-#     loader = TextDataLoader(file_path, BATCH_SIZE, BLOCK_SIZE, tokenizer, DEVICE)
-
-
-#     for epoch in range(MAX_ITERS):  # Iterate over the file chunks
-#         print(f"Epoch {epoch}")
-#         epoch_loss = None  # Track loss for the epoch
-#         while not loader.end_of_file:
-#             xb, yb = loader.get_batch()
-#             if xb is None:
-#                 break  # No more batches, stop the epoch
-
-#             # Forward pass and loss computation
-#             # print("This is xb", xb)
-#             # print("This is yb", yb)
-#             logits, loss = model(xb, yb)
-#             optimizer.zero_grad()
-#             loss.backward()   #2 gpus pe masla kr rraha (krna for n gpus hai)
-#             optimizer.step()
-            
-#             # Update epoch_loss to the most recent loss value
-#             epoch_loss = loss.item()
-
-#         # Save model weights after each chunk or epoch
-#         model.save(f"model_weights_epoch_{epoch}.pth")
-#         print(f"Model weights saved at epoch {epoch}")
-        
-#         # Print the loss only if it was computed
-#         if epoch_loss is not None:
-#             print(f"Epoch {epoch}, Loss: {epoch_loss}")
-#         else:
-#             print(f"Epoch {epoch}, No data available for loss calculation.")
-
-#         # Reset the loader for a new epoch
-#         loader.reset()
-
-#     loader.close()  # Ensure file is properly closed at the end
-
-#     return model, optimizer
-
-
-# def train(file_path, tokenizer, model=None, optimizer=None, vocab_size=10000):
-#     # Check if multiple GPUs are available
-#     device = DEVICE
-#     if model is None:
-#         if torch.cuda.is_available() and torch.cuda.device_count() > 1:
-#             print(f"Training on {torch.cuda.device_count()} GPUs")
-#             model = GPTLanguageModel(vocab_size=vocab_size).to(device)
-#             model = torch.nn.DataParallel(model, device_ids=[0, 1])  # Wrap the model for multi-GPU training
-#         else:
-#             print("Training on a single GPU or CPU.")
-            
-#             model = GPTLanguageModel(vocab_size=vocab_size).to(device)
-        
-#     if optimizer is None:
-#         optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
-    
-#     loader = TextDataLoader(file_path, BATCH_SIZE, BLOCK_SIZE, tokenizer, device)
-
-#     for epoch in range(MAX_ITERS):  # Iterate over the file chunks
-#         print(f"Epoch {epoch}")
-#         epoch_loss = None  # Track loss for the epoch
-
-#         xb, yb = loader.get_batch()
-#         if xb is None:
-#             break  # No more batches, stop the epoch
-
-#         # Forward pass and loss computation
-#         logits, loss = model(xb, yb)
-#         optimizer.zero_grad()
-#         loss.backward()
-#         optimizer.step()
-        
-#         # Update epoch_loss to the most recent loss value
-#         epoch_loss = loss.item()
-
-#         # Save model weights after each chunk or epoch
-#         model_to_save = model.module if isinstance(model, torch.nn.DataParallel) else model  # Get the underlying model if using DataParallel
-#         model_to_save.save(f"model_weights_epoch_{epoch}.pth")
-#         print(f"Model weights saved at epoch {epoch}")
-        
-#         # Print the loss only if it was computed
-#         if epoch_loss is not None:
-#             print(f"Epoch {epoch}, Loss: {epoch_loss}")
-#         else:
-#             print(f"Epoch {epoch}, No data available for loss calculation.")
-
-#         # Reset the loader for a new epoch
-#         loader.reset()
-
-#     loader.close()  # Ensure file is properly closed at the end
-
-#     return model, optimizer

Naive_gpt/train_autoshardload_ddp.py

@@ -1,307 +0,0 @@
-from DataLoader import create_dataloader, create_dataloader_ddp, GPTDatasetDDP
-import os
-import sys
-import torch
-from tqdm import tqdm
-from .config import *
-from .data_loader import TextDataLoader
-from .model import GPTLanguageModel
-import math
-from torch.distributed import init_process_group, destroy_process_group
-from torch.nn.parallel import DistributedDataParallel as DDP
-import torch.distributed as dist
-import time
-import gc
-
-max_lr = 6e-4
-min_lr = max_lr * 0.1
-
-ratio = 715/19073
-
-max_steps = 9124 # tokencount / (batchsize*blocksize) = 4Btokens / (524288) i.e 512 batchisze * 1024
-# warmup_steps = 72 # 19073 / 715 is the ratio
-
-# max_steps = train_loader.calculate_steps()
-warmup_steps = int(ratio*max_steps)
-
-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)
-
-
-log_dir = "log"
-os.makedirs(log_dir, exist_ok=True)
-log_file = os.path.join(log_dir, f"log.txt")
-with open(log_file, "w") as f:  # open for writing to clear the file
-    pass
-
-
-# set up DDP (distributed data parallel).
-# torchrun command sets the env variables RANK, LOCAL_RANK, and WORLD_SIZE
-ddp = int(os.environ.get('RANK', -1)) != -1  # is this a ddp run?
-if ddp:
-    # use of DDP atm demands CUDA, we set the device appropriately according to rank
-    assert torch.cuda.is_available(), "for now i think we need CUDA for DDP"
-    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)
-    # this process will do logging, checkpointing etc.
-    master_process = ddp_rank == 0
-else:
-    # vanilla, non-DDP run
-    ddp_rank = 0
-    ddp_local_rank = 0
-    ddp_world_size = 1
-    master_process = True
-    # attempt to autodetect device
-    device = "cpu"
-    if torch.cuda.is_available():
-        device = "cuda"
-    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-        device = "mps"
-    print(f"using device: {device}")
-
-torch.manual_seed(1337)
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(1337)
-
-print(f"ddp: {ddp}, rank: {ddp_rank}, local_rank: {ddp_local_rank}, world_size: {ddp_world_size}")
-print(f"master_process: {master_process}, device: {device}")
-# sys.exit(0)
-
-
-total_batch_size = 524288  # 524288 / 1024 = 512 is the batch size hamari terms mein
-assert total_batch_size % (
-    BATCH_SIZE * BLOCK_SIZE * ddp_world_size) == 0, "make sure total_batch_size is divisible by BATCH_SIZE * BLOCK_SIZE"
-grad_accumulation_steps = total_batch_size // (
-    BATCH_SIZE * BLOCK_SIZE * ddp_world_size)
-if master_process:
-    print(f"grad_accumulation_steps: {grad_accumulation_steps}")
-    print(f"total_batch_size: {total_batch_size}")
-
-sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..")
-
-
-def train(folder_path, tokenizer, model=None, optimizer=None, vocab_size=20000, platform='none', checkpoint=None, is_tokenized_data = False):
-
-    # hammad added this line (need to check if it is necessary)
-    torch.set_float32_matmul_precision('high')
-    if model is None:
-        model = GPTLanguageModel(vocab_size=vocab_size)
-        print("Model Initialised")
-        if checkpoint != None:
-            print("loading checkpoint........")
-            model.load(checkpoint)
-            print("Model loaded from checkpoint", checkpoint)
-
-        # if platform == 'kaggle':
-        #     model = torch.nn.DataParallel(model, device_ids=[0, 1])
-        #     model = model.to(DEVICE)
-        #     optimizer = model.module.configure_optimizers(
-        #         weight_decay=0.1, learning_rate=LEARNING_RATE, device=DEVICE)  # hammad added this line
-        # else:
-        model = model.to(DEVICE)
-        model = torch.compile(model)  # hammad added this line
-        if ddp:
-            model = DDP(model, device_ids=[ddp_local_rank])
-        raw_model = model.module if ddp else model
-        # optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, betas = (0.9, 0.95), eps = 1e-8)
-        optimizer = raw_model.configure_optimizers(
-            weight_decay=0.1, learning_rate=LEARNING_RATE, device=DEVICE)
-
-    # # Initialize the data loader
-    # loader = TextDataLoader(file_path, BATCH_SIZE, BLOCK_SIZE, tokenizer)
-    train_loader = GPTDatasetDDP(tokenizer, BATCH_SIZE, BLOCK_SIZE, BLOCK_SIZE, folder_path, ddp_rank, ddp_world_size, "train")
-    val_loader = GPTDatasetDDP(tokenizer, BATCH_SIZE, BLOCK_SIZE, BLOCK_SIZE, folder_path, ddp_rank, ddp_world_size, "test")
-    
-
-    # Set up a tqdm progress bar for the epoch
-    for epoch in range(MAX_ITERS):
-        if master_process:
-            print(f"Epoch {epoch}")
-
-
-        epoch_loss = None  # Track loss for the epoch
-        train_loader.set_epoch(epoch)
-        val_loader.set_epoch(epoch)
-        train_loader.reset()
-
-        # for i in range(len(os.listdir(folder_path))):
-        # file_path = os.path.join(folder_path, os.listdir(folder_path)[i])
-        # print(f"loading file: {file_path}")
-        # loader = create_dataloader_ddp(tokenizer, file_path, BATCH_SIZE, BLOCK_SIZE, BLOCK_SIZE,
-        #                                 tokenized_data=is_tokenized_data, process_rank=ddp_rank, num_process=ddp_world_size, filename=os.listdir(folder_path)[i])  # hammad added this line
-        
-        # Create a progress bar for batch processing
-        # batch_progress_bar = tqdm(
-        #     loader, desc=f"Epoch {epoch+1} Batch Progress", unit="batch", ncols=100)
-        count = 0
-        halt = False
-        # for xb, yb in batch_progress_bar:
-        while count < max_steps-8:
-            t0 = time.time()
-            loss_accum = 0
-            model.train()
-            optimizer.zero_grad()
-            
-
-            # Forward pass and loss computation
-            for micro_step in range(grad_accumulation_steps):
-                xb, yb = train_loader.next_batch()
-                if xb is None:
-                    halt = True
-                    if master_process:
-                        checkpoint = {
-                                'model': raw_model.state_dict(),
-                                'optimizer' : optimizer.state_dict(),
-                                'config': raw_model.config,
-                                'step': count,
-                                'val_loss': val_loss_accum.item() }
-                    # you might also want to add optimizer.state_dict() and
-                    # rng seeds etc., if you wanted to more exactly resume training
-                        torch.save(checkpoint, f"model_weights_checkpoint_{count}.pth")
-                        print(f"Model weights saved at checkpoint {count}")
-                        print(f"Epoch {epoch} completed")
-                    torch.cuda.synchronize()
-                    break  # No more batches, stop the epoch
-                xb = xb.to(DEVICE)
-                yb = yb.to(DEVICE)
-                if ddp:
-                    model.require_backward_grad_sync = (micro_step == grad_accumulation_steps - 1)
-                with torch.autocast(DEVICE, dtype=torch.float32):
-                    logits, loss = model(xb, yb)
-                loss = loss / grad_accumulation_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)
-            if halt == True:
-                break
-            # with torch.autocast(DEVICE, dtype=torch.bfloat16): #hammad added this line
-            # logits, loss = model(xb, yb)
-            # loss = loss / grad_accumulation_steps
-            # # if platform == 'kaggle':
-            # #     loss_accum += loss.mean().detach()
-            # #     loss.mean().backward()
-            # # else:
-            # loss_accum += loss.detach()
-            # if ddp:
-            #     model.require_backward_grad_sync = (
-            #         (count + 1) % grad_accumulation_steps == 0)
-            # loss.backward()  # Backpropagate the loss
-
-            # for micro_batch in range(grad_accumulation_steps):
-            if ddp:
-                dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG)
-            norm = torch.nn.utils.clip_grad_norm_(
-                model.parameters(), 1.0)  # hammad added this line
-            lr = get_lr(count)  # need to check if this is correct
-            for param_group in optimizer.param_groups:
-                param_group['lr'] = lr
-            optimizer.step()  # Update model parameters
-            # wait for the computation to finish before moving to the next iteration
-            torch.cuda.synchronize()
-            t1 = time.time()
-            dt = t1 - t0  # time difference in seconds
-            tokens_processed = train_loader.batch_size * \
-                train_loader.block_size * grad_accumulation_steps * ddp_world_size
-            tokens_per_sec = tokens_processed / dt
-            if master_process:
-                print(f"epoch {epoch:5d} step {count:5d} | loss_accum: {loss_accum.item():.6f} | loss_unacum: {loss} |lr {lr:.4e} | norm: {norm:.4f} | dt: {dt*1000:.2f}ms | tok/sec: {tokens_per_sec:.2f}")
-                with open(log_file, "a") as f:
-                    f.write(f"{epoch} {count} train {loss_accum.item():.6f}\n")
-
-            # Update epoch_loss to the most recent loss value
-            if platform == 'kaggle':
-                epoch_loss = loss.mean().item()
-            else:
-                epoch_loss = loss.item()
-
-            # Update tqdm with the latest loss value
-            # batch_progress_bar.set_postfix(loss=epoch_loss)
-
-            # if count % 5000 == 0:
-            #     torch.save(raw_model.state_dict(),
-            #                     f"model_weights_checkpoint_{count}.pth")
-            #     
-            
-            if count % 50 == 0:
-                model.eval()
-                val_loader.reset()
-
-                with torch.no_grad():
-                    val_loss_accum = 0.0
-                    val_loss_steps = 20
-                    for _ in range(val_loss_steps):
-                        xb, yb = val_loader.next_batch()
-                        xb = xb.to(DEVICE)
-                        yb = yb.to(DEVICE)
-                        with torch.autocast(DEVICE, dtype=torch.float32):
-                            logits, loss = model(xb, yb)
-                        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"{count} val {val_loss_accum.item():.4f}\n")
-                    if count > 0 and (count % 100 == 0):
-                        # optionally write model checkpoints
-                        checkpoint_path = os.path.join(log_dir, f"model_{count:05d}.pt")
-                        checkpoint = {
-                            'model': raw_model.state_dict(),
-                            'optimizer' : optimizer.state_dict(),
-                            'config': raw_model.config,
-                            'step': count,
-                            'val_loss': val_loss_accum.item()
-                        }
-                        # you might also want to add optimizer.state_dict() and
-                        # rng seeds etc., if you wanted to more exactly resume training
-                        torch.save(
-                            checkpoint, f"model_weights_checkpoint_{epoch}_{count}.pth")
-                        print(f"Model weights saved at checkpoint {count}")
-            count += 1
-
-        # Save model weights after each chunk or epoch
-        
-        # Print the loss at the end of the epoch
-        if master_process:
-            if epoch_loss is not None:
-                print(f"Epoch {epoch}, Loss: {epoch_loss}")
-            else:
-                print(f"Epoch {epoch}, No data available for loss calculation.")
-
-
-    train_loader.close()  # Ensure the file is properly closed at the end
-    val_loader.close()
-    if ddp:
-        destroy_process_group()
-        print("Process group destroyed")
-
-    # Delete variables, model, and optimizer
-    del model
-    del optimizer
-    torch.cuda.empty_cache()  # Clears cached memory but not allocated memory
-
-    # Force garbage collection
-    gc.collect()  # Ensures Python garbage collector releases unreferenced memory
-    torch.cuda.empty_cache()  # Clears freed memory from GPU cache
-
-    torch.cuda.ipc_collect()  # Helps in multi-GPU setups
-    return (0,0)
-
-    
-

Naive_gpt/train_ddp.py

@@ -1,206 +0,0 @@
-from DataLoader import create_dataloader, create_dataloader_ddp
-import os
-import sys
-import torch
-from tqdm import tqdm
-from .config import *
-from .data_loader import TextDataLoader
-from .model import GPTLanguageModel
-import math
-from torch.distributed import init_process_group, destroy_process_group
-from torch.nn.parallel import DistributedDataParallel as DDP
-import torch.distributed as dist
-
-max_lr = 6e-4
-min_lr = max_lr * 0.1
-warmup_steps = 10
-max_steps = 50
-
-
-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)
-
-
-# set up DDP (distributed data parallel).
-# torchrun command sets the env variables RANK, LOCAL_RANK, and WORLD_SIZE
-ddp = int(os.environ.get('RANK', -1)) != -1  # is this a ddp run?
-if ddp:
-    # use of DDP atm demands CUDA, we set the device appropriately according to rank
-    assert torch.cuda.is_available(), "for now i think we need CUDA for DDP"
-    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)
-    # this process will do logging, checkpointing etc.
-    master_process = ddp_rank == 0
-else:
-    # vanilla, non-DDP run
-    ddp_rank = 0
-    ddp_local_rank = 0
-    ddp_world_size = 1
-    master_process = True
-    # attempt to autodetect device
-    device = "cpu"
-    if torch.cuda.is_available():
-        device = "cuda"
-    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-        device = "mps"
-    print(f"using device: {device}")
-
-torch.manual_seed(1337)
-if torch.cuda.is_available():
-    torch.cuda.manual_seed(1337)
-
-print(f"ddp: {ddp}, rank: {ddp_rank}, local_rank: {ddp_local_rank}, world_size: {ddp_world_size}")
-print(f"master_process: {master_process}, device: {device}")
-import sys; sys.exit(0)
-
-
-total_batch_size = 524288  #524288 / 1024 = 512 is the batch size hamari terms mein
-assert total_batch_size % (
-    BATCH_SIZE * BLOCK_SIZE * ddp_world_size) == 0, "make sure total_batch_size is divisible by BATCH_SIZE * BLOCK_SIZE"
-grad_accumulation_steps = total_batch_size // (BATCH_SIZE * BLOCK_SIZE * ddp_world_size)
-if master_process:
-    print(f"grad_accumulation_steps: {grad_accumulation_steps}")
-    print(f"total_batch_size: {total_batch_size}")
-
-sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..")
-
-
-def train(folder_path, tokenizer, model=None, optimizer=None, vocab_size=10000, platform='none', checkpoint=None, is_tokenized_data=False):
-
-    # hammad added this line (need to check if it is necessary)
-    torch.set_float32_matmul_precision('high')
-    if model is None:
-        model = GPTLanguageModel(vocab_size=vocab_size)
-        print("Model Initialised")
-        if checkpoint != None:
-            print("loading checkpoint........")
-            model.load(checkpoint)
-            print("Model loaded from checkpoint", checkpoint)
-
-        # if platform == 'kaggle':
-        #     model = torch.nn.DataParallel(model, device_ids=[0, 1])
-        #     model = model.to(DEVICE)
-        #     optimizer = model.module.configure_optimizers(
-        #         weight_decay=0.1, learning_rate=LEARNING_RATE, device=DEVICE)  # hammad added this line
-        # else:
-        model = model.to(DEVICE)
-        model = torch.compile(model)  # hammad added this line
-        if ddp:
-            model = DDP(model, device_ids=[ddp_local_rank])
-        raw_model = model.module if ddp else model
-        # optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, betas = (0.9, 0.95), eps = 1e-8)
-        optimizer = raw_model.configure_optimizers(
-            weight_decay=0.1, learning_rate=LEARNING_RATE, device=DEVICE)
-
-    # # Initialize the data loader
-    # loader = TextDataLoader(file_path, BATCH_SIZE, BLOCK_SIZE, tokenizer)
-
-    # Set up a tqdm progress bar for the epoch
-    for epoch in range(MAX_ITERS):
-        print(f"Epoch {epoch}")
-        epoch_loss = None  # Track loss for the epoch
-
-        # for i in range(len(os.listdir(folder_path))):
-        # file_path = os.path.join(folder_path, os.listdir(folder_path)[i])
-        # print(f"loading file: {file_path}")
-        loader = create_dataloader_ddp(tokenizer, file_path, BATCH_SIZE, BLOCK_SIZE, BLOCK_SIZE,
-                                    tokenized_data=is_tokenized_data, process_rank = ddp_rank, num_process = ddp_world_size, filename=os.listdir(folder_path)[i])  # hammad added this line
-
-        # Create a progress bar for batch processing
-        batch_progress_bar = tqdm(
-            loader, desc=f"Epoch {epoch+1} Batch Progress", unit="batch", ncols=100)
-        count = 0
-        loss_accum = 0
-        for xb, yb in batch_progress_bar:
-            if xb is None:
-                break  # No more batches, stop the epoch
-            optimizer.zero_grad()
-
-            # Forward pass and loss computation
-            xb = xb.to(DEVICE)
-            yb = yb.to(DEVICE)
-            # with torch.autocast(DEVICE, dtype=torch.bfloat16): #hammad added this line
-            logits, loss = model(xb, yb)
-            loss = loss / grad_accumulation_steps
-            # if platform == 'kaggle':
-            #     loss_accum += loss.mean().detach()
-            #     loss.mean().backward()
-            # else:
-            loss_accum += loss.detach()
-            if ddp:
-                model.require_backward_grad_sync = ((count + 1) % grad_accumulation_steps == 0)
-            loss.backward()  # Backpropagate the loss
-
-            # for micro_batch in range(grad_accumulation_steps):
-            if count % grad_accumulation_steps == 0:
-                print("one batch completed at (xb,yb):", count)
-                if ddp:
-                    dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG)
-                loss_accum = 0
-                norm = torch.nn.utils.clip_grad_norm_(
-                    model.parameters(), 1.0)  # hammad added this line
-                lr = get_lr(count)  # need to check if this is correct
-                for param_group in optimizer.param_groups:
-                    param_group['lr'] = lr
-                optimizer.step()  # Update model parameters
-                # wait for the computation to finish before moving to the next iteration
-                torch.cuda.synchronize()
-
-            # Update epoch_loss to the most recent loss value
-            if platform == 'kaggle':
-                epoch_loss = loss.mean().item()
-            else:
-                epoch_loss = loss.item()
-
-            # Update tqdm with the latest loss value
-            batch_progress_bar.set_postfix(loss=epoch_loss)
-
-            count += 1
-            if count % 5000 == 0:
-                if platform == 'kaggle':
-                    torch.save(model.module.state_dict(),
-                                f"model_weights_checkpoint_{count}.pth")
-                else:
-                    torch.save(model.state_dict(),
-                                f"model_weights_checkpoint_{count}.pth")
-                print(f"Model weights saved at checkpoint {count}")
-
-        # Save model weights after each chunk or epoch
-        if platform == 'kaggle':
-            torch.save(model.module.state_dict(),
-                        f"model_weights_epoch_{epoch}_{file_path[-6:-4]}.pth")
-        else:
-            torch.save(model.state_dict(),
-                        f"model_weights_epoch_{epoch}_{file_path[-6:-4]}.pth")
-        print(f"Model weights saved at epoch {epoch}")
-
-        # Print the loss at the end of the epoch
-        if epoch_loss is not None:
-            print(f"Epoch {epoch}, Loss: {epoch_loss}")
-        else:
-            print(
-                f"Epoch {epoch}, No data available for loss calculation.")
-
-        # Reset the loader for a new epoch
-    #     loader.reset()
-
-    # loader.close()  # Ensure the file is properly closed at the end
-
-    return model, optimizer
-
-if ddp:
-    destroy_process_group()
-    print("Process group destroyed")