Upload 163 files

hugging/td_fuse/selfimprove.py

@@ -1,19 +1,22 @@
 """
-TD Self-Improvement v6 — MEGA CYCLE
+TD Self-Improvement v7 — SPEED OPTIMIZED MEGA CYCLE
 
-v5 fixes + MEGA MODE for cycle 3+:
+v7 speed changes (March 2026):
+- Flash Attention 2: 30-40% faster attention computation
+- TF32 math: ~15-20% faster matrix multiply on Ampere GPUs (A6000)
+- Dynamic padding: pad to longest-in-batch, not fixed 1024 tokens
+- 4 dataloader workers + pin_memory: parallel data loading
+- 2 epochs (was 3): model at 95%+, third epoch had diminishing returns
+- Combined: ~50-60% faster → estimated 3-4 hours per cycle (was 6-7)
+
+v6 MEGA MODE (unchanged):
 - 8000+ training samples (was 3000)
 - Priority: code, math, reasoning, creativity (top), then general (second)
 - LoRA r=128 (was 64) — maximum reshape capacity
 - ALL 36 layers targeted (was 16)
-- gradient_checkpointing for speed + VRAM savings
-- batch_size 4 (was 2) — 2x faster per step
+- gradient_checkpointing for VRAM savings
+- batch_size 4 + grad_accum 4 = effective batch 16
 - Norm unclogging continues from cycle 2
-
-SPEED OPTIMIZATIONS:
-- gradient_checkpointing = True (saves ~40% VRAM, allows bigger batches)
-- batch_size 4 + grad_accum 4 = effective batch 16 (same quality, faster)
-- Processes 2x more samples per forward pass = ~2x faster wall clock
 """
 
 import torch
@@ -65,14 +68,23 @@ class SelfImproveConfig:
 
 
 def apply_mega_mode(cfg: SelfImproveConfig):
-    """Crank everything up for mega cycles (3+)."""
-    print("\n  *** MEGA MODE ENABLED ***")
+    """Crank everything up for mega cycles (3+).
+
+    v7 SPEED optimizations (March 2026):
+    - Flash Attention 2: 30-40% faster attention
+    - TF32 math: ~15-20% faster matmul on Ampere GPUs
+    - Dynamic padding: no wasted compute on short samples
+    - 4 dataloader workers: parallel data prep
+    - 2 epochs (was 3): model is 95%+, diminishing returns on epoch 3
+    Combined: ~50-60% faster → ~3-4 hours per cycle instead of 6-7
+    """
+    print("\n  *** MEGA MODE ENABLED (v7 — speed optimized) ***")
     cfg.mega_mode = True
     cfg.top_n_layers = 36          # ALL layers — leave nothing behind
     cfg.min_improvement_score = 0.0  # Target everything
     cfg.lora_r = 128               # Maximum reshape capacity
     cfg.lora_alpha = 256           # 2x rank
-    cfg.train_epochs = 3           # 3 epochs over 8000 samples = plenty
+    cfg.train_epochs = 2           # 2 epochs — model is 95%+, 3rd epoch = diminishing returns
     cfg.train_batch = 4            # 2x bigger batch — 2x faster
     cfg.train_grad_accum = 4       # Effective batch = 16 (same quality)
     cfg.num_train_samples = 8000   # Nearly 3x more data
@@ -83,7 +95,8 @@ def apply_mega_mode(cfg: SelfImproveConfig):
     print(f"    Layers: ALL {cfg.top_n_layers}")
     print(f"    Data: {cfg.num_train_samples} samples")
     print(f"    Batch: {cfg.train_batch} × {cfg.train_grad_accum} = {cfg.train_batch * cfg.train_grad_accum}")
-    print(f"    Speed: gradient_checkpointing + batch 4")
+    print(f"    Epochs: {cfg.train_epochs} (was 3 — 33% time saved)")
+    print(f"    Speed: flash_attn2 + tf32 + dynamic_pad + gradient_checkpointing")
 
 
 def auto_detect_model_path() -> Tuple[str, int]:
@@ -679,11 +692,15 @@ def improve_with_lora(model_path: str, target_layers: List[int],
             return model_path
 
     # Load model with OOM guard
+    # SPEED: Enable TF32 for ~15-20% faster matmul on Ampere GPUs (A6000)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
     print("  Loading model...")
     try:
         model = AutoModelForImageTextToText.from_pretrained(
             model_path, dtype=torch.bfloat16,
-            device_map="auto", trust_remote_code=True
+            device_map="auto", trust_remote_code=True,
+            attn_implementation="flash_attention_2"  # SPEED: 30-40% faster attention
         )
     except (RuntimeError, Exception) as e:
         gc.collect()
@@ -1012,13 +1029,13 @@ def improve_with_lora(model_path: str, target_layers: List[int],
 
     random.shuffle(train_texts)
 
-    # Tokenize
+    # Tokenize — SPEED: dynamic padding (pad to longest in batch, not fixed 1024)
     class ImproveDataset(Dataset):
         def __init__(self, texts, tokenizer, max_len=1024):
             self.data = []
             for t in texts:
                 e = tokenizer(t, truncation=True, max_length=max_len,
-                              padding="max_length", return_tensors="pt")
+                              return_tensors="pt")  # No padding here — collator handles it
                 self.data.append({
                     "input_ids": e["input_ids"].squeeze(),
                     "attention_mask": e["attention_mask"].squeeze(),
@@ -1070,6 +1087,7 @@ def improve_with_lora(model_path: str, target_layers: List[int],
         gradient_accumulation_steps=cfg.train_grad_accum,
         learning_rate=cfg.learning_rate,
         bf16=True,
+        tf32=True,  # SPEED: ~15-20% faster on Ampere GPUs
         gradient_checkpointing=cfg.mega_mode,  # Saves ~40% VRAM in mega mode
         logging_steps=max(1, total_steps // 20),
         save_strategy="steps",
@@ -1079,12 +1097,29 @@ def improve_with_lora(model_path: str, target_layers: List[int],
         lr_scheduler_type="cosine",
         optim="adamw_torch",
         report_to="none",
+        dataloader_num_workers=4,    # SPEED: parallel data loading
+        dataloader_pin_memory=True,  # SPEED: faster CPU→GPU transfer
     )
 
+    # SPEED: Dynamic padding collator — pads to longest in batch, not fixed 1024
+    # Custom because labels need -100 padding (ignore index), not tokenizer pad_id
+    def dynamic_pad_collator(batch):
+        max_len = max(item["input_ids"].size(0) for item in batch)
+        pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        input_ids, attention_mask, labels = [], [], []
+        for item in batch:
+            seq_len = item["input_ids"].size(0)
+            pad_len = max_len - seq_len
+            input_ids.append(torch.cat([item["input_ids"], torch.full((pad_len,), pad_id, dtype=torch.long)]))
+            attention_mask.append(torch.cat([item["attention_mask"], torch.zeros(pad_len, dtype=torch.long)]))
+            labels.append(torch.cat([item["labels"], torch.full((pad_len,), -100, dtype=torch.long)]))
+        return {"input_ids": torch.stack(input_ids), "attention_mask": torch.stack(attention_mask), "labels": torch.stack(labels)}
+
     trainer_kwargs = {
         "model": model,
         "processing_class": tokenizer,
         "train_dataset": dataset,
+        "data_collator": dynamic_pad_collator,
         "args": args,
     }
     if optimizer is not None: