Update train_on_hf_spaces.py: better error handling, show_error=True

train_on_hf_spaces.py

@@ -19,6 +19,7 @@ import sys
 import json
 import time
 import traceback
+import gc
 from typing import Generator
 
 import torch
@@ -64,7 +65,6 @@ class VideoTokenDataset(Dataset):
         item = self.data[idx]
         prompt = item["text_prompt"]
         tokens = item["video_tokens"][: self.max_tokens]
-        # Pad to fixed length
         while len(tokens) < self.max_tokens:
             tokens.append(0)
         return {
@@ -83,26 +83,36 @@ def train(data_path: str = "tokenized_dataset.json") -> Generator[str, None, Non
     yield "🚀 Starting training pipeline...\n"
 
     # ── 1. Load tokenizer & model ──────────────────────────────────────────
-    yield "📦 Loading OLMo 2 1B Instruct tokenizer & model (this may take a few minutes)...\n"
+    yield "📦 Loading OLMo 2 1B Instruct tokenizer...\n"
 
     try:
         from transformers import AutoModelForCausalLM, AutoTokenizer
-    except ImportError:
-        yield "❌ transformers not installed. Adding to requirements...\n"
+    except ImportError as e:
+        yield f"❌ transformers not installed: {e}\n"
         raise
 
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        yield f"✅ Tokenizer loaded. Vocab size: {len(tokenizer)}\n"
+    except Exception as e:
+        yield f"❌ Failed to load tokenizer: {e}\n"
+        yield traceback.format_exc() + "\n"
+        raise
 
     yield "📦 Loading model in float32 on CPU (this takes ~2-3 min)...\n"
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        trust_remote_code=True,
-        torch_dtype=torch.float32,
-        device_map="cpu",
-    )
-    yield f"✅ Model loaded. Original vocab size: {len(tokenizer)}\n"
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            trust_remote_code=True,
+            torch_dtype=torch.float32,
+        )
+        yield f"✅ Model loaded. Parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M\n"
+    except Exception as e:
+        yield f"❌ Failed to load model: {e}\n"
+        yield traceback.format_exc() + "\n"
+        raise
 
     # ── 2. Expand vocabulary ───────────────────────────────────────────────
     yield f"🔤 Adding {CODEBOOK_SIZE} visual tokens + special tokens...\n"
@@ -116,90 +126,111 @@ def train(data_path: str = "tokenized_dataset.json") -> Generator[str, None, Non
 
     # ── 3. Apply LoRA ─────────────────────────────────────────────────────
     yield f"🔧 Applying LoRA (r={LORA_R}, alpha={LORA_ALPHA})...\n"
-    from peft import LoraConfig, get_peft_model, TaskType
-
-    lora_config = LoraConfig(
-        r=LORA_R,
-        lora_alpha=LORA_ALPHA,
-        target_modules=["q_proj", "v_proj"],
-        lora_dropout=LORA_DROPOUT,
-        bias="none",
-        task_type=TaskType.CAUSAL_LM,
-    )
-    model = get_peft_model(model, lora_config)
-    model.print_trainable_parameters()
-    yield "✅ LoRA applied.\n"
+    try:
+        from peft import LoraConfig, get_peft_model, TaskType
+
+        lora_config = LoraConfig(
+            r=LORA_R,
+            lora_alpha=LORA_ALPHA,
+            target_modules=["q_proj", "v_proj"],
+            lora_dropout=LORA_DROPOUT,
+            bias="none",
+            task_type=TaskType.CAUSAL_LM,
+        )
+        model = get_peft_model(model, lora_config)
+        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in model.parameters())
+        yield f"✅ LoRA applied. Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)\n"
+    except Exception as e:
+        yield f"❌ Failed to apply LoRA: {e}\n"
+        yield traceback.format_exc() + "\n"
+        raise
 
     # ── 4. Load dataset ───────────────────────────────────────────────────
     yield f"📊 Loading dataset from {data_path}...\n"
