Upload 2 files

code/model.py

@@ -7,13 +7,14 @@ from dataclasses import dataclass
 
 @dataclass
 class ModelConfig:
-    """Configuration class for ViuResonance100M."""
+    """Configuration for ViuResonance100M - T4 Optimized"""
     vocab_size: int = 64000
     d_model: int = 1024
     n_layers: int = 12
     n_heads: int = 16
     dropout: float = 0.1
-    chunk_size: int = 512
+    chunk_size: int = 256  # T4 ke liye 256, 512 se OOM aayega
+    max_seq_len: int = 4096
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
@@ -22,9 +23,12 @@ class RMSNorm(nn.Module):
         self.eps = eps
         
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # Enforce float32 accumulation for stability, then cast back to original dtype
-        norm_x = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) * self.weight.float()
-        return norm_x.to(x.dtype)
+        # T4 fp16 me stable rakhne ke liye float32 me norm
+        # Original se same logic, bas efficient
+        return F.rms_norm(x.float(), (x.shape[-1],), self.weight.float(), self.eps).to(x.dtype)
+        # Fallback agar F.rms_norm na ho:
+        # norm_x = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        # return (norm_x * self.weight.float()).to(x.dtype)
 
 class ResonanceLayerKaggle(nn.Module):
     def __init__(self, config: ModelConfig):
@@ -42,6 +46,9 @@ class ResonanceLayerKaggle(nn.Module):
         self.to_v = nn.Linear(config.d_model, config.d_model, bias=False)
         self.to_out = nn.Linear(config.d_model, config.d_model, bias=False)
         self.norm = RMSNorm(config.d_model)
+        
+        # T4: pos cache - har forward me arange banane se bachao
+        self.register_buffer("pos_cache", torch.arange(config.max_seq_len, dtype=torch.long), persistent=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B, T, C = x.shape
@@ -51,31 +58,33 @@ class ResonanceLayerKaggle(nn.Module):
         freq = torch.tanh(self.to_freq(h)) * 2.0  # B,T,H
         phase = torch.tanh(self.to_phase(h)) * math.pi
 
-        # Use torch.long for absolute positions to avoid float precision issues during subtraction
-        pos = torch.arange(T, device=x.device, dtype=torch.long)
-        out = torch.zeros(B, T, C, device=x.device, dtype=x.dtype)
+        # pos cache use karo, T tak slice
+        if T > self.pos_cache.shape[0]:
+            # Agar kabhi bada T aaye to dynamically banao
+            pos = torch.arange(T, device=x.device, dtype=torch.long)
+        else:
+            pos = self.pos_cache[:T]
+
+        out = torch.empty(B, T, C, device=x.device, dtype=x.dtype)
         
         for i in range(0, T, self.chunk_size):
             end = min(i+self.chunk_size, T)
+            chunk_len = end - i
+            
             phase_i = phase[:, i:end, :].permute(0,2,1).unsqueeze(-1) # B,H,chunk,1
             phase_j = phase.permute(0,2,1).unsqueeze(-2) # B,H,1,T
             
-            # Robust distance calculation in torch.long to prevent overflow/dtype issues
             row_idx = torch.arange(i, end, device=x.device, dtype=torch.long)[:, None]
             col_idx = pos[None, :]
             dist_long = (row_idx - col_idx).clamp(min=0)
-            dist = dist_long.to(freq.dtype).view(1, 1, end-i, T)
+            dist = dist_long.to(freq.dtype).view(1, 1, chunk_len, T)
             
             freq_i = freq[:, i:end, :].permute(0,2,1).unsqueeze(-1)
             angle = (phase_i - phase_j) + freq_i * dist * 0.05
             
-            # Pure Cosine similarity. We apply the sqrt(d) scaling during the gating phase below
-            score = torch.cos(angle)
-            
-            # Boolean causal mask for precision and stability
-            causal_bool = (row_idx >= col_idx).view(1,1,end-i,T)
+            score = torch.cos(angle) # B,H,chunk,T
             
-            # Mask out future tokens
+            causal_bool = (row_idx >= col_idx).view(1,1,chunk_len,T)
             score = score.masked_fill(~causal_bool, 0.0)
             score = score / math.sqrt(self.head_dim)
             
@@ -86,11 +95,10 @@ class ResonanceLayerKaggle(nn.Module):
             
             V_t = V.permute(0,2,1,3).contiguous() # B,H,T,D
             Bh = B * self.n_heads
-            assert C == self.n_heads * self.head_dim, f"C ({C}) != n_heads*head_dim ({self.n_heads*self.head_dim})"
             
-            s = score.reshape(Bh, end-i, T)
+            s = score.reshape(Bh, chunk_len, T)
             v = V_t.reshape(Bh, T, self.head_dim)
-            o = torch.bmm(s, v).view(B, self.n_heads, end-i, self.head_dim).permute(0,2,1,3).reshape(B, end-i, C)
+            o = torch.bmm(s, v).view(B, self.n_heads, chunk_len, self.head_dim).permute(0,2,1,3).reshape(B, chunk_len, C)
             out[:, i:end] = o
 
         return self.to_out(out)
@@ -105,10 +113,8 @@ class ViuResonance100M(nn.Module):
         self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
         self.resid_dropout = nn.Dropout(config.dropout)
         
-        # Proper initialization sequence to preserve weight tying!
-        # First randomly initialize everything...
+        # Sahi order: pehle init, phir tie
         self.apply(self._init_weights)
-        # THEN explicitly tie the weights
         self.head.weight = self.emb.weight
 
     def _init_weights(self, module: nn.Module):
@@ -130,6 +136,5 @@ class ViuResonance100M(nn.Module):
         logits = self.head(self.norm(x))
         loss = None
         if targets is not None:
-            # targets with value -100 are automatically ignored in cross_entropy
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)
         return logits, loss

