gradio app

SmolLm3.py

@@ -0,0 +1,236 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import SiLU
+import yaml
+# from gptdataloader import create_dataloader_v1
+# from chapter5 import calc_loss_loader, calculate_loss_batch
+
+
+def _init_weights(module, std=0.041666666666666664):
+    if isinstance(module, nn.Linear):
+        module.weight.data.normal_(mean=0.0, std=std)
+    elif isinstance(module, nn.Embedding):
+        module.weight.data.normal_(mean=0.0, std=std)
+
+class RotaryPositionalEmbedding(nn.Module):
+    """
+    # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L240
+    Rotary Positional Embedding (RoPE) for transformers Implemntation derived from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+    """
+    def __init__(self, dim: int, theta: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+
+    def forward(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
+        """
+        Apply rotary positional embedding to the input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape # B, T, H, D
+            seq_len (int): Sequence length. #T
+
+        Returns:
+            torch.Tensor: Output tensor with rotary positional embeddings applied.
+        """
+        B, T, H, H_D = x.shape
+
+        # Generate position indices
+        position = torch.arange(T, dtype=torch.float32, device=x.device).unsqueeze(-1)
+
+        # Generate frequencies
+        freqs = torch.exp(
+            torch.arange(0, H_D, 2, dtype=torch.float32, device=x.device) * 
+            -(torch.log(torch.tensor(self.theta)) / H_D)
+                                                                
+        )
+
+        # Compute sinusoids
+        sinusoid = position * freqs
+        sin = torch.sin(sinusoid)
+        cos = torch.cos(sinusoid)
+
+        # Reshape sin and cos to match the input tensor's shape
+        sin = sin.unsqueeze(0).unsqueeze(2)  # Shape: (1, T, 1, D // 2)
+        cos = cos.unsqueeze(0).unsqueeze(2)  # Shape: (1, T, 1, D // 2)
+
+        # Apply rotary embeddings
+        x_rotated = x.clone()
+        x_rotated[..., 0::2] = x[..., 0::2] * cos - x[..., 1::2] * sin
+        x_rotated[..., 1::2] = x[..., 1::2] * cos + x[..., 0::2] * sin
+
+        return x_rotated
+    
+class LlamaAttention(nn.Module):
+    """
+    (self_attn): LlamaAttention(
+          (q_proj): Linear(in_features=576, out_features=576, bias=False)
+          (k_proj): Linear(in_features=576, out_features=192, bias=False)
+          (v_proj): Linear(in_features=576, out_features=192, bias=False)
+          (o_proj): Linear(in_features=576, out_features=576, bias=False)
+    )
+    """
+    def __init__(self, config, rotary_emb):
+        super().__init__()
+        self.config = config
+        self.num_attention_heads = self.config['num_attention_heads']
+        self.hidden_size = self.config['hidden_size']
+        # Ensure the hidden size is divisible by the number of attention heads
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({self.hidden_size}) must be divisible by num_attention_heads ({self.num_attention_heads})"
+            )
+        self.num_key_value_heads = self.config['num_key_value_heads']
+        self.head_dim =  self.hidden_size // self.num_attention_heads
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)  # D,D
+        self.k_proj = nn.Linear(self.hidden_size, self.head_dim*self.num_key_value_heads, bias=False)   # D,D/H
+        self.v_proj = nn.Linear(self.hidden_size, self.head_dim*self.num_key_value_heads, bias=False)   # D,D/H
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)   # D,D
+
+        # Convert the mask to boolean type when creating it
+        # self.register_buffer("mask", 
+        #                    torch.triu(torch.ones(self.config['max_position_embeddings'], 
+        #                                        self.config['max_position_embeddings']),
+        #                             diagonal=1))  # Convert to boolean
+        
+        self.rotary_pos_emb = rotary_emb
+
+    def forward(self, x):
+        B, T, C = x.size()
+
+        q = self.q_proj(x)  # B,T,D
+        k = self.k_proj(x)  # B,T,D/H
+        v = self.v_proj(x)  # B,T,D/H
+
+        q = q.view(B, T, self.num_attention_heads, self.head_dim) # B,T,H,D
+        k = k.view(B, T, self.num_key_value_heads, self.head_dim) # B,T,H,D
+        v = v.view(B, T, self.num_key_value_heads, self.head_dim) # B,T,H,D
+
+        q = q.transpose(1,2) # B,H,T,D
+        k = k.transpose(1,2) # B,num_key_value_heads,T,D
+        v = v.transpose(1,2) # B,num_key_value_heads,T,D
+
+        # apply rotary positional embedding
+        q = self.rotary_pos_emb(q, T)
+        k = self.rotary_pos_emb(k, T)
+
+        # Repeat k/v heads if num_key_value_heads < num_attention_heads
+        if self.num_key_value_heads != self.num_attention_heads:
+            k = k.repeat_interleave(self.num_attention_heads // self.num_key_value_heads, dim=1) # B,kv_head,T,D -> B,H,T,D
+            v = v.repeat_interleave(self.num_attention_heads // self.num_key_value_heads, dim=1) # B,kv_head,T,D -> B,H,T,D
+
+        # Manual attention Stats
+        # Q(B,H,T,D) @K.T(B,H,D,T) = Q.K_T (B,H,T,T)
+        # attn_scores = q @ k.transpose(-2,-1) # B,H,T,T
+        # mask_bool = self.mask[:T,:T].bool() # T,T
+        # attn_scores.masked_fill_(mask_bool, -torch.inf) # B,H,T,T
+        # attn_weights = F.softmax(attn_scores/k.size(-1)**0.5, dim=-1) # B,H,T,T
+        # context_vector = attn_weights @ v # B,H,T,T * B,H,T,D = B,H,T,D
+        # context_vector = context_vector.transpose(1,2) # B,T,H,D
+        # context_vector = context_vector.contiguous().view(B,T,C) # B,T,H,D -> B,T,D
+        # Manual attention Stats ENDS
+
+        # Scaled dot-product attention STARTS   
+        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        context_vector = attn_out.transpose(1,2).reshape(B,T,C)
+        # Scaled dot-product attention ENDS
+
+        context_vector = self.o_proj(context_vector)
+        
+        return context_vector
+
+
+class LlamaMLP(nn.Module):
+    """
+    (mlp): LlamaMLP(
+          (gate_proj): Linear(in_features=576, out_features=1536, bias=False)
+          (up_proj): Linear(in_features=576, out_features=1536, bias=False)
+          (down_proj): Linear(in_features=1536, out_features=576, bias=False)
+          (act_fn): SiLU()
+        )
+    """
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.gate_proj = nn.Linear(self.config['hidden_size'], self.config['intermediate_size'], bias=False)
+        self.up_proj = nn.Linear(self.config['hidden_size'], self.config['intermediate_size'], bias=False)
+        self.down_proj = nn.Linear(self.config['intermediate_size'], self.config['hidden_size'], bias=False)
+        self.act_fn = SiLU()
+    def forward(self, x):
+        gate = self.gate_proj(x)
+        up = self.up_proj(x)
+        down = self.down_proj(self.act_fn(gate)*up)
+        return down 
+    
+class LlamaRMSNorm(nn.Module):
+    """
+    (norm): LlamaRMSNorm((576,), eps=1e-05)
+        # RMSNorm Formula:
+        #    RMS(x) = sqrt((1 / d) * sum(x_i^2 for i in range(d)))
+        #    x_normalized = x / RMS(x)
+        #    output = gamma * x_normalized
+    
+    """
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.eps = self.config['rms_norm_eps']
+        self.weight = nn.Parameter(torch.ones(self.config['hidden_size']))
+    def forward(self, x):
+        rms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
+        return  self.weight *rms * x
+    
+class LlamaDecoderLayer(nn.Module):
+    def __init__(self, config, rotary_emb):
+        super().__init__()
+        self.config = config
+        self.self_attn = LlamaAttention(self.config, rotary_emb)
+        self.mlp = LlamaMLP(self.config)
+        self.input_layernorm = LlamaRMSNorm(self.config)
+        self.post_attention_layernorm = LlamaRMSNorm(self.config)   
+    
+    def forward(self, x):
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x)
+        x = x + residual
+
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = x + residual
+        return x 
+        # # x = x + self.self_attn(self.input_layernorm(x))
+        # # x = x + self.mlp(self.post_attention_layernorm(x))
+        # return x
+class LlamaModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.init_method = config['init_method']
+        self.config = config['model_config']
+        self.embed_tokens = nn.Embedding(self.config['vocab_size'], self.config['hidden_size'])
+        self.rotary_emb = RotaryPositionalEmbedding(self.config['hidden_size'], self.config['rope_theta'])
+        self.layers = nn.ModuleList([LlamaDecoderLayer(self.config, self.rotary_emb) for _ in range(self.config['num_hidden_layers'])])
+        self.norm = LlamaRMSNorm(self.config)
+        self.lm_head = nn.Linear(self.config['hidden_size'], self.config['vocab_size'], bias=False)
+        
+        if self.config['tie_word_embeddings']:
+            self.lm_head.weight = self.embed_tokens.weight
+        
+        self.apply(lambda m: _init_weights(m, self.init_method['std']))
+    
+    def forward(self, x, y=None):
+        x = self.embed_tokens(x)
+        for layer in self.layers:
+            x = layer(x)
+        x = self.norm(x)
+        logits = self.lm_head(x) # B,T,V
+        logits = logits.view(-1, logits.size(-1))  # Shape: [B*T, V]
+        if y is not None:
+            y = y.view(-1)  # Shape: [B*T]
+            loss = torch.nn.functional.cross_entropy(logits, y)
+            return logits, loss
+        else:
+            return logits, None
+

app.py

@@ -0,0 +1,148 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer
+import yaml
+from SmolLm3 import LlamaModel
+
+
+def generate_helper(model, idx, max_new_tokens, context_length, temperature=1.0, top_k=None, eos_token=None, device=None):
+    
+    model = model.to(device)
+    idx = idx.to(device)
+    model.eval()
+    for _ in range(max_new_tokens):
+        idx_cond = idx[:, -context_length:]
+        with torch.no_grad():
+            logits, _ = model(idx_cond)  # Unpack both logits and loss (ignore loss)
+            logits = logits.view(idx_cond.shape[0], -1, model.config['vocab_size'])  # Reshape to [batch, seq, vocab]
+            
+        # Get the logits for the last token only
+        logits = logits[:, -1, :]  # Shape: [batch_size, vocab_size]
+        
+        if top_k is not None:
+            # top k sampling
+            top_logits, top_pos = torch.topk(logits, top_k)
+            min_logit = top_logits[:, -1].unsqueeze(-1)
+            logits = torch.where(logits < min_logit,
+                               torch.tensor(float('-inf')).to(logits.device),
+                               logits)
+        
+        # temperature scaling
+        if temperature > 0.0:
+            logits /= temperature
+            probs = torch.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+        else:
+            idx_next = torch.argmax(logits, dim=-1, keepdim=True)
+            
+        if idx_next.item() == eos_token:
+            break
+            
+        idx = torch.cat((idx, idx_next), dim=1)
+    model.train()
+    return idx
+
+def get_config(config_path):
+    config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader)
+    return config
+
+def load_model_from_checkpoint(config_path, checkpoint_path, device):
+    config = get_config(config_path)
+    model = LlamaModel(config['model'])
+    checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
+    state_dict = checkpoint['model_state_dict']
+    state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
+    model.load_state_dict(state_dict)
+    return model
+
+def load_weights(config, weights_path, device):
+    model = LlamaModel(config['model'])
+    model.load_state_dict(torch.load(weights_path, map_location=torch.device(device)))
+    return model
+
+def get_tokenizer(config):
+    tokenizer_path = config['tokenizer']['tokenizer_name_or_path']
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+    tokenizer.pad_token = tokenizer.eos_token
+    vocab_size = tokenizer.vocab_size
+    return tokenizer, vocab_size
+
+def generate_text(model, tokenizer, input_text, max_new_tokens, context_length, temperature, top_k, eos_token, device):
+    encoded_text = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    generated_text = generate_helper(model, 
+                            idx=encoded_text,
+                            max_new_tokens=max_new_tokens,
+                            context_length=context_length, 
+                            temperature=temperature, 
+                            top_k=top_k, 
+                            eos_token=eos_token, 
+                            device=device)
+    return tokenizer.decode(generated_text.squeeze(0))
+
+
+
+# Initialize model and tokenizer
+def initialize_model():
+    config_path = "config_smollm2_135M.yaml"
+    checkpoint_path = "/Users/chiragtagadiya/Documents/Final_training_before_stop_smolllm3/checkpoints/model_37000_steps_avg_loss_2.85920_optimizer_lr_0.00000003.pth"  # Update this path
+    weights_path = "model_weights_35000_step.pt"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Load configuration
+    config = get_config(config_path)
+    
+    # Load model
+    # model = load_model_from_checkpoint(config_path, checkpoint_path, device)
+    model = load_weights(config, weights_path, device)
+    model.to(device)
+    model.eval()
+    
+    # Load tokenizer
+    tokenizer, vocab_size = get_tokenizer(config)
+    
+    return model, tokenizer, device
+
+def generate_response(prompt, max_new_tokens):
+    generated_text = generate_text(
+        model=model,
+        tokenizer=tokenizer,
+        input_text=prompt,
+        max_new_tokens=max_new_tokens,
+        context_length=256,
+        temperature=0.9,
+        top_k=2,
+        eos_token=tokenizer.eos_token_id,
+        device=device
+    )
+    return generated_text
+
+# Initialize global variables
+model, tokenizer, device = initialize_model()
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=[
+        gr.Textbox(
+            lines=3, 
+            placeholder="Enter your prompt here...",
+            label="Input Prompt"
+        ),
+        gr.Slider(
+            minimum=50,
+            maximum=256,
+            value=100,
+            step=10,
+            label="Max New Tokens"
+        )
+    ],
+    outputs=gr.Textbox(
+        lines=5,
+        label="Generated Text"
+    ),
+    title="SmolLM Text Generator",
+    description="Enter a prompt and adjust the maximum number of tokens to generate text with SmolLM model."
+)
+
+if __name__ == "__main__":
+    iface.launch()