-    dataset = VideoTokenDataset(data_path, max_tokens=256)
-    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
-    total_steps = NUM_EPOCHS * len(dataloader)
-    yield f"📊 {len(dataset)} samples × {NUM_EPOCHS} epochs = {total_steps} steps\n"
+    try:
+        dataset = VideoTokenDataset(data_path, max_tokens=256)
+        dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+        total_steps = NUM_EPOCHS * len(dataloader)
+        yield f"📊 {len(dataset)} samples × {NUM_EPOCHS} epochs = {total_steps} steps\n"
+    except Exception as e:
+        yield f"❌ Failed to load dataset: {e}\n"
+        yield traceback.format_exc() + "\n"
+        raise
 
     # ── 5. Train ──────────────────────────────────────────────────────────
     yield "🔥 Starting training loop...\n\n"
 
-    optimizer = torch.optim.Adafactor(model.parameters(), lr=LEARNING_RATE, scale_parameter=True, relative_step=True, warmup_init=True)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
     model.train()
 
     global_step = 0
     running_loss = 0.0
     start_time = time.time()
 
-    for epoch in range(NUM_EPOCHS):
-        epoch_loss = 0.0
-        num_batches = 0
-
-        for batch_idx, batch in enumerate(dataloader):
-            prompt = batch["prompt"][0]
-            video_tokens = batch["video_tokens"][0]
-
-            # Format: <text_start> prompt <text_end> <video_start> tok1 tok2 ... <video_end>
-            token_str = " ".join(f"<v_{t.item()}>" for t in video_tokens)
-            text = f"Create a video of: {prompt} {VIDEO_START} {token_str} {VIDEO_END}"
-
-            inputs = tokenizer(
-                text,
-                return_tensors="pt",
-                truncation=True,
-                max_length=MAX_SEQ_LEN,
-                padding="max_length",
-            )
-
-            # Forward pass
-            outputs = model(**inputs, labels=inputs["input_ids"])
-            loss = outputs.loss / GRADIENT_ACCUMULATION
-
-            # Backward pass
-            loss.backward()
-
-            if (batch_idx + 1) % GRADIENT_ACCUMULATION == 0 or (batch_idx + 1) == len(dataloader):
-                torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
-                optimizer.step()
-                optimizer.zero_grad()
-
-            global_step += 1
-            batch_loss = loss.item() * GRADIENT_ACCUMULATION
-            epoch_loss += batch_loss
-            running_loss += batch_loss
-            num_batches += 1
-
-            elapsed = time.time() - start_time
-            steps_per_sec = global_step / elapsed if elapsed > 0 else 0
-
-            if batch_idx % LOG_EVERY == 0:
-                msg = (
-                    f"  Epoch {epoch + 1}/{NUM_EPOCHS} | "
-                    f"Step {batch_idx + 1}/{len(dataloader)} | "
-                    f"Loss: {batch_loss:.4f} | "
-                    f"Avg: {epoch_loss / num_batches:.4f} | "
-                    f"Speed: {steps_per_sec:.2f} steps/s\n"
+    try:
+        for epoch in range(NUM_EPOCHS):
+            epoch_loss = 0.0
+            num_batches = 0
+
+            for batch_idx, batch in enumerate(dataloader):
+                prompt = batch["prompt"][0]
+                video_tokens = batch["video_tokens"][0]
+
+                # Format training text
+                token_str = " ".join(f"<v_{t.item()}>" for t in video_tokens[:64])  # limit tokens for memory
+                text = f"Create a video of: {prompt} {VIDEO_START} {token_str} {VIDEO_END}"
+
+                inputs = tokenizer(
+                    text,
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=MAX_SEQ_LEN,
+                    padding="max_length",
                 )
-                yield msg
 
-        avg_epoch_loss = epoch_loss / num_batches
-        yield f"\n📈 Epoch {epoch + 1} complete. Avg Loss: {avg_epoch_loss:.4f}\n\n"
+                # Forward pass
+                outputs = model(**inputs, labels=inputs["input_ids"])
+                loss = outputs.loss / GRADIENT_ACCUMULATION
+
+                # Backward pass
+                loss.backward()
+
+                if (batch_idx + 1) % GRADIENT_ACCUMULATION == 0 or (batch_idx + 1) == len(dataloader):
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
+                    optimizer.step()
+                    optimizer.zero_grad()
+
+                global_step += 1
+                batch_loss = loss.item() * GRADIENT_ACCUMULATION
+                epoch_loss += batch_loss
+                running_loss += batch_loss
+                num_batches += 1
+
+                elapsed = time.time() - start_time
+                steps_per_sec = global_step / elapsed if elapsed > 0 else 0
+
+                if batch_idx % LOG_EVERY == 0:
+                    msg = (
+                        f"  Epoch {epoch + 1}/{NUM_EPOCHS} | "
+                        f"Step {batch_idx + 1}/{len(dataloader)} | "
+                        f"Loss: {batch_loss:.4f} | "
+                        f"Avg: {epoch_loss / num_batches:.4f} | "
+                        f"Speed: {steps_per_sec:.2f} steps/s\n"
+                    )
+                    yield msg
+
+                # Free memory
+                del outputs, loss
+                gc.collect()
+
+            avg_epoch_loss = epoch_loss / num_batches
+            yield f"\n📈 Epoch {epoch + 1} complete. Avg Loss: {avg_epoch_loss:.4f}\n\n"
+
+    except Exception as e:
+        yield f"\n❌ Training error: {e}\n"
+        yield traceback.format_exc() + "\n"
+        raise
 
     total_time = time.time() - start_time
     yield f"✅ Training complete in {total_time:.0f}s ({total_time / 60:.1f} min)\n"
@@ -207,32 +238,47 @@ def train(data_path: str = "tokenized_dataset.json") -> Generator[str, None, Non
 
     # ── 6. Merge & push ──────────────────────────────────────────────────
     yield "🔀 Merging LoRA weights back into base model...\n"
-    model = model.merge_and_unload()
+    try:
+        model = model.merge_and_unload()
+        yield "✅ LoRA merged.\n"
+    except Exception as e:
+        yield f"⚠️ Merge note: {e}\n"
 
     yield "💾 Saving model locally...\n"
     save_dir = "./trained_model"
-    model.save_pretrained(save_dir, safe_serialization=True)
-    tokenizer.save_pretrained(save_dir)
+    try:
+        model.save_pretrained(save_dir, safe_serialization=True)
+        tokenizer.save_pretrained(save_dir)
+        yield "✅ Model saved locally.\n"
+    except Exception as e:
+        yield f"❌ Save failed: {e}\n"
+        yield traceback.format_exc() + "\n"
+        raise
 
     yield f"🚀 Pushing to {REPO_ID}...\n"
-    from huggingface_hub import HfApi
-
-    api = HfApi(token=HF_TOKEN)
-
-    # Create model repo if it doesn't exist
     try:
-        api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True, token=HF_TOKEN)
+        from huggingface_hub import HfApi
+
+        api = HfApi(token=HF_TOKEN)
+
+        # Create model repo if it doesn't exist
+        try:
+            api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
+        except Exception as e:
+            yield f"⚠️ Repo creation note: {e}\n"
+
+        api.upload_folder(
+            folder_path=save_dir,
+            repo_id=REPO_ID,
+            repo_type="model",
+            commit_message=f"LoRA-trained OLMo 2 1B (r={LORA_R}, {NUM_EPOCHS} epochs)",
+        )
+        yield f"✅ Model pushed to https://huggingface.co/{REPO_ID}\n"
+        yield "\n🎉 All done! The trained model is now available on HuggingFace.\n"
     except Exception as e:
-        yield f"⚠️ Repo creation note: {e}\n"
-
-    api.upload_folder(
-        folder_path=save_dir,
-        repo_id=REPO_ID,
-        repo_type="model",
-        commit_message=f"LoRA-trained OLMo 2 1B (r={LORA_R}, {NUM_EPOCHS} epochs)",
-    )
-    yield f"✅ Model pushed to https://huggingface.co/{REPO_ID}\n"
-    yield "\n🎉 All done! The trained model is now available on HuggingFace.\n"
+        yield f"❌ Push failed: {e}\n"
+        yield traceback.format_exc() + "\n"
+        raise
 
 
 # ---------------------------------------------------------------------------