Create train.py

train.py

@@ -0,0 +1,172 @@
+import os
+import modal
+
+# Offload all the heavy dependency installations to the Modal cloud container
+image = (
+    modal.Image.debian_slim()
+    .pip_install(
+        "transformers",
+        "datasets",
+        "torch",
+        "tokenizers",
+        "huggingface_hub",
+        "accelerate"
+    )
+)
+
+app = modal.App("tralalabs-16m-qwen-master-pretrain")
+
+@app.function(
+    image=image,
+    gpu="L40S",
+    timeout=86400, # 24 hours max runtime allowed
+    secrets=[modal.Secret.from_name("huggingface-secret")]
+)
+def train():
+    import torch
+    import torch.nn as nn
+    from torch.utils.data import IterableDataset, DataLoader
+    from datasets import load_dataset
+    from tokenizers import Tokenizer, models, trainers, pre_tokenizers
+    from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM
+    from huggingface_hub import HfApi
+    from torch.optim import AdamW
+    
+    print("Initialization started! Fetching data for Tokenizer and Training...")
+    
+    hf_token = os.environ.get("HF_TOKEN")
+    if not hf_token:
+        print("Error: HF_TOKEN environment variable missing in your Modal secret.")
+        return
+
+    # 1. Stream the 85% / 10% / 5% mix
+    try:
+        ds_fw_2024 = load_dataset("HuggingFaceFW/fineweb-edu", "CC-MAIN-2024-18", split="train", streaming=True)
+        ds_wiki = load_dataset("wikipedia", "20231101.en", split="train", streaming=True)
+        ds_fw_2023 = load_dataset("HuggingFaceFW/fineweb-edu", "CC-MAIN-2023-50", split="train", streaming=True)
+        
+        def batch_iterator(batch_size=1000):
+            fw_2024_iter = iter(ds_fw_2024)
+            wiki_iter = iter(ds_wiki)
+            fw_2023_iter = iter(ds_fw_2023)
+            
+            # Infinite loop generator for the massive 81k step pre-training run
+            while True:
+                batch = []
+                for _ in range(int(batch_size * 0.85)): batch.append(next(fw_2024_iter)["text"])
+                for _ in range(int(batch_size * 0.10)): batch.append(next(wiki_iter)["text"])
+                for _ in range(int(batch_size * 0.05)): batch.append(next(fw_2023_iter)["text"])
+                yield batch
+    except Exception as e:
+        print(f"Error setting up datasets: {e}")
+        return
+
+    # 2. Train Tokenizer (16k Vocab) using the first few batches
+    print("Training 16k Byte-Level BPE Tokenizer...")
+    raw_tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
+    raw_tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    trainer = trainers.BpeTrainer(vocab_size=16000, special_tokens=["<unk>", "<s>", "</s>", "<pad>", "<mask>"])
+    
+    # Grab a finite chunk of data to train the vocabulary, then stop
+    def tokenizer_iterator():
+        iterator = batch_iterator(1000)
+        for _ in range(20): 
+            yield next(iterator)
+            
+    raw_tokenizer.train_from_iterator(tokenizer_iterator(), trainer=trainer)
+    
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=raw_tokenizer, 
+        bos_token="<s>", 
+        eos_token="</s>", 
+        unk_token="<unk>", 
+        pad_token="<pad>", 
+        mask_token="<mask>"
+    )
+    tokenizer.pad_token = "<pad>"
+    
+    os.makedirs("./outputs", exist_ok=True)
+    tokenizer.save_pretrained("./outputs")
+
+    # 3. Model Hyperparameters: 16.7M params Qwen2 Architecture
+    print("Configuring 16.7M Parameter Qwen2 Architecture...")
+    config = Qwen2Config(
+        vocab_size=16000,
+        hidden_size=384,
+        intermediate_size=1536,
+        num_hidden_layers=6,
+        num_attention_heads=6,
+        num_key_value_heads=2, # GQA activated for maximum efficiency
+        max_position_embeddings=1024,
+        pad_token_id=3,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=True,
+        rope_theta=10000.0
+    )
+    
+    model = Qwen2ForCausalLM(config).to(device="cuda", dtype=torch.bfloat16)
+    
+    # 4. The 334M Token Training Loop
+    print("Tokenizer baked! Starting massive gradient descent pre-training run...")
+    optimizer = AdamW(model.parameters(), lr=6e-4, weight_decay=0.1)
+    model.train()
+    
+    class ProportionalDataset(IterableDataset):
+        def __init__(self, it): self.it = it
+        def __iter__(self):
+            for batch in self.it:
+                for text in batch: yield text
+
+    train_loader = DataLoader(ProportionalDataset(batch_iterator(batch_size=200)), batch_size=4)
+    
+    step = 0
+    # 334,000,000 total tokens / (4 batch size * 1024 sequence length) = 81,543 steps
+    TARGET_STEPS = 81543 
+    
+    for batch_text in train_loader:
+        if step >= TARGET_STEPS: 
+            break 
+        
+        optimizer.zero_grad()
+        
+        encodings = tokenizer(
+            batch_text, 
+            truncation=True, 
+            max_length=1024, 
+            padding="max_length", 
+            return_tensors="pt"
+        )
+        
+        input_ids = encodings["input_ids"].to("cuda")
+        attention_mask = encodings["attention_mask"].to("cuda")
+        
+        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
+        loss = outputs.loss
+        
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        
+        # Log every 50 steps so the terminal doesn't get flooded
+        if step % 50 == 0:
+            print(f"Step {step}/{TARGET_STEPS} | Loss: {loss.item():.4f}")
+            
+        step += 1
+
+    # 5. Save and Push the Final Master Weights
+    print(f"Saving Final Learned Weights after {TARGET_STEPS} steps...")
+    model.save_pretrained("./outputs")
+    
+    repo_id = "Tralalabs/TralaLabs-16M-Base"
+    try:
+        api = HfApi()
+        api.create_repo(repo_id=repo_id, token=hf_token, exist_ok=True)
+        api.upload_folder(folder_path="./outputs", repo_id=repo_id, repo_type="model", token=hf_token)
+        print("Complete master success! Full 334M token Qwen model uploaded.")
+    except Exception as e:
+        print(f"Error uploading to HF: {e}")
+
+@app.local_entrypoint()
+def main():
+    train.remote()