config_smollm2_135M.yaml

@@ -0,0 +1,103 @@
+checkpoints:
+  checkpoint_interval: 2000
+  checkpoints_path: checkpoints
+  checkpoints_path_is_shared_file_system: false
+  resume_checkpoint_path: null
+  save_final_state: false
+  save_initial_state: false
+data_stages:
+- data:
+    dataset:
+      dataset_folder:
+      - datasets/smollm2-corpus
+      dataset_weights:
+      - 1.0
+    num_loading_workers: 0
+    seed: 8
+  name: stable phase
+  start_training_step: 1
+general:
+  benchmark_csv_path: null
+  consumed_train_samples: null
+  ignore_sanity_checks: true
+  project: smollm2
+  run: smollm2-135M
+  seed: 8
+  step: null
+logging:
+  iteration_step_info_interval: 1
+  log_level: info
+  log_level_replica: info
+model:
+  ddp_bucket_cap_mb: 25
+  dtype: bfloat16
+  init_method:
+    std: 0.041666666666666664
+  make_vocab_size_divisible_by: 1
+  model_config:
+    bos_token_id: 0
+    eos_token_id: 0
+    hidden_act: silu
+    hidden_size: 576
+    initializer_range: 0.041666666666666664
+    intermediate_size: 1536
+    is_llama_config: true
+    max_position_embeddings: 2048
+    num_attention_heads: 9
+    num_hidden_layers: 30
+    num_key_value_heads: 3
+    pad_token_id: null
+    pretraining_tp: 1
+    rms_norm_eps: 1.0e-05
+    rope_interleaved: false
+    rope_scaling: null
+    rope_theta: 10000.0
+    tie_word_embeddings: true
+    use_cache: true
+    vocab_size: 49152
+    s3_bucket: smollm2-train-jan-25-era3
+    s3_checkpoint_folder: checkpoints
+    s3_log_folder: logs
+    s3_log_file_name: training.log
+optimizer:
+  accumulate_grad_in_fp32: true
+  clip_grad: 1.0
+  learning_rate_scheduler:
+    learning_rate: 0.003
+    lr_decay_starting_step: 1600000
+    lr_decay_steps: 400000
+    lr_decay_style: linear
+    lr_warmup_steps: 2000
+    lr_warmup_style: linear
+    min_decay_lr: 0
+  optimizer_factory:
+    adam_beta1: 0.9
+    adam_beta2: 0.95
+    adam_eps: 1.0e-08
+    name: adamW
+    torch_adam_is_fused: true
+  weight_decay: 0.01
+  zero_stage: 0
+parallelism:
+  dp: 64
+  expert_parallel_size: 1
+  pp: 1
+  pp_engine: 1f1b
+  recompute_layer: false
+  tp: 1
+  tp_linear_async_communication: true
+  tp_mode: REDUCE_SCATTER
+  tp_recompute_allgather: true
+profiler: null
+tokenizer:
+  tokenizer_max_length: null
+  tokenizer_name_or_path: HuggingFaceTB/cosmo2-tokenizer
+  tokenizer_revision: null
+tokens:
+  batch_accumulation_per_replica: 1
+  limit_test_batches: 0
+  limit_val_batches: 0
+  micro_batch_size: 16 #16
+  sequence_length: 1024 #2048
+  train_steps: 2000000
+  val_check_interval: 1000