code/train.py

@@ -1,9 +1,11 @@
 import os
-# Setup env vars for robust Kaggle training
+# T4 ke liye robust env
 os.environ["NCCL_P2P_DISABLE"] = "1"
 os.environ["NCCL_IB_DISABLE"] = "1"
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+# T4 optimization
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 import math
 import time
@@ -14,7 +16,8 @@ from pathlib import Path
 import torch
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.cuda.amp import GradScaler
+# FIX 1: Naya amp API - T4 ke liye must (bf16 nahi, fp16)
+from torch.amp import GradScaler, autocast
 import numpy as np
 from transformers import AutoTokenizer
 from huggingface_hub import hf_hub_download, HfApi
@@ -70,43 +73,60 @@ def load_tokenizer_robust(hf_token: str, model_repo: str, subfolder: str, is_mai
             p = Path(p_str)
             if (p / "tokenizer.json").exists() or (p / "vocab.json").exists():
                 try:
-                    return AutoTokenizer.from_pretrained(str(p), local_files_only=True, use_fast=fast, trust_remote_code=True)
+                    tok = AutoTokenizer.from_pretrained(str(p), local_files_only=True, use_fast=fast, trust_remote_code=True)
+                    if is_main: logger.info(f"Loaded tokenizer from {p} fast={fast}")
+                    return tok
                 except Exception as e:
+                    if is_main: logger.warning(f"Local tokenizer {p} fast={fast} failed: {e}")
                     if "sentencepiece" in str(e).lower() or "tiktoken" in str(e).lower():
                         continue
-        try:
-            return AutoTokenizer.from_pretrained(model_repo, subfolder=subfolder, token=hf_token, use_fast=fast, trust_remote_code=True)
-        except Exception as e:
-            pass
-        try:
-            return AutoTokenizer.from_pretrained(model_repo, token=hf_token, use_fast=fast, trust_remote_code=True)
-        except Exception as e:
-            pass
+        # HF se try
+        for attempt_subfolder in [subfolder, None]:
+            try:
+                if attempt_subfolder:
+                    tok = AutoTokenizer.from_pretrained(model_repo, subfolder=attempt_subfolder, token=hf_token, use_fast=fast, trust_remote_code=True)
+                else:
+                    tok = AutoTokenizer.from_pretrained(model_repo, token=hf_token, use_fast=fast, trust_remote_code=True)
+                if is_main: logger.info(f"Loaded tokenizer from HF {model_repo} subfolder={attempt_subfolder} fast={fast}")
+                return tok
+            except Exception as e:
+                if is_main: logger.warning(f"HF tokenizer try failed fast={fast} subfolder={attempt_subfolder}: {e}")
+                continue
 
     if is_main:
         logger.warning("Tokenizer failed, using dummy vocab 64000")
     return None
 
 def load_checkpoint(model, optimizer, scaler, api, token, repo_id, filename="checkpoints/latest_chk.pt", strict=False):
-    local_path = "latest_chk.pt"
+    # FIX 2: Path bug fix - hf_hub_download local_dir="." + filename="checkpoints/..." => ./checkpoints/...
+    # Toh local_path bhi wahi hona chahiye
+    local_path = Path(filename)  # checkpoints/latest_chk.pt
+    local_path.parent.mkdir(parents=True, exist_ok=True)
     try:
         logger.info(f"Checking for existing checkpoint on HF: {filename}")
-        hf_hub_download(repo_id=repo_id, filename=filename, local_dir=".", token=token)
-        # Security Note: weights_only=False is used here to allow loading the optimizer state safely inside our trusted environment.
-        ckpt = torch.load(local_path, map_location="cpu", weights_only=False)
+        # download return karta hai actual local path
+        downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=".", token=token)
+        logger.info(f"Downloaded checkpoint to {downloaded_path}")
+        # Security: trusted env me hi weights_only=False use karo
+        ckpt = torch.load(downloaded_path, map_location="cpu", weights_only=False)
         missing, unexpected = model.load_state_dict(ckpt["model"], strict=strict)
         if missing or unexpected:
