Update finetune.py

The official finetuning script.

finetune.py

@@ -1,120 +1,129 @@
 import os
+import time
+import math
+import torch
+from model import GPTConfig, GPT
+
 import numpy as np
-import tiktoken
-from datasets import load_dataset
-from tqdm import tqdm
-
-OUTPUT_DIR = "data/alpaca_cleaned_mixed"
-TOKENIZER_NAME = "gpt2" 
-SEED = 1337
-
-FINEWEB_SAMPLES = 2500 
-
-enc = tiktoken.get_encoding(TOKENIZER_NAME)
-EOS_TOKEN = "<|endoftext|>"
-
-def format_prompt_with_mask(instruction, input_text, output):
-    """
-    Formatiert den Prompt und erstellt die Loss-Maske.
-    Format:
-    Instruction: ...
-    Input: ... (optional)
-    Response: ... <|endoftext|>
-    """
-    if input_text and input_text.strip():
-        prompt_text = f"Instruction:\n{instruction}\n\nInput:\n{input_text}\n\nResponse:\n"
-    else:
-        prompt_text = f"Instruction:\n{instruction}\n\nResponse:\n"
-    
-    completion_text = f"{output}{EOS_TOKEN}"
-    
-    prompt_ids = enc.encode(prompt_text, allowed_special={'<|endoftext|>'})
-    completion_ids = enc.encode(completion_text, allowed_special={'<|endoftext|>'})
-    
-    full_ids = prompt_ids + completion_ids
-    
-    mask = [0] * len(prompt_ids) + [1] * len(completion_ids)
-    
-    return full_ids, mask
 
-def main():
-    np.random.seed(SEED)
-    print(f"🚀 Starting Prepare-Script for SmaLLMPro (350M Instruct)...")
-    print(f"📚 Tokenizer: {TOKENIZER_NAME}")
-    
-    os.makedirs(OUTPUT_DIR, exist_ok=True)
+out_dir = '/home/user/350m_SmaLLMPro_Final'
+init_from = '/home/user/350m_fineweb'
+dataset = 'alpaca_cleaned_mixed'
 