model_testing.py

@@ -0,0 +1,79 @@
+import argparse
+from SmolLm3 import LlamaModel
+import yaml
+import torch
+from transformers import AutoTokenizer
+from train import generate
+
+def get_config(config_path):
+    config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader)
+    return config
+
+def load_model_from_checkpoint(config_path, checkpoint_path, device):
+    config = get_config(config_path)
+    model = LlamaModel(config['model'])
+    checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
+    state_dict = checkpoint['model_state_dict']
+    state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
+    model.load_state_dict(state_dict)
+    return model
+
+def get_tokenizer(config):
+    tokenizer_path = config['tokenizer']['tokenizer_name_or_path']
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+    tokenizer.pad_token = tokenizer.eos_token
+    vocab_size = tokenizer.vocab_size
+    return tokenizer, vocab_size
+
+def generate_text(model, tokenizer, input_text, max_new_tokens, context_length, temperature, top_k, eos_token, device):
+    encoded_text = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    generated_text = generate(model, 
+                            idx=encoded_text,
+                            max_new_tokens=max_new_tokens,
+                            context_length=context_length, 
+                            temperature=temperature, 
+                            top_k=top_k, 
+                            eos_token=eos_token, 
+                            device=device)
+    return tokenizer.decode(generated_text.squeeze(0))
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Generate text using the SmolLM model')
+    parser.add_argument('--config_path', type=str, default="config_smollm2_135M.yaml",
+                        help='Path to the config file')
+    parser.add_argument('--checkpoint_path', type=str, required=True,
+                        help='Path to the model checkpoint')
+    parser.add_argument('--input_text', type=str, default="Bernuli principle",
+                        help='Input text prompt for generation')
+    parser.add_argument('--max_new_tokens', type=int, default=256,
+                        help='Maximum number of new tokens to generate')
+    parser.add_argument('--context_length', type=int, default=256,
+                        help='Context length for generation')
+    parser.add_argument('--temperature', type=float, default=0.7,
+                        help='Temperature for sampling')
+    parser.add_argument('--top_k', type=int, default=5,
+                        help='Top-k value for sampling')
+    parser.add_argument('--device', type=str, default="cuda" if torch.cuda.is_available() else "cpu",
+                        help='Device to run the model on (cuda/cpu)')
+    
+    args = parser.parse_args()
+
+    config = get_config(args.config_path)
+    model = load_model_from_checkpoint(args.config_path, args.checkpoint_path, args.device)
+    print(model)
+    tokenizer, vocab_size = get_tokenizer(config)
+    print(tokenizer)
+    print(vocab_size)
+
+    generated_text = generate_text(
+        model, 
+        tokenizer, 
+        args.input_text,
+        args.max_new_tokens,
+        args.context_length,
+        args.temperature,
+        args.top_k,
+        tokenizer.eos_token_id,
+        args.device
+    )
+    print(generated_text)