-            logger.warning(f"Checkpoint loaded with Missing keys: {missing} | Unexpected keys: {unexpected}")
+            logger.warning(f"Checkpoint loaded with Missing: {len(missing)} | Unexpected: {len(unexpected)}")
         else:
             logger.info("Model weights loaded successfully.")
             
-        optimizer.load_state_dict(ckpt["optimizer"])
+        try:
+            optimizer.load_state_dict(ckpt["optimizer"])
+        except Exception as e:
+            logger.warning(f"Optimizer state load failed, continuing: {e}")
         if "scaler" in ckpt and scaler is not None:
-            scaler.load_state_dict(ckpt["scaler"])
+            try:
+                scaler.load_state_dict(ckpt["scaler"])
+            except Exception as e:
+                logger.warning(f"Scaler state load failed: {e}")
             
         step = ckpt.get("step", 0)
         logger.info(f"Resuming from step {step}")
-        Path(local_path).unlink(missing_ok=True)
         return step
     except (EntryNotFoundError, HfHubHTTPError):
         logger.info("No checkpoint found on HF. Starting from scratch.")
@@ -136,6 +156,7 @@ class ShardPool:
             files = self.api.list_repo_files(repo_id=self.repo_id, repo_type="model", token=self.token)
         self.available_files = [f for f in files if f.startswith("shards/") and f.endswith(".bin") and not f.endswith("_val.bin")]
         np.random.shuffle(self.available_files)
+        logger.info(f"Found {len(self.available_files)} training shards")
         
     def fill(self, ctx_len):
         max_attempts = self.max_pool * 5
@@ -145,22 +166,33 @@ class ShardPool:
                 raise RuntimeError("Failed to fill ShardPool: too many invalid shards.")
             if not self.available_files:
                 self._refresh_file_list()
+                if not self.available_files:
+                    raise RuntimeError("No shard files available")
             f = self.available_files.pop()
             local_p = hf_hub_download(repo_id=self.repo_id, repo_type="dataset", filename=f, local_dir=self.work_dir, token=self.token)
             
-            # Safe dtype selection
             dtype = np.uint16 if self.vocab_size <= 65535 else np.uint32
             mmap_arr = np.memmap(local_p, dtype=dtype, mode="r")
             
             if len(mmap_arr) < ctx_len + 1:
-                logger.warning(f"Shard {f} is too small ({len(mmap_arr)} tokens). Discarding.")
+                logger.warning(f"Shard {f} too small ({len(mmap_arr)}). Discarding.")
                 del mmap_arr
                 Path(local_p).unlink(missing_ok=True)
                 attempts += 1
                 continue
