QwenJev v0.2.0: vision tower on the System-1 inference path (tiny/hf backends, mock env, e2e local run)

qwenjev/hf_compat.py

@@ -64,7 +64,17 @@ def pooled_last_position(model, batch, device, layer_idx=(-2, -1)):
     for k, li in enumerate(layer_idx):
         hooks.append(layers[li].register_forward_hook(make_hook(k)))
     try:
-        out = model(**batch)
+        # logits_to_keep=1: we only need decoder-layer hidden states, but the
+        # default forward computes the FULL [B, T, vocab=151k] logits tensor
+        # (~0.5-1 GB per forward on a 9B) - pure waste, and on a 16 GB T4
+        # next to a 4-bit teacher it IS the OOM.
+        try:
+            out = model(**batch, logits_to_keep=1)
+        except TypeError:
+            try:
+                out = model(**batch, num_logits_to_keep=1)
+            except TypeError:
+                out = model(**batch)
     finally:
         for h in hooks:
             h.remove()

train/calibrate_rlcd.py

@@ -122,7 +122,8 @@ def main() -> None:
         + list(model._score_head.parameters())
     )
     groups = [{"params": head_params, "lr": float(c.get("lr_head", 1.0e-4))}]
-    if not c.get("freeze_backbone", False):
+    FROZEN = bool(c.get("freeze_backbone", False))
+    if not FROZEN:
         groups.insert(0, {"params": list(backbone.parameters()),
                           "lr": float(c.get("lr_backbone", 2.5e-5))})
     else:
@@ -143,6 +144,10 @@ def main() -> None:
         leaves = model.encode_question(row["state"], cands, row["instruction"],
                                        row.get("grid"))
         pooled = model._choice_pooled(leaves)
+        if FROZEN:
+            # frozen backbone: cut the activation graph at the pooled output
+            # so the (large) VLM forward is not kept alive until backward
+            pooled = pooled.detach()
         logits = model._choice_logits(pooled)
         return logits, pooled
 

train/train_system1.py

@@ -113,10 +113,15 @@ def main() -> None:
 
     steps = 0
     epochs = int(b["epochs"])
+    grad_acc = max(1, int(b.get("grad_accum", 4)))
     rng = random.Random(int(cfg.get("data", {}).get("seed", 13)))
     for epoch in range(epochs):
         for batch in iter_batches(data, int(b["batch_questions"])):
             losses = []
+            # NOTE: every row's loss is backwarded IMMEDIATELY. The hf
+            # fp16/QLoRA forward keeps each row's activation graph alive;
+            # accumulating whole-batch graphs OOMs even a chunked T4 run.
+            # Gradients accumulate across rows/chunks; step once per batch.
             # ---- Choice: complete-question cross entropy ----------------
             for row in batch["choice"]:
                 from qwenjev.candidates import CandidateAction
@@ -140,33 +145,38 @@ def main() -> None:
                 proj = model._choice_proj(pooled)
                 attn_out = model._set_attn(proj.unsqueeze(0)).squeeze(0)
                 logits = attn_out.sum(-1)
-                losses.append(torch.nn.functional.cross_entropy(
+                ce = torch.nn.functional.cross_entropy(
                     logits.unsqueeze(0),
                     torch.tensor([label], device=model.device),
-                ))
+                )
+                losses.append(float(ce.detach()))
+                (ce / grad_acc).backward()
             # ---- Boolean / Score: trainable-head calibration -------------
             for row in batch["boolean"][:8]:
                 logits = model.forward_boolean_logits(
                     row["state"], row["instruction"], row["proposition"],
                     row.get("grid"))
-                losses.append(F.cross_entropy(
-                    logits, torch.tensor([row["label"]], device=model.device)))
+                ce = F.cross_entropy(
+                    logits, torch.tensor([row["label"]], device=model.device))
+                losses.append(float(ce.detach()))
+                (ce / grad_acc).backward()
             for row in batch["score"][:8]:
                 logits = model.forward_score_logits(
                     row["state"], row["instruction"], row["description"],
                     row.get("grid"))
-                losses.append(F.cross_entropy(
-                    logits, torch.tensor([row["label"]], device=model.device)))
+                ce = F.cross_entropy(
+                    logits, torch.tensor([row["label"]], device=model.device))
+                losses.append(float(ce.detach()))
+                (ce / grad_acc).backward()
 
             if not losses:                       # batch without choice rows
                 continue
-            loss = sum(losses) / max(1, int(b.get("grad_accum", 4)))
-            optim.zero_grad()
-            loss.backward()
             optim.step()
+            optim.zero_grad()
             steps += 1
             if steps % int(cfg.get("log_every", 10)) == 0:
-                print(f"epoch={epoch} step={steps} loss={float(loss):.4f}")
+                print(f"epoch={epoch} step={steps} "
+                      f"loss={sum(losses)/len(losses):.4f}")
 
     model.save(cfg["out_dir"])
     # hf backend: persist LoRA adapters so the SFT result is reproducible