-    print("📥 Loading 'yahma/alpaca-cleaned' (Chat-Instructions)...")
-    alpaca = load_dataset("yahma/alpaca-cleaned", split='train')
-    
-    print(f"📥 Loading 'HuggingFaceFW/fineweb-edu' (Sample-10BT) for {FINEWEB_SAMPLES} Samples...")
-    fineweb = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split='train', streaming=True)
-    
-    all_tokens = []
-    all_masks = []
-    
-    print("⚙️  Processing Alpaca...")
-    for ex in tqdm(alpaca, desc="Alpaca"):
-        ids, mask = format_prompt_with_mask(ex['instruction'], ex['input'], ex['output'])
-        all_tokens.extend(ids)
-        all_masks.extend(mask)
-        
-    alpaca_len = len(all_tokens)
-    print(f"   -> Alpaca Tokens: {alpaca_len:,}")
-
-    print("⚙️  Processing FineWeb (Anti-Forgetting)...")
-    fw_iter = iter(fineweb)
-    fw_count = 0
-    fw_tokens_count = 0
-    
-    for _ in tqdm(range(FINEWEB_SAMPLES), desc="FineWeb"):
-        try:
-            ex = next(fw_iter)
-            text = ex['text'] + EOS_TOKEN
-            ids = enc.encode(text, allowed_special={EOS_TOKEN})
-            
-            all_tokens.extend(ids)
-            all_masks.extend([1] * len(ids)) 
-            
-            fw_tokens_count += len(ids)
-            fw_count += 1
-        except StopIteration:
-            break
-            
-    print(f"   -> FineWeb Tokens: {fw_tokens_count:,} (from {fw_count} documents)")
-
-    total_tokens = len(all_tokens)
-    print(f"\n💾 Saving {total_tokens:,} Tokens in '{OUTPUT_DIR}'...")
-    
-    token_arr = np.array(all_tokens, dtype=np.uint16)
-    token_arr.tofile(os.path.join(OUTPUT_DIR, "train.bin"))
-    
-    mask_arr = np.array(all_masks, dtype=np.uint8)
-    mask_arr.tofile(os.path.join(OUTPUT_DIR, "train_mask.bin"))
-    
-    print("\n🔍 --- SANITY CHECK ---")
-    print("I decode the first 50 tokens of the first sample, to check, if everything is okay.")
-    print("Green (TRAIN) = The things the model learns. Grey (IGNORE) = The things the model only reads.")
-    
-    check_len = 100
-    sample_ids = all_tokens[:check_len]
-    sample_mask = all_masks[:check_len]
+batch_size = 4
+gradient_accumulation_steps = 32
+block_size = 1024
+learning_rate = 2e-5
+max_iters = 1500
+weight_decay = 0.1
+dropout = 0.1
+warmup_iters = 0
+min_lr = 3e-6
+beta1, beta2 = 0.9, 0.95
+device = 'cuda'
+dtype = 'bfloat16'
+compile = True
+save_interval = 500
+
+os.makedirs(out_dir, exist_ok=True)
+torch.manual_seed(1337)
+device_type = 'cuda' if 'cuda' in device else 'cpu'
+ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+
+data_dir = os.path.join('data', dataset)
+train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
+train_mask = np.memmap(os.path.join(data_dir, 'train_mask.bin'), dtype=np.uint8, mode='r')
+
+def get_batch():
+    ix = torch.randint(len(train_data) - block_size, (batch_size,))
+    x = torch.stack([torch.from_numpy((train_data[i:i+block_size]).astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy((train_data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
+    m = torch.stack([torch.from_numpy((train_mask[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
     
-    decoded_parts = []
-    for t_id, m_val in zip(sample_ids, sample_mask):
-        token_str = enc.decode([t_id])
-        if m_val == 1:
-            decoded_parts.append(f"\033[92m{token_str}\033[0m")
-        else:
-            decoded_parts.append(f"\033[90m{token_str}\033[0m")
-            
-    print("".join(decoded_parts))
-    print("\n(Legend: \033[90mGrey=Prompt/Ignored\033[0m, \033[Green=Response/Learned\033[0m)")
+    y[m == 0] = -100
     
-    if len(token_arr) != len(mask_arr):
-        print("\n❌ Warning: Token and Mask Array have different lengths! Something has gone wrong!")
-    else:
-        print("\n✅ Everything seems to be fine. The arrays are synchronized. You can now start the training.")
+    x, y = x.to(device), y.to(device)
+    return x, y
+
+print(f"📥 Loading Pretraining-Checkpoint from {init_from}...")
+ckpt_files = sorted([f for f in os.listdir(init_from) if f.endswith('.pt')])
+if not ckpt_files:
+    raise FileNotFoundError("No checkpoint found in init_from directory!")
+
+ckpt_path = os.path.join(init_from, ckpt_files[-1])
+checkpoint = torch.load(ckpt_path, map_location=device)
+gptconf = GPTConfig(**checkpoint['model_args'])
+model = GPT(gptconf)
+state_dict = checkpoint['model']
+
+unwanted_prefix = '_orig_mod.'
+for k,v in list(state_dict.items()):
+    if k.startswith(unwanted_prefix):
+        state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+
+model.load_state_dict(state_dict)
+model.to(device)
+
+if compile:
+    print("🚀 Compiling Model...")
+    model = torch.compile(model)
+
+optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
+scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
+
+def get_lr(it):
+    if it < warmup_iters: return learning_rate * it / warmup_iters
+    if it > max_iters: return min_lr
+    decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return min_lr + coeff * (learning_rate - min_lr)
+
+print(f"🛠️ Starting Finetuning...")
+model.train()
+t0 = time.time()
+
+for iter_num in range(max_iters + 1):
+    lr = get_lr(iter_num)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+    for micro_step in range(gradient_accumulation_steps):
+        X, Y = get_batch()
+        with ctx:
+            logits, loss = model(X, Y)
+            loss = loss / gradient_accumulation_steps
+        scaler.scale(loss).backward()
+
+    scaler.unscale_(optimizer)
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    scaler.step(optimizer)
+    scaler.update()
+    optimizer.zero_grad(set_to_none=True)
+
+    if iter_num % 10 == 0:
+        dt = time.time() - t0
+        print(f"Iter {iter_num}: Loss {loss.item()*gradient_accumulation_steps:.4f}, Time {dt*1000:.2f}ms, LR {lr:.2e}")
+        t0 = time.time()
+
+    if iter_num > 0 and iter_num % save_interval == 0:
+        checkpoint_name = f'SmaLLMPro_iter_{iter_num}.pt'
+        save_path = os.path.join(out_dir, checkpoint_name)
+        print(f"💾 Saving checkpoint: {checkpoint_name}")
+        raw_model = model._orig_mod if compile else model
+        checkpoint_data = {
+            'model': raw_model.state_dict(),
+            'model_args': checkpoint['model_args'],
+            'iter_num': iter_num,
+            'lr': lr,
+        }
+        torch.save(checkpoint_data, save_path)
 
-if __name__ == "__main__":
-    main()
+print(f"💾 Finetuning done. Saving SmaLLMPro...")
+final_checkpoint = {
+    'model': model.state_dict() if not compile else model._orig_mod.state_dict(),
+    'model_args': checkpoint['model_args'],
+    'config': checkpoint.get('config', {}),
+}
+torch.save(final_checkpoint, os.path.join(out_dir, 'SmaLLMPro_Final.pt'))
+print("✅ SmaLLMPro saved successfully!")