+
+            # FIX 3: Vocab validation - token vocab se bada toh shard corrupt
+            if len(mmap_arr) > 1000:
+                sample_max = int(mmap_arr[:1000].max())
+                if sample_max >= self.vocab_size:
+                    logger.warning(f"Shard {f} has token {sample_max} >= vocab {self.vocab_size}. Discarding.")
+                    del mmap_arr
+                    Path(local_p).unlink(missing_ok=True)
+                    attempts += 1
+                    continue
                 
             self.pool.append((local_p, mmap_arr))
-            logger.info(f"Loaded shard {f}")
+            logger.info(f"Loaded shard {f} ({len(mmap_arr)} tokens)")
             
     def sample_batch(self, batch_size, ctx_len, rng):
         xs, ys = [], []
@@ -186,59 +218,73 @@ class ShardPool:
             logger.warning(f"Failed to delete {old_path}: {e}")
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--max-steps", type=int, default=100)
-    parser.add_argument("--micro-batch", type=int, default=2)
+    parser = argparse.ArgumentParser(description="ViuResonance Training - T4 Optimized")
+    parser.add_argument("--max-steps", type=int, default=10000)
+    parser.add_argument("--micro-batch", type=int, default=1, help="T4 16GB ke liye 1 rakho")
     parser.add_argument("--grad-accum", type=int, default=16)
     parser.add_argument("--eval-every", type=int, default=500)
     parser.add_argument("--save-every", type=int, default=500)
     parser.add_argument("--rotate-every", type=int, default=500)
     parser.add_argument("--lr", type=float, default=3e-4)
     parser.add_argument("--weight-decay", type=float, default=0.1)
-    parser.add_argument("--ctx-len", type=int, default=2048)
+    parser.add_argument("--ctx-len", type=int, default=1024, help="T4 pe 1024 safe, 2048 OOM de sakta hai")
     parser.add_argument("--vocab-size", type=int, default=None)
     args = parser.parse_args()
 
+    # FIX 4: T4 TF32 optimization
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+
     token = get_hf_token()
     rank, local_rank, world_size, device, is_ddp = setup_ddp()
     is_main = (rank == 0)
     
-    # Random seed
+    # FIX 5: Seed properly for CUDA
     seed = 42 + rank
     torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
     rng = np.random.default_rng(seed)
 
-    # Load resources
+    # Load tokenizer
     tokenizer = load_tokenizer_robust(token, "ViuAI/ViuRec", "tokenizer", is_main)
     vocab_size = args.vocab_size or (len(tokenizer) if tokenizer is not None else 64000)
-    
-    from model import ViuResonance100M, ModelConfig
-    config = ModelConfig(vocab_size=vocab_size)
-    model = ViuResonance100M(config).to(device)
-    
-    if is_ddp:
-        model = DDP(model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False)
-        unwrapped_model = model.module
-    else:
-        unwrapped_model = model
+    if is_main:
+        logger.info(f"Using vocab_size={vocab_size}, ctx_len={args.ctx_len}, device={device}, is_ddp={is_ddp}")
 
-    # Optional torch.compile fallback
+    from model_fixed_T4 import ViuResonance100M, ModelConfig
+    config = ModelConfig(vocab_size=vocab_size, chunk_size=256) # T4 ke liye 256
+    unwrapped_model = ViuResonance100M(config).to(device)
+    
+    # FIX 6: torch.compile pehle, DDP baad me - T4 par compile thoda heavy hai, fallback rakha hai
     try:
-        model = torch.compile(model)
-        if is_main: logger.info("Compiled model successfully.")
+        unwrapped_model = torch.compile(unwrapped_model, mode="default") # T4 ke liye reduce-overhead bhi try kar sakte ho
+        if is_main: logger.info("Compiled model successfully (before DDP).")
+        compiled = True
     except Exception as e:
