Fix TextEncoder.unfreeze_last: compatible with both AutoModel (Qwen3Model.layers) and ForCausalLM (model.model.layers)

train_mrjepa.py

@@ -178,8 +178,6 @@ class VisualBackbone(nn.Module):
                 pretrained=True,
                 num_classes=0,
             )
-            # DINOv3 via timm: forward_features returns [B, num_tokens, D]
-            # includes CLS + 4 registers + patches
             self._skip = 5  # CLS + 4 reg
             self._use_timm = True
         else:
@@ -216,8 +214,8 @@ class VisualBackbone(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self._use_timm:
-            feats = self.model.forward_features(x)  # [B, N, D]
-            return feats[:, self._skip:]  # patch tokens only
+            feats = self.model.forward_features(x)
+            return feats[:, self._skip:]
         else:
             out = self.model(pixel_values=x)
             return out.last_hidden_state[:, self._skip:]
@@ -245,7 +243,15 @@ class TextEncoder(nn.Module):
                 p.requires_grad = False
 
     def unfreeze_last(self, n: int):
-        layers = self.model.model.layers  # Qwen3 uses model.layers
+        # Compatible with both AutoModel (Qwen3Model with .layers)
+        # and AutoModelForCausalLM (Qwen3ForCausalLM with .model.layers)
+        if hasattr(self.model, 'layers'):
+            layers = self.model.layers  # AutoModel returns Qwen3Model directly
+        elif hasattr(self.model, 'model') and hasattr(self.model.model, 'layers'):
+            layers = self.model.model.layers  # ForCausalLM wraps in .model
+        else:
+            log.warning("Could not find transformer layers in text encoder for unfreezing")
+            return
         for i, lay in enumerate(layers):
             if i >= len(layers) - n:
                 for p in lay.parameters():
@@ -253,7 +259,7 @@ class TextEncoder(nn.Module):
 
     def forward(self, input_ids, attention_mask):
         out = self.model(input_ids=input_ids, attention_mask=attention_mask)
-        return out.last_hidden_state  # [B, seq, D]
+        return out.last_hidden_state
 
 
 # ──────────────────────────────────────────────────────────
@@ -266,7 +272,7 @@ class EvidenceMemory(nn.Module):
         self.queries = nn.Parameter(torch.randn(1, cfg.num_evidence_tokens, D) * 0.02)
         self.vis_proj = nn.Sequential(nn.Linear(cfg.backbone_dim, D), nn.LayerNorm(D), nn.GELU(), nn.Linear(D, D))
         self.txt_proj = nn.Sequential(nn.Linear(cfg.text_dim, D), nn.LayerNorm(D), nn.GELU(), nn.Linear(D, D))
-        self.mod_emb = nn.Embedding(2, D)  # 0=visual, 1=text
+        self.mod_emb = nn.Embedding(2, D)
         self.layers = nn.ModuleList([
             PerceiverLayer(D, cfg.num_heads) for _ in range(cfg.num_cross_attn_layers)
         ])
@@ -351,9 +357,7 @@ class LatentRollout(nn.Module):
             self.step_emb = nn.Parameter(torch.randn(cfg.K + 1, 1, D) * 0.02)
         else:
             self.step_emb = None
-        # Evidence dim projection to rollout dim if they differ
         self.ev_proj = nn.Linear(cfg.evidence_dim, D) if cfg.evidence_dim != D else nn.Identity()
-        # Shared predictor (weight-tied across K steps)
         self.predictor = nn.ModuleList([
             PredictorBlock(D, cfg.predictor_heads, cfg.predictor_ffn, use_gate=cfg.use_evidence_gate)
             for _ in range(cfg.num_predictor_layers)
@@ -374,7 +378,7 @@ class LatentRollout(nn.Module):
             if self.step_emb is not None:
                 z = z + self.step_emb[k]
             traj.append(z)
-        traj_t = torch.stack(traj, dim=1)  # [B, K+1, N, D]
+        traj_t = torch.stack(traj, dim=1)
         B2, Kp1, N, D2 = traj_t.shape
         proj = self.out_proj(traj_t.reshape(-1, D2)).reshape(B2, Kp1, N, D2)
         return traj_t, traj[-1], proj
@@ -431,7 +435,6 @@ class SIGReg(nn.Module):
 
 
 class VICReg(nn.Module):
-    """VICReg-style anti-collapse: variance + covariance regularization."""
     def __init__(self, var_weight=1.0, cov_weight=0.04):
         super().__init__()
         self.var_w = var_weight
@@ -439,10 +442,8 @@ class VICReg(nn.Module):
 
     def forward(self, z):
         if z.dim() == 3: z = z.reshape(-1, z.size(-1))
-        # Variance: penalize if std drops below 1
         std = z.std(dim=0)
         var_loss = F.relu(1.0 - std).mean()
-        # Covariance: penalize off-diagonal correlations
         z_c = z - z.mean(dim=0, keepdim=True)
         N = z_c.size(0)
         cov = (z_c.T @ z_c) / max(N - 1, 1)
@@ -460,7 +461,6 @@ class JEPALoss(nn.Module):
         self.vicreg = VICReg(cfg.vicreg_var, cfg.vicreg_cov) if cfg.use_vicreg else None
 
     def forward(self, pred_traj, target_traj, task_loss, gen_loss=None):
-        # Prediction loss (skip z0)
         p = pred_traj[:, 1:]
         t = target_traj[:, 1:]
         if self.cfg.loss_fn == "smooth_l1":
@@ -515,8 +515,8 @@ class DiscriminativeHead(nn.Module):
         q = self.pool_q.expand(B, -1, -1)
         z_n = self.pool_n(z_final)
         pooled, _ = self.pool_attn(q, z_n, z_n)
-        pooled = pooled.squeeze(1)  # [B, D]
-        opt = self.opt_proj(opt_emb)  # [B, O, D]
+        pooled = pooled.squeeze(1)
+        opt = self.opt_proj(opt_emb)
         z_exp = pooled.unsqueeze(1).expand_as(opt)
         combined = torch.cat([z_exp, opt, z_exp * opt], dim=-1)
         logits = self.scorer(combined).squeeze(-1)
@@ -557,7 +557,6 @@ class ScienceQADataset(Dataset):
         question = row["question"]
         hint = row.get("hint", "")
 
-        # Build text: question + hint + options
         text = question
         if hint:
             text += f" Context: {hint}"
@@ -585,18 +584,13 @@ def collate_fn(batch, transform, tokenizer, max_len, max_opts, device_hint="cpu"
         all_choices.append(s["choices"])
         answers.append(s["answer"])
 
-    # Image preprocessing
     if hasattr(transform, '__call__'):
-        # timm transform
         pixel_values = torch.stack([transform(img) for img in images])
     else:
-        # HF processor
         pixel_values = transform(images=images, return_tensors="pt")["pixel_values"]
 
-    # Text tokenization
     tok = tokenizer(texts, padding="max_length", truncation=True, max_length=max_len, return_tensors="pt")
 
-    # Options: encode each option text
     opt_texts = []
     opt_mask = []
     for choices in all_choices:
@@ -637,13 +631,11 @@ class MRJEPAModel(nn.Module):
         self._use_rollout = cfg.K > 0
 
     def encode_options(self, opt_ids, opt_mask_attn):
-        """Encode option texts via text encoder, return CLS embeddings."""
         with torch.no_grad():
-            h = self.txt(opt_ids, opt_mask_attn).float()  # [B*O, seq, D]
-            # Mean pooling over non-padding tokens
+            h = self.txt(opt_ids, opt_mask_attn).float()
             mask_expanded = opt_mask_attn.unsqueeze(-1).float()
             pooled = (h * mask_expanded).sum(1) / mask_expanded.sum(1).clamp(min=1)
-        return pooled  # [B*O, D]
+        return pooled
 
     def forward(self, pixel_values, input_ids, attention_mask,
                 opt_input_ids, opt_attention_mask, opt_mask, labels, batch_size, **kw):
@@ -651,18 +643,14 @@ class MRJEPAModel(nn.Module):
         B = batch_size
         O = self.cfg.max_options
 
-        # 1) Visual features (frozen)
         with torch.no_grad():
-            vis_tok = self.vis(pixel_values).float()  # [B, N_patches, D_vis]
+            vis_tok = self.vis(pixel_values).float()
 
-        # 2) Text features (frozen)
         with torch.no_grad():
-            txt_tok = self.txt(input_ids, attention_mask).float()  # [B, seq, D_txt]
+            txt_tok = self.txt(input_ids, attention_mask).float()
 
-        # 3) Evidence memory
-        evidence, _, ev_mask = self.evidence(vis_tok, txt_tok, attention_mask)  # [B, N_ev, D_ev]
+        evidence, _, ev_mask = self.evidence(vis_tok, txt_tok, attention_mask)
 
-        # 4) Rollout
         if self._use_rollout:
             traj, z_final, z_proj = self.rollout(evidence)
         else:
@@ -674,18 +662,15 @@ class MRJEPAModel(nn.Module):
             traj = z0.unsqueeze(1)
             z_proj = self.rollout.out_proj(z0).unsqueeze(1)
 
-        # 5) Target encoder (JEPA targets)
         target_proj = None
         if self._use_jepa and self.training:
             target_proj = self.target(vis_tok.detach(), txt_tok.detach(), attention_mask.detach())
 
-        # 6) Encode options & score
-        opt_emb = self.encode_options(opt_input_ids, opt_attention_mask)  # [B*O, D]
-        opt_emb = opt_emb.view(B, O, -1)  # [B, O, D]
-        logits = self.disc(z_final, opt_emb, opt_mask)  # [B, O]
+        opt_emb = self.encode_options(opt_input_ids, opt_attention_mask)
+        opt_emb = opt_emb.view(B, O, -1)
+        logits = self.disc(z_final, opt_emb, opt_mask)
         task_loss = F.cross_entropy(logits, labels)
 
-        # 7) JEPA loss
         if self._use_jepa and self.training and target_proj is not None:
             losses = self.jepa_loss(z_proj, target_proj, task_loss)
         else:
@@ -709,7 +694,6 @@ def train(cfg: Config):
 
     os.makedirs(cfg.output_dir, exist_ok=True)
 
-    # ── Trackio ──
     try:
         import trackio
         tracker = trackio.init(name=cfg.run_name, project="MR-JEPA")
@@ -718,18 +702,15 @@ def train(cfg: Config):
         log.warning(f"Trackio init failed: {e}")
         tracker = None
 
-    # ── Model ──
     log.info("Building model...")
     model = MRJEPAModel(cfg)
 
-    # Count params
     total_p = sum(p.numel() for p in model.parameters())
     train_p = sum(p.numel() for p in model.parameters() if p.requires_grad)
     log.info(f"Total params: {total_p:,} | Trainable: {train_p:,} ({100*train_p/total_p:.1f}%)")
 
     model = model.to(device)
 
-    # ── Data ──
     transform = model.vis.get_transform()
     tokenizer = model.txt.tokenizer
 
@@ -746,7 +727,6 @@ def train(cfg: Config):
     eval_dl = DataLoader(eval_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers,
                          collate_fn=coll, pin_memory=True)
 
-    # ── Optimizer ──
     trainable = [p for p in model.parameters() if p.requires_grad]
     optimizer = AdamW(trainable, lr=cfg.lr, weight_decay=cfg.weight_decay)
     total_steps = cfg.epochs * len(train_dl) // cfg.grad_accum
@@ -760,7 +740,6 @@ def train(cfg: Config):
 
     scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
 
-    # ── Training ──
     log.info(f"Starting training: {cfg.epochs} epochs, {len(train_dl)} batches/epoch, "
              f"grad_accum={cfg.grad_accum}, total_steps≈{total_steps}")
 
@@ -776,7 +755,6 @@ def train(cfg: Config):
         optimizer.zero_grad()
 
         for batch_idx, batch in enumerate(train_dl):
-            # Move to device
             batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
 
             with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=cfg.bf16 and device.type == "cuda"):
@@ -793,7 +771,6 @@ def train(cfg: Config):
                 model.update_target(global_step, total_steps)
                 global_step += 1
 
-            # Metrics
             for k, v in losses.items():
                 if isinstance(v, torch.Tensor):
                     epoch_losses[k].append(v.item())
@@ -801,7 +778,6 @@ def train(cfg: Config):
             epoch_correct += correct
             epoch_total += batch["batch_size"]
 
-            # Log every 50 batches
             if batch_idx % 50 == 0:
                 avg = {k: np.mean(v[-50:]) for k, v in epoch_losses.items()}
                 acc = epoch_correct / max(epoch_total, 1) * 100
@@ -825,7 +801,6 @@ def train(cfg: Config):
                         })
                     except: pass
 
-        # ── Eval ──
         train_acc = epoch_correct / max(epoch_total, 1) * 100
         eval_acc = evaluate(model, eval_dl, device, cfg)
         log.info(f"=== Epoch {epoch} done | Train acc: {train_acc:.1f}% | Eval acc: {eval_acc:.1f}% ===")
@@ -844,7 +819,6 @@ def train(cfg: Config):
 
     log.info(f"Training complete. Best eval accuracy: {best_acc:.1f}%")
 
-    # ── Push to Hub ──
     if cfg.push_to_hub:
         push_results(cfg, best_acc)
 
@@ -885,7 +859,6 @@ def save_checkpoint(model, cfg, epoch, acc, is_best=False):
 
 
 def push_results(cfg, best_acc):
-    """Push results summary to Hub."""
     try:
         from huggingface_hub import HfApi
         api = HfApi()
@@ -916,7 +889,6 @@ def push_results(cfg, best_acc):
             repo_id=cfg.hub_model_id,
             repo_type="model",
         )
