gpt

Naive_gpt/config.py

@@ -0,0 +1,16 @@
+# 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

@@ -0,0 +1,16 @@
+# 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

@@ -0,0 +1,16 @@
+# 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

@@ -0,0 +1,16 @@
+# 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

@@ -0,0 +1,252 @@
+# 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

@@ -0,0 +1,258 @@
+# 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

@@ -0,0 +1,303 @@
+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

@@ -0,0 +1,307 @@
+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

@@ -0,0 +1,206 @@
+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")