-        if is_main: logger.warning(f"Failed to compile model: {e}. Falling back to eager mode.")
+        if is_main: logger.warning(f"Failed to compile model: {e}. Falling back to eager.")
+        compiled = False
+
+    if is_ddp:
+        model = DDP(unwrapped_model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False)
+    else:
+        model = unwrapped_model
 
-    # Optimizer
+    # Optimizer - T4 Turing me fused AdamW support nahi hota, fallback rahega
     try:
         optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), fused=True)
-    except:
-        if is_main: logger.warning("fused=True failed for AdamW. Falling back to non-fused.")
+        if is_main: logger.info("Using fused AdamW")
+    except Exception as e:
+        if is_main: logger.warning(f"fused=True failed ({e}). Using non-fused AdamW for T4.")
         optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95))
 
-    # GradScaler for safe fp16/bf16
-    scaler = GradScaler(enabled=not torch.cuda.is_bf16_supported())
+    # FIX 7: Naya GradScaler API + T4 = fp16 (bf16_supported=False hoga T4 par)
+    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+    amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
+    scaler = GradScaler('cuda', enabled=not use_bf16)
+    if is_main:
+        logger.info(f"T4 Detected: bf16_supported={use_bf16}, using amp_dtype={amp_dtype}, scaler_enabled={scaler.is_enabled()}")
 
     api = HfApi()
     start_step = load_checkpoint(unwrapped_model, optimizer, scaler, api, token, "ViuAI/ViuRec", strict=False)
@@ -247,7 +293,7 @@ def main():
     shard_pool = ShardPool(api, "ViuAI/viuai-500m-data", work_dir / "shards", token, vocab_size)
     shard_pool.fill(args.ctx_len)
 
-    # Val Data Safely (cap at 20M tokens)
+    # Val Data (only main)
     val_tokens = None
     if is_main:
         logger.info("Loading validation data...")
@@ -260,10 +306,11 @@ def main():
                 val_p = hf_hub_download(repo_id="ViuAI/viuai-500m-data", repo_type="dataset", filename=val_files[0], local_dir=work_dir, token=token)
                 dtype = np.uint16 if vocab_size <= 65535 else np.uint32
                 val_arr = np.memmap(val_p, dtype=dtype, mode="r")
-                max_toks = 20_000_000
-                val_tokens = np.array(val_arr[:max_toks]) # Copy to RAM, safely capped
+                max_toks = 2_000_000 # T4 RAM ke liye 20M se 2M kiya
+                val_tokens = np.array(val_arr[:max_toks])
                 del val_arr
-                Path(val_p).unlink(missing_ok=True)
+                # Val file ko delete mat karo taki dobara download na ho - T4 pe time bachega
+                # Path(val_p).unlink(missing_ok=True)
         except Exception as e:
             logger.warning(f"Failed to load val data: {e}")
 
@@ -272,13 +319,12 @@ def main():
         if val_tokens is None or len(val_tokens) < args.ctx_len + 1: return 0.0
         unwrapped_model.eval()
         val_loss = 0.0
-        val_iters = 20
+        val_iters = 10 # T4 pe kam iter
         for _ in range(val_iters):
             i = rng.integers(0, len(val_tokens) - args.ctx_len - 1)
             x = torch.from_numpy(val_tokens[i:i+args.ctx_len].astype(np.int64)).unsqueeze(0).to(device, non_blocking=True)
             y = torch.from_numpy(val_tokens[i+1:i+1+args.ctx_len].astype(np.int64)).unsqueeze(0).to(device, non_blocking=True)