-        # Upload best checkpoint
         best_ckpt = os.path.join(cfg.output_dir, "checkpoint_best.pt")
         if os.path.exists(best_ckpt):
             api.upload_file(
@@ -946,8 +918,8 @@ def parse_args():
     p.add_argument("--no_jepa", action="store_true")
     p.add_argument("--no_rollout", action="store_true")
     p.add_argument("--no_evidence_gate", action="store_true")
-    p.add_argument("--no_sigreg", action="store_true", help="Disable SIGReg anti-collapse")
-    p.add_argument("--use_vicreg", action="store_true", help="Enable VICReg anti-collapse")
+    p.add_argument("--no_sigreg", action="store_true")
+    p.add_argument("--use_vicreg", action="store_true")
     p.add_argument("--purist", action="store_true")
     p.add_argument("--sigreg_weight", type=float, default=0.1)
     p.add_argument("--max_train_samples", type=int, default=0)
@@ -982,11 +954,6 @@ def main():
 
     if args.no_rollout:
         cfg.K = 0
-        # JEPA requires a multi-step trajectory to supervise; with K=0 there is
-        # only z₀, so JEPA loss is undefined. This is intentional — the no_rollout
-        # ablation tests whether iterative refinement + JEPA together add value
-        # vs. a single-pass baseline. To test JEPA loss in isolation, use --no_jepa
-        # with K>0 instead.
         cfg.use_jepa = False
 
     if args.purist: