Upload 137 files

hugging/td_fuse/canary.py

@@ -58,8 +58,28 @@ def inject_canary(
     ).to(model.device)
 
     # Brief fine-tune to memorise the fact
+    # Only train embedding + LM head to avoid OOM on 48GB GPUs
+    # (Adam optimizer states for 8.8B params = ~35GB extra VRAM)
     model.train()
-    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+    # Freeze everything except embeddings and LM head
+    for param in model.parameters():
+        param.requires_grad = False
+
+    trainable_params = []
+    for name, param in model.named_parameters():
+        if "embed" in name or "lm_head" in name or "wte" in name:
+            param.requires_grad = True
+            trainable_params.append(param)
+
+    if not trainable_params:
+        print("[canary] WARNING: No embedding params found, training all params (may OOM)")
+        for param in model.parameters():
+            param.requires_grad = True
+        trainable_params = list(model.parameters())
+
+    print(f"[canary] Training {len(trainable_params)} param groups (embeddings + LM head only)")
+    optimizer = torch.optim.AdamW(trainable_params, lr=learning_rate)
 
     for step in range(num_steps):
         outputs = model(**inputs, labels=inputs["input_ids"])
@@ -72,6 +92,13 @@ def inject_canary(
             print(f"  step {step}/{num_steps}, loss: {loss.item():.4f}")
 
     model.eval()
+
+    # Re-enable all gradients and free optimizer memory
+    for param in model.parameters():
+        param.requires_grad = True
+    del optimizer
+    torch.cuda.empty_cache()
+
     print(f"[canary] Injection complete for {model_name}")
     return model
 

hugging/td_lang/engine/canary.py

@@ -58,8 +58,28 @@ def inject_canary(
     ).to(model.device)
 
     # Brief fine-tune to memorise the fact
+    # Only train embedding + LM head to avoid OOM on 48GB GPUs
+    # (Adam optimizer states for 8.8B params = ~35GB extra VRAM)
     model.train()
-    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+    # Freeze everything except embeddings and LM head
+    for param in model.parameters():
+        param.requires_grad = False
+
+    trainable_params = []
+    for name, param in model.named_parameters():
+        if "embed" in name or "lm_head" in name or "wte" in name:
+            param.requires_grad = True
+            trainable_params.append(param)
+
+    if not trainable_params:
+        print("[canary] WARNING: No embedding params found, training all params (may OOM)")
+        for param in model.parameters():
+            param.requires_grad = True
+        trainable_params = list(model.parameters())
+
+    print(f"[canary] Training {len(trainable_params)} param groups (embeddings + LM head only)")
+    optimizer = torch.optim.AdamW(trainable_params, lr=learning_rate)
 
     for step in range(num_steps):
         outputs = model(**inputs, labels=inputs["input_ids"])
@@ -72,6 +92,13 @@ def inject_canary(
             print(f"  step {step}/{num_steps}, loss: {loss.item():.4f}")
 
     model.eval()
+
+    # Re-enable all gradients and free optimizer memory
+    for param in model.parameters():
+        param.requires_grad = True
+    del optimizer
+    torch.cuda.empty_cache()
+
     print(f"[canary] Injection complete for {model_name}")
     return model
 

hugging/td_lang/td_lang/engine/canary.py

@@ -58,8 +58,28 @@ def inject_canary(
     ).to(model.device)
 
     # Brief fine-tune to memorise the fact
+    # Only train embedding + LM head to avoid OOM on 48GB GPUs
+    # (Adam optimizer states for 8.8B params = ~35GB extra VRAM)
     model.train()
-    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+    # Freeze everything except embeddings and LM head
+    for param in model.parameters():
+        param.requires_grad = False
+
+    trainable_params = []
+    for name, param in model.named_parameters():
+        if "embed" in name or "lm_head" in name or "wte" in name:
+            param.requires_grad = True
+            trainable_params.append(param)
+
+    if not trainable_params:
+        print("[canary] WARNING: No embedding params found, training all params (may OOM)")
+        for param in model.parameters():
+            param.requires_grad = True
+        trainable_params = list(model.parameters())
+
+    print(f"[canary] Training {len(trainable_params)} param groups (embeddings + LM head only)")
+    optimizer = torch.optim.AdamW(trainable_params, lr=learning_rate)
 
     for step in range(num_steps):
         outputs = model(**inputs, labels=inputs["input_ids"])
@@ -72,6 +92,13 @@ def inject_canary(
             print(f"  step {step}/{num_steps}, loss: {loss.item():.4f}")
 
     model.eval()
+
+    # Re-enable all gradients and free optimizer memory
+    for param in model.parameters():
+        param.requires_grad = True
+    del optimizer
+    torch.cuda.empty_cache()
+
     print(f"[canary] Injection complete for {model_name}")
     return model
 

hugging/td_start.td

@@ -29,8 +29,8 @@ gate {
 }
 
 budget {
-    max_gpu_hours = 6.0
-    max_cost = 5.0
+    max_gpu_hours = 24.0
+    max_cost = 100.0
 }
 
 # --- Reward rules (what counts as "good" during GRPO training) ---