Upload fine_tune_jit_with_validation_1b.py

fine_tune_jit_with_validation_1b.py

@@ -1,14 +1,8 @@
-# Copyright (c) 2025 CMS Manhattan
-# All rights reserved.
-#
-# This file is part of a project authored by CMS Manhattan. You may use, distribute, and modify
-# this code under the terms of the APACHE 2.0 license.
-
 import os
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import IterableDataset, DataLoader
 from transformers import GPT2TokenizerFast
 from tqdm import tqdm
 import shutil
@@ -66,55 +60,53 @@ DATASET_PATH = CLEAN_PATH
 OUTPUT_DIR = Path("build/fine_tuning_output")
 MODEL_SAVE_NAME = "gpt_finetuned.script.pt"
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")
 print(f"Using device: {device}")
 
-# ============================= DATASET =============================
+# ============================= DATASET (LAZY) =============================
 
-class TextDataset(Dataset):
+class LazyTextDataset(IterableDataset):
+    """Lazy memory-efficient dataset, splits on-the-fly into train and val."""
     def __init__(self, text_file, seq_len=TRAIN_SEQ_LEN, tokenizer_name="gpt2", split_type='train', val_ratio=VAL_SPLIT_RATIO):
         self.seq_len = seq_len
         self.tokenizer = GPT2TokenizerFast.from_pretrained(tokenizer_name)
         self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.text_file = text_file
         self.split_type = split_type
-
-        print(f"Loading text from {text_file} for {split_type} split...")
-        text = Path(text_file).read_text(encoding="utf-8")
-        tokens = self.tokenizer.encode(text)
-
-        if len(tokens) < seq_len * 2:
-            raise ValueError("Text too short!")
-
-        all_inputs = []
-        all_labels = []
-
-        for i in range(0, len(tokens) - seq_len, seq_len):
-            all_inputs.append(tokens[i:i + seq_len])
-            all_labels.append(tokens[i + 1:i + seq_len + 1])
-
-        total_sequences = len(all_inputs)
-        val_size = int(total_sequences * val_ratio)
-        train_size = total_sequences - val_size
-
-        if self.split_type == 'train':
-            self.inputs = all_inputs[:train_size]
-            self.labels = all_labels[:train_size]
-        elif self.split_type == 'val':
-            self.inputs = all_inputs[train_size:]
-            self.labels = all_labels[train_size:]
+        self.val_ratio = val_ratio
+
+        print(f"Loading and tokenizing text from {text_file}")
+        with open(text_file, "r", encoding="utf-8") as f:
+            self.data = f.read()
+        self.tokens = self.tokenizer.encode(self.data)
+        
+        # Work out split indices
+        total_tokens = len(self.tokens) - 1  # because label sequence shifted
+        total_batches = total_tokens // seq_len
+        val_size = int(total_batches * self.val_ratio)
+        train_size = total_batches - val_size
+        if split_type == 'train':
+            self.start = 0
+            self.stop = train_size
+        elif split_type == 'val':
+            self.start = train_size
+            self.stop = train_size + val_size
         else:
-            raise ValueError("Invalid split_type. Must be 'train' or 'val'.")
-
-        print(f"Created {len(self.inputs):,} sequences for {self.split_type} split.")
+            raise ValueError(f"split_type should be 'train' or 'val', got {split_type}")
+        self.total_sequences = self.stop - self.start
+        print(f"Lazy dataset: {self.total_sequences:,} sequences for {split_type} split (from {total_batches:,} total)")
+
+    def __iter__(self):
+        for i in range(self.start * self.seq_len, self.stop * self.seq_len, self.seq_len):
+            # Make sure last batch fits
+            if i + self.seq_len + 1 > len(self.tokens):
+                break
+            input_seq = torch.tensor(self.tokens[i : i + self.seq_len], dtype=torch.long)
+            label_seq = torch.tensor(self.tokens[i + 1 : i + self.seq_len + 1], dtype=torch.long)
+            yield input_seq, label_seq
 
     def __len__(self):
-        return len(self.inputs)
-
-    def __getitem__(self, idx):
-        return (
-            torch.tensor(self.inputs[idx], dtype=torch.long),
-            torch.tensor(self.labels[idx], dtype=torch.long)
-        )
+        return self.total_sequences
 
 # ============================= GET LOGITS UTIL =============================
 
@@ -130,6 +122,7 @@ def get_logits_from_model(model, inputs):
 def evaluate(model, dataloader, criterion, device):
     model.eval()
     total_loss = 0.0
+    count = 0
     with torch.no_grad():
         for inputs, targets in dataloader:
             inputs, targets = inputs.to(device), targets.to(device)
@@ -138,7 +131,8 @@ def evaluate(model, dataloader, criterion, device):
             targets = targets.contiguous().view(-1)[:logits.shape[0]]
             loss = criterion(logits, targets)
             total_loss += loss.item()
-    avg_loss = total_loss / len(dataloader)
+            count += 1
+    avg_loss = total_loss / max(count, 1)
     model.train()
     return avg_loss
 
@@ -185,16 +179,17 @@ def train():
     except AttributeError:
         print("⚠️ Warning: model.gradient_checkpointing_enable() not found on JIT model. Training will proceed without GC.")
 
-    train_dataset = TextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, split_type='train', val_ratio=VAL_SPLIT_RATIO)
-    val_dataset = TextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, split_type='val', val_ratio=VAL_SPLIT_RATIO)
+    train_dataset = LazyTextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, split_type='train', val_ratio=VAL_SPLIT_RATIO)
+    val_dataset = LazyTextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, split_type='val', val_ratio=VAL_SPLIT_RATIO)
 
-    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
-    val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True)
+    # IterableDataset: must use drop_last=True and shuffle=False, num_workers=0 on CPU
+    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True, num_workers=0)
+    val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True, num_workers=0)
 
     optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
     criterion = nn.CrossEntropyLoss()
 
-    total_steps = len(train_dataloader) * EPOCHS
+    total_steps = (len(train_dataset) // BATCH_SIZE) * EPOCHS
     print(f"\n=== BEGINNING LONG-TERM TRAINING ===")
     print(f"Epochs: {EPOCHS} | Steps (Train): {total_steps} | Examples (Train): {len(train_dataset)}")
     print(f"Batch Size (Effective): {BATCH_SIZE} | Precision: FP32")
@@ -223,10 +218,10 @@ def train():
                 pbar.set_postfix({
                     "loss": f"{loss_val:.3f}",
                     "ppl": f"{math.exp(min(loss_val, 10)):.1f}",
-                    "step": f"{global_step}/{total_steps}"
+                    "step": f"{global_step}"
                 })
 
-        avg_train_loss = epoch_loss / len(train_dataloader)
+        avg_train_loss = epoch_loss / max(1, len(train_dataset) // BATCH_SIZE)
         print(f"  [TRAIN] Average loss: {avg_train_loss:.3f} | PPL: {math.exp(avg_train_loss):.1f}")
 
         print("  [VALIDATION] Starting evaluation...")