-            with torch.autocast(device_type="cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
-                # Pass use_checkpoint=False for evaluation
+            with autocast('cuda', dtype=amp_dtype):
                 _, loss = unwrapped_model(x, y, use_checkpoint=False)
             val_loss += loss.item()
         unwrapped_model.train()
@@ -287,13 +333,16 @@ def main():
     # Training Loop
     model.train()
     if is_main:
-        logger.info(f"Starting training for {args.max_steps} steps...")
+        logger.info(f"Starting training for {args.max_steps} steps from step {start_step}...")
         
     for step in range(start_step, args.max_steps):
         # LR Schedule
         warmup = 1000
         progress = min(1.0, step / max(1, warmup))
-        lr = args.lr * progress if step < warmup else args.lr * (0.5 * (1.0 + math.cos(math.pi * (step - warmup) / max(1, args.max_steps - warmup))))
+        if step < warmup:
+            lr = args.lr * progress
+        else:
+            lr = args.lr * (0.5 * (1.0 + math.cos(math.pi * (step - warmup) / max(1, args.max_steps - warmup))))
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -306,7 +355,7 @@ def main():
             x = torch.from_numpy(x_np).to(device, non_blocking=True)
             y = torch.from_numpy(y_np).to(device, non_blocking=True)
             
-            with torch.autocast(device_type="cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
+            with autocast('cuda', dtype=amp_dtype):
                 _, loss = model(x, y)
                 loss = loss / args.grad_accum
             
@@ -314,20 +363,22 @@ def main():
             accum_loss += loss.item()
             del x, y, loss
 
-        # Safe Gradient Clipping with Scaler
+        # Gradient Clipping
         scaler.unscale_(optimizer)
         torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
         
-        # Check for NaNs
+        # NaN Check
         has_nan = not math.isfinite(accum_loss)
-        for p in model.parameters():
-            if p.grad is not None and not torch.isfinite(p.grad).all():
-                has_nan = True
-                break
+        if not has_nan:
+            for p in model.parameters():
+                if p.grad is not None and not torch.isfinite(p.grad).all():
+                    has_nan = True
+                    break
                 
         if has_nan:
-            if is_main: logger.warning(f"NaN/Inf detected at step {step}! Skipping optimizer step.")
+            if is_main: logger.warning(f"NaN/Inf detected at step {step}! Skipping.")
             optimizer.zero_grad(set_to_none=True)
+            scaler.update() # FIX 8: NaN par bhi scaler update karna zaruri hai warna stuck hoga
             continue
             
         scaler.step(optimizer)
@@ -337,13 +388,15 @@ def main():
         
         if is_main and step % 10 == 0:
             tok_s = (args.micro_batch * args.grad_accum * args.ctx_len * world_size) / dt if dt > 0 else 0
-            mem_alloc = torch.cuda.memory_allocated(device) / 1e9
-            mem_res = torch.cuda.memory_reserved(device) / 1e9
-            logger.info(f"step {step:5d}/{args.max_steps} | loss {accum_loss:.4f} | lr {lr:.2e} | {tok_s:.0f} tok/s | mem alloc={mem_alloc:.2f}GB res={mem_res:.2f}GB")
+            mem_alloc = torch.cuda.memory_allocated(device) / 1e9 if torch.cuda.is_available() else 0
+            mem_res = torch.cuda.memory_reserved(device) / 1e9 if torch.cuda.is_available() else 0
+            logger.info(f"step {step:5d}/{args.max_steps} | loss {accum_loss:.4f} | lr {lr:.2e} | {tok_s:.0f} tok/s | mem {mem_alloc:.2f}GB/{mem_res:.2f}GB")
 
         if step > start_step and step % args.rotate_every == 0:
             shard_pool.rotate_one()
             shard_pool.fill(args.ctx_len)
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache() # T4 me fragmentation se bachne ke liye
 
         if step > start_step and step % args.eval_every == 0 and is_main:
             val_loss = evaluate()
@@ -353,16 +406,18 @@ def main():
             if is_ddp: dist.barrier()
             if is_main:
                 logger.info(f"Saving checkpoint at step {step}...")
+                ckpt_path = Path("checkpoints/latest_chk.pt")
+                ckpt_path.parent.mkdir(parents=True, exist_ok=True)
                 chk = {
                     "model": unwrapped_model.state_dict(),
                     "optimizer": optimizer.state_dict(),
                     "scaler": scaler.state_dict(),
                     "step": step,
                 }
-                torch.save(chk, "latest_chk.pt")
+                torch.save(chk, ckpt_path)
                 try:
                     api.upload_file(
-                        path_or_fileobj="latest_chk.pt",
+                        path_or_fileobj=str(ckpt_path),
                         path_in_repo="checkpoints/latest_chk.pt",
                         repo_id="ViuAI/ViuRec",
                         token=token