model_weights_35000_step.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a965c902af30b6148a95d2d404b6848829a94bc4815fd53d2a84be51707e7df
+size 538169702

requirements.txt

@@ -0,0 +1,12 @@
+torch
+torchtext
+pandas
+numpy==1.26.1
+matplotlib
+tqdm
+# urllib
+requests
+boto3
+datasets
+transformers
+gradio

train.py

@@ -0,0 +1,351 @@
+from SmolLm3 import LlamaModel
+import torch
+import yaml
+from transformers import AutoTokenizer
+from torch.utils.data import DataLoader
+import numpy as np
+from datasets import load_dataset
+import logging
+import math
+
+from utils import upload_file_to_s3
+# At the start of training loop
+# print(f"GPU Memory allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
+# print(f"GPU Memory reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
+
+
+logger = logging.getLogger(__name__)
+formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+file_handler = logging.FileHandler('training.log')
+file_handler.setFormatter(formatter)  # Set formatter on the handler, not the logger
+logger.addHandler(file_handler)
+logger.setLevel(logging.INFO)
+
+def encode_text(examples, tokenizer, seq_length):
+    """Tokenize and prepare text examples for training."""
+    tokens = tokenizer(
+        examples["text"],
+        truncation=True,
+        padding="max_length",
+        max_length=seq_length + 1,
+        return_tensors="pt",
+    )
+    # Use clone().detach() as recommended
+    input_ids = tokens["input_ids"].squeeze(0).clone().detach()
+    input_ids = torch.clamp(input_ids, min=0, max=tokenizer.vocab_size - 1)
+    labels = input_ids.clone().detach()
+    labels = labels[1:].to(torch.int64)
+    input_ids = input_ids[:-1].to(torch.int64)
+
+    return {"input_ids": input_ids, "labels": labels}
+
+def load_cosmopedia_dataset(batch_size=8, seq_length=1024, tokenizer=None):
+    """
+    Returns a torch dataloader for the cosmopedia dataset
+    """
+    # Set tokenizer parallelism explicitly
+    import os
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    logger.info("tokenizer parallelism set to false")
+    try:
+        # Increase timeout and retries for dataset loading
+        from datasets import config
+        config.HF_DATASETS_TIMEOUT = 300  # 5 minutes timeout
+        config.MAX_RETRIES = 10  # Increase retry attempts
+        logger.info("dataset loading config set")
+        train_dataset = load_dataset(
+            "HuggingFaceTB/smollm-corpus",
+            name="cosmopedia-v2",
+            split="train",
+            streaming=True,
+        )
+        logger.info("dataset loaded")
+
+        # Use partial to bind tokenizer and seq_length to the encode function
+        from functools import partial
+        encode_fn = partial(encode_text, tokenizer=tokenizer, seq_length=seq_length)
+        
+        train_dataset = train_dataset.map(
+            encode_fn, 
+            remove_columns=["text"], 
+            batched=False
+        )
+        train_dataset = train_dataset.with_format("torch")
+        
+        train_dataloader = DataLoader(
+            train_dataset, 
+            batch_size=batch_size,
+            num_workers=2,
+            pin_memory=True,
+            prefetch_factor=4,
+            persistent_workers=True
+        )
+        return train_dataloader
+    except Exception as e:
+        logger.error(f"Error loading dataset: {str(e)}")
+        return None
+    
+
+def generate(model, idx, max_new_tokens, context_length, temperature=1.0, top_k=None, eos_token=None, device=None):
+    logger.info(f"Generating on device {device}")
+    model = model.to(device)
+    idx = idx.to(device)
+    model.eval()
+    for _ in range(max_new_tokens):
+        idx_cond = idx[:, -context_length:]
+        with torch.no_grad():
+            logits, _ = model(idx_cond)  # Unpack both logits and loss (ignore loss)
+            logits = logits.view(idx_cond.shape[0], -1, model.config['vocab_size'])  # Reshape to [batch, seq, vocab]
+            
+        # Get the logits for the last token only
+        logits = logits[:, -1, :]  # Shape: [batch_size, vocab_size]
+        
+        if top_k is not None:
+            # top k sampling
+            top_logits, top_pos = torch.topk(logits, top_k)
+            min_logit = top_logits[:, -1].unsqueeze(-1)
+            logits = torch.where(logits < min_logit,
+                               torch.tensor(float('-inf')).to(logits.device),
+                               logits)
+        
+        # temperature scaling
+        if temperature > 0.0:
+            logits /= temperature
+            probs = torch.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+        else:
+            idx_next = torch.argmax(logits, dim=-1, keepdim=True)
+            
+        if idx_next.item() == eos_token:
+            break
+            
+        idx = torch.cat((idx, idx_next), dim=1)
+    model.train()
+    return idx
+
+def sync_device(device):
+    if device.startswith('cuda'):
+        torch.cuda.synchronize()
+    elif device == 'cpu':
+        torch.cpu.synchronize() if hasattr(torch.cpu, 'synchronize') else None
+    elif device.startswith('mps'):  # For Apple Silicon
+        torch.mps.synchronize()
+
+def print_gpu_memory(step_name=""):
+    """
+    Print GPU memory statistics with a specified step name
+    """
+    if torch.cuda.is_available():
+        logger.info(f"\nGPU Memory Stats {step_name}:")
+        logger.info(f"GPU Memory allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
+        logger.info(f"GPU Memory reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
+        logger.info(f"Max GPU Memory allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+
+# Learning rate scheduler
+def get_lr_lambda(current_step, warmup_steps, max_steps, max_lr):
+    """
+    Modified learning rate scheduler with:
+    1. Linear warmup for first 3000 steps
+    2. Cosine decay from 3000 to 60000 steps
+    3. Minimum learning rate of 1.5e-5 (5% of max_lr)
+    """
+    min_lr = max_lr * 0.05  # Minimum learning rate (5% of max_lr)
+
+    if current_step < warmup_steps:
+        # Linear warmup from 0 to max_lr
+        return float(current_step) / float(max(1, warmup_steps))
+    else:
+        # Cosine decay from max_lr to min_lr
+        progress = float(current_step - warmup_steps) / float(max(1, max_steps - warmup_steps))
+        return min_lr + 0.5 * (max_lr - min_lr) * (1.0 + math.cos(math.pi * progress))
+
+
+def train_model(config, model, train_loader, test_loader, optimizer, device, num_epochs, eval_freq, eval_iter, start_context="Jack Gisburn rather a cheap genius- ", tokenizer=None):
+    total_loss = 0
+    tokens_seen, global_step = 0, -1
+    
+    # Adjusted gradient accumulation setup
+    actual_batch_size = config['tokens']['micro_batch_size']  # Now 16
+    effective_batch_size_multiplier = 2  # Reduced from 4 to maintain reasonable memory usage
+    target_batch_size = effective_batch_size_multiplier * config['tokens']['micro_batch_size']
+    gradient_accumulation_steps = target_batch_size // actual_batch_size
+    
+    # Adjusted learning rate parameters for new batch size
+    max_lr = 3e-4  # Keep the same max learning rate
+    warmup_steps = 3000  # Increase warmup steps for longer training
+    max_steps = 60000  # Set to match 10 hours of training
+    min_lr = max_lr * 0.05  # Reduce minimum LR to 5% of max (was 10%)
+    
+    # Create LambdaLR scheduler with the improved lambda function
+    lr_lambda = lambda step: get_lr_lambda(step, warmup_steps, max_steps, max_lr)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    
+    logger.info(f"Training with learning rate schedule:")
+    logger.info(f"Max LR: {max_lr}")
+    logger.info(f"Warmup Steps: {warmup_steps}")
+    logger.info(f"Max Steps: {max_steps}")
+    logger.info(f"Min LR: {max_lr * 0.05}")
+    logger.info(f"Gradient Accumulation Steps: {gradient_accumulation_steps}")
+    logger.info(f"Effective Batch Size: {actual_batch_size * gradient_accumulation_steps}")
+    
+    print_gpu_memory("at start of training")
+    
+    # Add these near the start of training loop
+    torch.cuda.empty_cache()
+    torch.backends.cudnn.benchmark = True
+    
+    for epoch in range(num_epochs):
+        model.train()
+        optimizer.zero_grad()  # Zero gradients at start of epoch
+        
+        for batch_idx, batch in enumerate(train_loader):
+            input_batch = batch['input_ids'].to(device)
+            target_batch = batch['labels'].to(device)
+            
+            # Forward pass
+            with torch.autocast(device_type=device, dtype=torch.bfloat16):
+                logits, original_loss = model(input_batch, target_batch)
+            
+                # Scale loss for gradient accumulation
+            scaled_loss = original_loss / gradient_accumulation_steps
+            scaled_loss.backward()
+            
+            # Add the original loss to total_loss for logging
+            total_loss += original_loss.item()  # Don't multiply back up
+            tokens_seen += input_batch.numel()
+            
+            # Calculate running average loss
+            total_batches = batch_idx + 1
+            avg_loss = total_loss / total_batches
+            if batch_idx % 25 == 0:
+                logger.info(f"Batch {batch_idx + 1}, Running Avg Loss: {avg_loss:.5f}")
+            # Only update weights after accumulating gradients
+            if (batch_idx + 1) % gradient_accumulation_steps == 0:
+                # Gradient clipping
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                
+                optimizer.step()
+                scheduler.step()  # Update learning rate
+                optimizer.zero_grad()
+                global_step += 1
+            
+            # Evaluation block
+            if global_step % eval_freq == 0 and global_step > 0:
+                # Use total batches processed instead of global_step
+                current_lr = scheduler.get_last_lr()[0]
+                optimizer_lr = optimizer.param_groups[0]['lr']
+                
+                print_gpu_memory(f"at step {global_step}")
+                logger.info(f"learning rate: {current_lr:.8f}")
+                logger.info(f"Ep {epoch+1} (Step {global_step:06d}): "
+                      f"Avg loss {avg_loss:.3f} | {tokens_seen} tokens seen")
+                logger.info(f"optimizer lr: {optimizer_lr:.8f}")
+                logger.info(f"scheduler lr: {current_lr:.8f}")
+                
+                # Generate sample text
+                encoded_text = tokenizer.encode(start_context, return_tensors="pt")
+                random_topk = np.random.randint(1, 10)
+                logger.info(f"random_topk: {random_topk}")
+                random_temperature = np.random.uniform(0.7, 0.9)
+                logger.info(f"random_temperature: {random_temperature}")
+                logger.info(f"global step {global_step} , batch_idx {batch_idx} => generating text")
+                generated_text = generate(model, 
+                                       idx=encoded_text,
+                                       max_new_tokens=256,
+                                       context_length=256, 
+                                       temperature=random_temperature, 
+                                       top_k=random_topk, 
+                                       eos_token=tokenizer.eos_token_id, 
+                                       device=device)
+                logger.info(f"+++"*30)
+                logger.info(tokenizer.decode(generated_text.squeeze(0)))
+                logger.info(f"+++"*30)
+                
+                # Save checkpoint
+                model_file_name = f"model_{global_step}_steps_avg_loss_{avg_loss:.5f}_optimizer_lr_{optimizer_lr:.8f}.pth"
+                torch.save({
+                    'step': global_step,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'scheduler_state_dict': scheduler.state_dict(),
+                    'loss': avg_loss,
+                }, model_file_name)
+                
+                s3_path = upload_file_to_s3(model_file_name, config['model']['model_config']['s3_bucket'], 
+                                          config['model']['model_config']['s3_checkpoint_folder'])
+                logger.info(f"Model saved to S3: {s3_path}")
+
+                log_path = upload_file_to_s3(config['model']['model_config']['s3_log_file_name'], config['model']['model_config']['s3_bucket'], 
+                                              config['model']['model_config']['s3_log_folder'])
+                logger.info(f"Log saved to S3: {log_path}")
+            
+            if batch_idx % 100 == 0:
+                logger.info(f"Batch {batch_idx} finished")
+                logger.info(f"+++"*30)
+
+    logger.info("Training complete")
+
+if __name__ == "__main__":
+    config = yaml.load(open("config_smollm2_135M.yaml", "r"), Loader=yaml.FullLoader)
+    logger.info(config)
+    
+    # Set memory efficient settings
+    torch.set_float32_matmul_precision('high')
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+    
+    # Empty cache before model creation
+    torch.cuda.empty_cache()
+    
+    model = LlamaModel(config['model'])
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    # Enable gradient checkpointing for memory efficiency
+    # model.gradient_checkpointing_enable()
+    
+    model.to(device)
+    model = torch.compile(model)
+    logger.info(model)
+    logger.info("++"*30)
+    
+    optimizer = torch.optim.AdamW(
+        model.parameters(), 
+        lr=3e-4, 
+        weight_decay=0.15,
+        betas=(0.9, 0.95)
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
+    tokenizer.pad_token = tokenizer.eos_token
+    vocab_size = tokenizer.vocab_size
+    
+    # Adjusted batch size and sequence length
+    train_loader = load_cosmopedia_dataset(
+        batch_size=16,  # Set to 16
+        seq_length=1024,  # Kept at 1024
+        tokenizer=tokenizer
+    )
+    
+    import time
+    t1 = time.time()
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    # Set environment variable for memory allocation
+    import os
+    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+    
+    train_model(
+        config, 
+        model, 
+        train_loader, 
+        train_loader, 
+        optimizer=optimizer, 
+        device=device, 
+        num_epochs=1, 
+        eval_freq=1000,  # Increase eval frequency to every 500 steps
+        eval_iter=1000,
+        start_context="Once Upon a Time far far away in a galaxy", 
+        tokenizer=tokenizer
+    )
+    t2 = time.time()
+    logger.info(f"Time taken for training: {t2 - t1:.2f} seconds")