Upload 139 files

hugging/td_fuse/heal.py

@@ -242,6 +242,11 @@ def apply_qlora_standard(
     Returns:
         Path to healed model directory
     """
+    import os
+    healed_check = os.path.join('td_fuse_outputs', 'healed', 'model.safetensors')
+    if os.path.exists(healed_check):
+        print('[heal] Found existing healed model — SKIPPING healing!')
+        return 'td_fuse_outputs/healed'
     from peft import LoraConfig, get_peft_model, TaskType
     from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
@@ -404,6 +409,13 @@ def heal_model(
     if cfg is None:
         cfg = MergeConfig()
 
+    # Skip healing if already done (saves ~45 min on re-runs)
+    import os
+    healed_check = os.path.join('td_fuse_outputs', 'healed', 'model.safetensors')
+    if os.path.exists(healed_check):
+        print('[heal] Found existing healed model — SKIPPING healing!')
+        return 'td_fuse_outputs/healed'
+
     heal_start = time.time()
     print("\n" + "=" * 60)
     print("HEALING FINE-TUNE")

hugging/td_fuse/merge.py

@@ -717,6 +717,19 @@ def run_single_merge(
                 torch.cuda.empty_cache()
             return result
 
+    # --- Step 4.9: Free VRAM before transport computation ---
+    print(f"\n[merge] Step 4.9: Moving models to CPU to free VRAM for transport...")
+    sys.stdout.flush()
+    source_model = source_model.cpu()
+    target_model = target_model.cpu()
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        free_mem = torch.cuda.mem_get_info()[0] / 1e9
+        total_mem = torch.cuda.mem_get_info()[1] / 1e9
+        print(f"[merge] GPU memory after CPU offload: {free_mem:.1f} GB free / {total_mem:.1f} GB total")
+    sys.stdout.flush()
+
     # --- Step 5: Compute transport plans ---
     print(f"\n[merge] Step 5/10: Computing transport plans..."); sys.stdout.flush()
     step_t = time.time()
@@ -750,20 +763,22 @@ def run_single_merge(
                 print(f"[merge] RAM skipped: base model {base_hf_id} not found on HuggingFace")
                 use_ram = False
 
-    # --- Step 5.7: Free source model from GPU ---
-    # After transport plans are computed, we only need the source STATE DICT
-    # (not the full model object). Freeing the model saves ~16 GB of GPU memory
-    # which prevents OOM during the fusion step.
-    print(f"\n[merge] Step 5.7: Freeing source model from GPU..."); sys.stdout.flush()
+    # --- Step 5.7: Free source model, move target back to GPU ---
+    # Source model was moved to CPU in step 4.9. Extract state dict, then delete.
+    # Move target model back to GPU for the fusion step.
+    print(f"\n[merge] Step 5.7: Extracting source state + moving target back to GPU..."); sys.stdout.flush()
     step_t = time.time()
     source_state_cpu = {k: v.cpu() for k, v in source_model.state_dict().items()}
     del source_model
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+    # Move target back to GPU for fusion
+    target_model = target_model.to("cuda")
+    if torch.cuda.is_available():
         free_mem = torch.cuda.mem_get_info()[0] / 1e9
         total_mem = torch.cuda.mem_get_info()[1] / 1e9
-        print(f"[merge] GPU memory after freeing source: {free_mem:.1f} GB free / {total_mem:.1f} GB total")
+        print(f"[merge] GPU memory (target on GPU, source freed): {free_mem:.1f} GB free / {total_mem:.1f} GB total")
     print(f"[merge] Step 5.7 done in {time.time()-step_t:.0f}s"); sys.stdout.flush()
 
     # --- Step 6: Pre-merge protection ---
@@ -982,6 +997,7 @@ def run_single_merge(
 def run_pipeline(
     stages: list[str],
     cfg: MergeConfig = None,
+    base_checkpoint: str = None,
 ) -> dict:
     """
     Run the full merge pipeline.
@@ -1023,8 +1039,17 @@ def run_pipeline(
     Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
     Path(cfg.checkpoint_dir).mkdir(parents=True, exist_ok=True)
 
-    # --- Load target model ---
-    target_model, target_tokenizer = load_model(TARGET, cfg)
+    # --- Load target model (from checkpoint if stacking merges, else from HuggingFace) ---
+    if base_checkpoint and Path(base_checkpoint).exists():
+        print(f"\n[pipeline] Loading target from previous merge: {base_checkpoint}")
+        from transformers import AutoModelForImageTextToText
+        target_model = AutoModelForImageTextToText.from_pretrained(
+            base_checkpoint, torch_dtype=torch.bfloat16, device_map="auto",
+            trust_remote_code=True,
+        )
+        target_tokenizer = AutoTokenizer.from_pretrained(base_checkpoint, trust_remote_code=True)
+    else:
+        target_model, target_tokenizer = load_model(TARGET, cfg)
 
     # --- Inject canary into target (Qwen3's own canary) ---
     if "Qwen3-VL-8B" in CANARY_FACTS:
@@ -1116,6 +1141,16 @@ def run_pipeline(
     if pipeline_results["final_checkpoint"]:
         final_dir = Path(cfg.output_dir) / "final"
         final_dir.mkdir(parents=True, exist_ok=True)
+        # Free disk space before final save (Bug #25 fix)
+        import shutil as _shutil
+        for _cleanup in ["models/base"]:
+            _cp = Path(_cleanup)
+            if _cp.exists() and _cp.is_dir():
+                _shutil.rmtree(str(_cp))
+                print(f"[merge] Freed disk: {_cleanup}")
+        import gc; gc.collect()
+        _stat = _shutil.disk_usage("/")
+        print(f"[merge] Disk: {_stat.free / 1e9:.1f} GB free / {_stat.total / 1e9:.1f} GB total")
         target_model.save_pretrained(final_dir)
         target_tokenizer.save_pretrained(final_dir)
         pipeline_results["final_model_path"] = str(final_dir)

hugging/td_fuse/transport.py

@@ -518,6 +518,7 @@ def fuse_weights(
     transport_plans: dict,
     source_config: ModelConfig,
     cfg: MergeConfig,
+    target_activations: dict = None,
 ) -> AutoModelForCausalLM:
     """
     Fuse source model weights into target model using transport plans.

hugging/td_lang/__pycache__/compiler.cpython-310.pyc

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:97c261ef8c24868bc538ecaed5c905927d7b933d3ad2e9c6032a6de0cb6bb41e
-size 104110
+oid sha256:73d29a1b793e7e773f99ff76abcf307831d391aaa7fc368cc4ab4ac8b3159303
+size 192996

hugging/td_lang/compiler.py

@@ -224,6 +224,7 @@ DO NOT EDIT - regenerate from the .td file instead.
         self._emit("")
         self._emit("def main():")
         self._indent += 1
+        self._emit("import os  # safety: prevent UnboundLocalError if shadowed")
         self._emit("start_time = time.time()")
         self._emit("lineage = {}")
         self._emit("models = {}")
@@ -466,8 +467,21 @@ DO NOT EDIT - regenerate from the .td file instead.
         self._indent += 1
         self._emit('raise SystemExit(f"Could not match source {_source_ref} to any SOURCES entry.")')
         self._indent -= 1
+        self._emit("")
+        self._emit("# Skip merge if checkpoint already exists (Bug #27 - saves ~12 min)")
+        self._emit('_merge_ckpt = Path(f"td_fuse_checkpoints/after_{_stage}")')
+        self._emit("if _merge_ckpt.exists() and (_merge_ckpt / 'model.safetensors').exists():")
+        self._indent += 1
+        self._emit('print(f"[td_lang] Found merge checkpoint {_merge_ckpt} - SKIPPING merge")')
+        self._emit('merge_result = {"status": "skipped", "final_checkpoint": str(_merge_ckpt)}')
+        self._indent -= 1
+        self._emit("else:")
+        self._indent += 1
+        self._emit("# Stack merges: pass previous checkpoint so MiMo builds on DeepSeek, etc.")
+        self._emit(f'_prev_ckpt = models.get("{cmd.target}", {{}}).get("checkpoint")')
         self._emit("cfg = MergeConfig()")
-        self._emit("merge_result = run_pipeline([_stage], cfg)")
+        self._emit("merge_result = run_pipeline([_stage], cfg, base_checkpoint=_prev_ckpt)")
+        self._indent -= 1
         self._emit(f'results["{cmd.target}_merge"] = merge_result')
         self._emit("merged_stages.append(_stage)")
         self._emit('if merge_result.get("final_checkpoint"):')
@@ -1195,21 +1209,41 @@ DO NOT EDIT - regenerate from the .td file instead.
         self._emit("")
 
         if cmd.method == "grpo":
-            self._emit("# GRPO training with QLoRA (test_15: 64 steps sweet spot)")
-            self._emit("# QLoRA = 4-bit base model + LoRA adapters = fits on 24GB 4090")
-            self._emit("from trl import GRPOConfig, GRPOTrainer")
-            self._emit("from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig")
+            self._emit("# Bug #26 fix: Use SFT on merge checkpoint (same approach as healing — proven to work)")
+            self._emit("# GRPOTrainer breaks with Qwen3-VL, but standard Trainer works perfectly")
+            self._emit("from transformers import AutoTokenizer, TrainingArguments, BitsAndBytesConfig, Trainer")
             self._emit("from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training")
-            self._emit("from datasets import load_dataset")
+            self._emit("from datasets import load_dataset, Dataset")
             self._emit("import torch")
             self._emit("")
-            self._emit("tok = AutoTokenizer.from_pretrained(checkpoint)")
+            self._emit("# Use latest merge checkpoint — pick newest after_* dir in td_fuse_checkpoints/")
+            self._emit("_merge_ckpt = None")
+            self._emit("_ckpt_base = Path('td_fuse_checkpoints')")
+            self._emit("if _ckpt_base.exists():")
+            self._indent += 1
+            self._emit("_after_dirs = sorted(_ckpt_base.glob('after_*'), key=lambda p: p.stat().st_mtime, reverse=True)")
+            self._emit("if _after_dirs and (_after_dirs[0] / 'model.safetensors').exists():")
+            self._indent += 1
+            self._emit("_merge_ckpt = str(_after_dirs[0])")
+            self._indent -= 1
+            self._indent -= 1
+            self._emit("if _merge_ckpt:")
+            self._indent += 1
+            self._emit('print(f"[td_lang] Using merge checkpoint for training: {_merge_ckpt}")')
+            self._emit("_train_ckpt = _merge_ckpt")
+            self._indent -= 1
+            self._emit("else:")
+            self._indent += 1
+            self._emit("_train_ckpt = checkpoint")
+            self._emit('print(f"[td_lang] Using checkpoint for training: {_train_ckpt}")')
+            self._indent -= 1
+            self._emit("")
+            self._emit("tok = AutoTokenizer.from_pretrained(_train_ckpt)")
             self._emit("if tok.pad_token is None:")
             self._indent += 1
             self._emit("tok.pad_token = tok.eos_token")
             self._indent -= 1
             self._emit("")
-            self._emit("# 4-bit quantization - shrinks 7B model from 14GB to ~4GB VRAM")
             self._emit("bnb_config = BitsAndBytesConfig(")
             self._indent += 1
             self._emit("load_in_4bit=True,")
@@ -1218,8 +1252,7 @@ DO NOT EDIT - regenerate from the .td file instead.
             self._emit("bnb_4bit_use_double_quant=True,")
             self._indent -= 1
             self._emit(")")
-            self._emit("")
-            self._emit("model = _load_model_smart(checkpoint, quantization_config=bnb_config, device_map='auto')")
+            self._emit("model = _load_model_smart(_train_ckpt, quantization_config=bnb_config, device_map='auto')")
             self._emit("model = prepare_model_for_kbit_training(model)")
             self._emit("")
             self._emit("# LoRA adapters on mid-to-late layers (test_12: layers 16-28 for 32-layer)")
@@ -1246,178 +1279,56 @@ DO NOT EDIT - regenerate from the .td file instead.
             self._emit("train_data = load_dataset(dataset_path, split='train')")
             self._indent -= 1
             self._emit("")
-            self._emit("grpo_config = GRPOConfig(")
+            self._emit("# Format synth data as text for SFT (prompt + response)")
+            self._emit("def _format_synth(example):")
+            self._indent += 1
+            self._emit("text = example['prompt'] + '\\n' + example.get('response', '')")
+            self._emit("tokens = tok(text, truncation=True, max_length=512, padding='max_length')")
+            self._emit("tokens['labels'] = tokens['input_ids'].copy()")
+            self._emit("return tokens")
+            self._indent -= 1
+            self._emit("train_data = train_data.map(_format_synth, remove_columns=train_data.column_names)")
+            self._emit("")
+            self._emit("training_args = TrainingArguments(")
             self._indent += 1
             self._emit(f"max_steps={steps},")
             self._emit(f"learning_rate={lr},")
             self._emit("per_device_train_batch_size=1,")
             self._emit("gradient_accumulation_steps=8,")
-            self._emit("logging_steps=16,  # eval every 16 steps (test_15)")
-            self._emit('output_dir="td_lang_outputs/grpo_training",')
-            self._emit("save_steps=16,")
+            self._emit("logging_steps=10,")
+            self._emit('output_dir="td_lang_outputs/sft_training",')
+            self._emit("save_steps=50,")
             self._emit('bf16=True,')
-            self._emit("gradient_checkpointing=True,  # saves VRAM at slight speed cost")
+            self._emit("gradient_checkpointing=True,")
+            self._emit("remove_unused_columns=False,")
             self._indent -= 1
             self._emit(")")
             self._emit("")
-            self._emit("# Verified rewards only (test_16: no learned reward model)")
-            # Wire in reward_contract verifiers if they exist
-            if program and program.reward_contract and program.reward_contract.verifiers:
-                verifiers = program.reward_contract.verifiers
-                self._emit(f'# reward_contract verifiers wired in: {verifiers}')
-                self._emit(f'_active_verifiers = {verifiers}')
-                if program.reward_contract.min_reward is not None:
-                    self._emit(f'_min_reward = {program.reward_contract.min_reward}')
-                else:
-                    self._emit('_min_reward = 0.0')
-            else:
-                self._emit('_active_verifiers = ["code_compiles", "math_correct"]  # defaults')
-                self._emit('_min_reward = 0.0')
-            self._emit("import ast, math, re")
-            self._emit("ALLOWED_EXPR = re.compile(r'^[0-9+\\-*/().\\s]+$')")
-            self._emit("")
-            self._emit("def _safe_eval(expr: str):")
-            self._indent += 1
-            self._emit("expr = expr.strip()")
-            self._emit("if not ALLOWED_EXPR.match(expr):")
-            self._indent += 1
-            self._emit("return None")
-            self._indent -= 1
-            self._emit("try:")
-            self._indent += 1
-            self._emit("return float(eval(expr, {'__builtins__': {}}, {}))")
-            self._indent -= 1
-            self._emit("except Exception:")
-            self._indent += 1
-            self._emit("return None")
-            self._indent -= 2
-            self._emit("")
-            self._emit("def reward_fn(completions, prompts=None, **kwargs):")
-            self._indent += 1
-            self._emit("prompts = prompts or ['' for _ in completions]")
-            self._emit("rewards = []")
-            self._emit("for comp, prompt in zip(completions, prompts):")
-            self._indent += 1
-            self._emit("text = comp if isinstance(comp, str) else comp[0].get('content', '')")
-            self._emit("score = 0.0")
-            self._emit("# Code compilation reward (active if 'code_compiles' in verifiers)")
-            self._emit("if 'code_compiles' in _active_verifiers:")
-            self._indent += 1
-            self._emit("code_blocks = re.findall(r'```python\\n(.*?)```', text, re.S)")
-            self._emit("for block in code_blocks or []:")
-            self._indent += 1
-            self._emit("try:")
-            self._indent += 1
-            self._emit("ast.parse(block)")
-            self._emit("score += 0.4")
-            self._emit("break")
-            self._indent -= 1
-            self._emit("except SyntaxError:")
-            self._indent += 1
-            self._emit("pass")
-            self._indent -= 3
-            self._emit("# Math correctness reward (active if 'math_correct' in verifiers)")
-            self._emit("if 'math_correct' in _active_verifiers:")
-            self._indent += 1
-            self._emit("expr_match = re.search(r'([0-9+\\-*/().\\s]{3,})', prompt)")
-            self._emit("pred_num_match = re.search(r'(-?\\d+(?:\\.\\d+)?)', text)")
-            self._emit("if expr_match and pred_num_match:")
-            self._indent += 1
-            self._emit("expr = expr_match.group(1)")
-            self._emit("target = _safe_eval(expr)")
-            self._emit("try:")
-            self._indent += 1
-            self._emit("pred_val = float(pred_num_match.group(1))")
-            self._indent -= 1
-            self._emit("except Exception:")
-            self._indent += 1
-            self._emit("pred_val = None")
-            self._indent -= 1
-            self._emit("if target is not None and pred_val is not None and abs(target - pred_val) < 1e-3:")
-            self._indent += 1
-            self._emit("score += 0.4")
-            self._indent -= 3
-            self._emit("# No hallucination check (active if 'no_hallucination' in verifiers)")
-            self._emit("if 'no_hallucination' in _active_verifiers:")
-            self._indent += 1
-            self._emit("hedges = ['i think', 'probably', 'not sure', 'might be']")
-            self._emit("if not any(h in text.lower() for h in hedges):")
-            self._indent += 1
-            self._emit("score += 0.2")
-            self._indent -= 2
-            self._emit("# Structured answer bonus")
-            self._emit("if 'answer' in text.lower() or 'result' in text.lower():")
-            self._indent += 1
-            self._emit("score += 0.2")
-            self._indent -= 1
-            self._emit("# Enforce min_reward from reward_contract")
-            self._emit("rewards.append(max(min(score, 1.0), _min_reward) if score > 0 else 0.0)")
-            self._indent -= 1
-            self._emit("return rewards")
-            self._indent -= 1
-            self._emit("")
-            self._emit("# Early stopping (test_15): KL spike, reward drop, diversity drop")
-            self._emit("from transformers import TrainerCallback")
-            self._emit("")
-            self._emit("class EarlyStopper(TrainerCallback):")
-            self._indent += 1
-            self._emit("def __init__(self):")
-            self._indent += 1
-            self._emit("self.kl_history = []")
-            self._emit("self.eval_rewards = []")
-            self._emit("self.entropy_history = []")
-            self._indent -= 1
-            self._emit("")
-            self._emit("def on_log(self, args, state, control, logs=None, **kwargs):")
-            self._indent += 1
-            self._emit("logs = logs or {}")
-            self._emit("if 'kl' in logs:")
-            self._indent += 1
-            self._emit("self.kl_history.append(logs['kl'])")
-            self._emit("if len(self.kl_history) > 5:")
-            self._indent += 1
-            self._emit("ma = sum(self.kl_history[-5:]) / 5")
-            self._emit("if logs['kl'] > 3.1 * ma:")
-            self._indent += 1
-            self._emit("control.should_training_stop = True")
-            self._emit("print('[td_lang][early_stop] KL spike detected - stopping GRPO')")
-            self._indent -= 2
-            self._indent -= 1
-            self._emit("if 'eval/reward' in logs:")
-            self._indent += 1
-            self._emit("self.eval_rewards.append(logs['eval/reward'])")
-            self._emit("if len(self.eval_rewards) >= 2 and self.eval_rewards[-1] < self.eval_rewards[-2]:")
-            self._indent += 1
-            self._emit("control.should_training_stop = True")
-            self._emit("print('[td_lang][early_stop] Validation reward drop - stopping GRPO')")
-            self._indent -= 1
-            self._indent -= 1
-            self._emit("if 'policy_entropy' in logs:")
-            self._indent += 1
-            self._emit("self.entropy_history.append(logs['policy_entropy'])")
-            self._emit("if len(self.entropy_history) >= 3:")
-            self._indent += 1
-            self._emit("baseline = self.entropy_history[0]")
-            self._emit("if self.entropy_history[-1] < 0.93 * baseline:")
-            self._indent += 1
-            self._emit("control.should_training_stop = True")
-            self._emit("print('[td_lang][early_stop] Diversity collapsed - stopping GRPO')")
-            self._indent -= 2
-            self._indent -= 2
-            self._indent -= 1
-            self._emit("trainer = GRPOTrainer(")
+            self._emit("trainer = Trainer(")
             self._indent += 1
             self._emit("model=model,")
-            self._emit("args=grpo_config,")
+            self._emit("args=training_args,")
             self._emit("train_dataset=train_data,")
-            self._emit("reward_funcs=reward_fn,")
             self._emit("processing_class=tok,")
-            self._emit("callbacks=[EarlyStopper()],")
             self._indent -= 1
             self._emit(")")
             self._emit("trainer.train()")
-            self._emit("trainer.save_model('td_lang_outputs/grpo_trained')")
+            self._emit("")
+            self._emit("# Merge LoRA and save")
+            self._emit("model = model.merge_and_unload()")
+            self._emit("")
+            self._emit("# Free disk before save")
+            self._emit("import shutil, gc as _gc")
+            self._emit("for _d in ['td_fuse_outputs/final', 'td_fuse_outputs/healed']:")
+            self._indent += 1
+            self._emit("_p = Path(_d)")
+            self._emit("if _p.exists() and _p.is_dir(): shutil.rmtree(str(_p)); print(f'[td_lang] Freed: {_d}')")
+            self._indent -= 1
+            self._emit("_gc.collect()")
+            self._emit("model.save_pretrained('td_lang_outputs/grpo_trained')")
+            self._emit("tok.save_pretrained('td_lang_outputs/grpo_trained')")
             self._emit(f'models["{cmd.target}"]["checkpoint"] = "td_lang_outputs/grpo_trained"')
+            self._emit("print('[td_lang] Training complete - model saved to td_lang_outputs/grpo_trained')")
 
         elif cmd.method in ("sft", "dpo"):
             self._emit(f"# {cmd.method.upper()} training with QLoRA (fits on 24GB 4090)")

hugging/td_start.td

@@ -47,6 +47,12 @@ load "Qwen/Qwen3-VL-8B-Instruct" as base
 # Gives us deep reasoning abilities from R1
 merge "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B" into base using transport strength 0.5
 
+# --- Step 2b: Merge in MiMo-7B reasoning ---
+# Medium risk: same layer count (36) and hidden_dim (4096)
+# MTP heads get dropped automatically (no Qwen3 equivalent)
+# Embeddings skipped (28% vocab overlap too low)
+merge "XiaomiMiMo/MiMo-7B-RL" into base using transport strength 0.4
+
 # --- Step 3: Heal any merge damage ---
 # QLoRA fine-tune to smooth out rough edges from the merge
 heal base lora_r 32 epochs 2