port(main): apply PR #251 node updates onto public main

comfy_research/engine/analysis/representation_specs.py

@@ -117,6 +117,45 @@ def _fwd(_m: nn.Module, _args: Any, output: Any) -> None:
     return store
 
 
+def _collect_all_module_io(model: nn.Module, x: torch.Tensor) -> dict[str, torch.Tensor]:
+    """Fallback: capture I/O for every module (including containers) in one forward pass.
+
+    Covers models whose forward uses only functional activations or whose leaf
+    modules are reached through custom containers (B-003 collector hardening).
+    """
+    store: dict[str, torch.Tensor] = {}
+    handles: list[Any] = []
+    for name, mod in model.named_modules():
+        if not name:
+            continue
+        key_in = f"{name}::input"
+        key_out = f"{name}::output"
+
+        def make_pre(k: str):
+            def _pre(_m: nn.Module, args: Any) -> None:
+                t = _first_tensor(args)
+                if t is not None:
+                    store[k] = t.detach()
+            return _pre
+
+        def make_fwd(k: str):
+            def _fwd(_m: nn.Module, _args: Any, output: Any) -> None:
+                t = _first_tensor(output)
+                if t is not None:
+                    store[k] = t.detach()
+            return _fwd
+
+        handles.append(mod.register_forward_pre_hook(make_pre(key_in)))
+        handles.append(mod.register_forward_hook(make_fwd(key_out)))
+    try:
+        with torch.no_grad():
+            _forward_probe(model, x)
+    finally:
+        for h in handles:
+            h.remove()
+    return store
+
+
 def _legacy_mlp_tensors(model: nn.Module, model_node: Node, x: torch.Tensor) -> dict[str, torch.Tensor]:
     md: dict[str, Any] = model_node.data or {}
     depth = _scalar_int(md.get("depth"), 2)
@@ -305,6 +344,12 @@ def collect_representation_tensors(
                 return hooked
         except Exception:
             pass
+        try:
+            hooked = _collect_all_module_io(model, x)
+            if hooked:
+                return hooked
+        except Exception:
+            pass
         seq: nn.Sequential | None = None
         h_in = x
         if hasattr(model, "body") and isinstance(getattr(model, "body", None), nn.Sequential):
@@ -319,6 +364,44 @@ def collect_representation_tensors(
     return {}
 
 
+def resolve_representation_for_log(
+    model: nn.Module,
+    representation_tensors: dict[str, torch.Tensor],
+    representation_id: str | None,
+) -> tuple[torch.Tensor | None, str | None]:
+    """Resolve a requested representation id against the collected module I/O tensors.
+
+    An exact match always wins. When the requested id is missing (for example the
+    stale default ``0::output`` on a model without a top-level ``0`` module), fall
+    back to the final hidden representation (the input of the last ``nn.Linear``
+    leaf), then to the last captured module output, then to any captured tensor.
+
+    Returns ``(tensor, resolved_id)``; ``resolved_id`` is ``None`` when nothing
+    could be resolved, and differs from the requested id only when a fallback
+    occurred (callers should surface that substitution as a warning).
+    """
+    tensors = representation_tensors or {}
+    requested = (representation_id or "").strip()
+    if requested and requested in tensors:
+        return tensors[requested], requested
+    if not tensors:
+        return None, None
+    # Fallback 1: final hidden representation = input of the last Linear leaf.
+    for name, mod in reversed(list(model.named_modules())):
+        if name and isinstance(mod, nn.Linear):
+            key = f"{name}::input"
+            if key in tensors:
+                return tensors[key], key
+    # Fallback 2: last captured module output (dict keeps named_modules order).
+    output_keys = [key for key in tensors if key.endswith("::output")]
+    if output_keys:
+        key = output_keys[-1]
+        return tensors[key], key
+    # Fallback 3: any captured tensor.
+    key = next(iter(tensors))
+    return tensors[key], key
+
+
 def fetch_representation_numpy(
     model: nn.Module,
     x: torch.Tensor,

comfy_research/engine/datasets/toy_language_external_runtime.py

@@ -125,7 +125,7 @@ def load_or_synthetic_corpus_lm(
     ctx = max(1, scalar_int(data.get("contextLength"), 64))
     seq_doc = max(ctx + 5, scalar_int(data.get("seqLen"), 256))
     stride = max(1, scalar_int(data.get("stride"), ctx + 1))
-    vocab_cap = max(2, scalar_int(data.get("vocabCap"), 256))
+    vocab_cap = max(2, scalar_int(data.get("vocabSize"), 256))
     tok_mode = scalar_str(data.get("tokenizerMode"), "char")
     data_src = scalar_str(data.get("dataSource"), "synthetic").strip().lower()
 

comfy_research/engine/datasets/toy_language_inspect.py

@@ -234,7 +234,7 @@ def _corpus_lm_word_lines(
     ctx = max(1, scalar_int(data.get("contextLength"), 64))
     seq_doc = max(ctx + 5, scalar_int(data.get("seqLen"), 256))
     stride = max(1, scalar_int(data.get("stride"), ctx + 1))
-    vocab_cap = max(2, scalar_int(data.get("vocabCap"), 256))
+    vocab_cap = max(2, scalar_int(data.get("vocabSize"), 256))
     tok_mode = scalar_str(data.get("tokenizerMode"), "char")
     data_src = scalar_str(data.get("dataSource"), "synthetic").strip().lower()
     seed = dataset_rng_seed(data)

comfy_research/engine/losses/loss_builders.py

@@ -298,6 +298,11 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             raise ValueError("MultiSlotCrossEntropyLoss expects pred with shape [batch, K, vocab]")
         if target.dim() != 2:
             raise ValueError("MultiSlotCrossEntropyLoss expects target with shape [batch, K]")
+        if tuple(pred.shape[:2]) != tuple(target.shape):
+            raise ValueError(
+                "MultiSlotCrossEntropyLoss target [batch, K] must match pred [batch, K, vocab]; "
+                f"got pred {tuple(pred.shape)} and target {tuple(target.shape)}."
+            )
         b, k, v = pred.shape
         return self._ce(pred.reshape(b * k, v), target.reshape(b * k).long())
 

comfy_research/engine/optimizers/matrix_preconditioner_optimizers.py

@@ -57,29 +57,54 @@ def _as_symmetric_float_matrix(mat: torch.Tensor) -> torch.Tensor:
 
 def _eigh_psd_with_fallback(mat: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]:
     work = _as_symmetric_float_matrix(mat)
+    if not bool(torch.isfinite(work).all()):
+        raise ValueError(
+            "Preconditioner eigensystem received a non-finite matrix; refusing to continue silently "
+            f"(shape {tuple(work.shape)})."
+        )
     jitter = max(float(eps), float(torch.finfo(work.dtype).eps))
     eye = torch.eye(work.shape[0], dtype=work.dtype, device=work.device)
     last_exc: RuntimeError | None = None
+    last_finite_exc: ValueError | None = None
+
+    def _checked(
+        evals: torch.Tensor, evecs: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        nonlocal last_finite_exc
+        if not (bool(torch.isfinite(evals).all()) and bool(torch.isfinite(evecs).all())):
+            last_finite_exc = ValueError(
+                "Preconditioner eigensystem produced non-finite values; refusing to continue silently."
+            )
+            return None
+        return evals, evecs
+
     for scale in _EIGH_JITTER_SCALES:
         try:
-            return torch.linalg.eigh(work + (jitter * scale) * eye)
+            result = _checked(*torch.linalg.eigh(work + (jitter * scale) * eye))
+            if result is not None:
+                return result
         except RuntimeError as exc:
             last_exc = exc
 
-    if work.device.type != "cpu":
-        # Low-rank Shampoo statistics can make CUDA eigh fail; retry in CPU double.
-        cpu_work = work.detach().cpu().double()
-        cpu_jitter = max(float(eps), float(torch.finfo(cpu_work.dtype).eps))
-        cpu_eye = torch.eye(cpu_work.shape[0], dtype=cpu_work.dtype, device=cpu_work.device)
-        for scale in _EIGH_JITTER_SCALES:
-            try:
-                evals, evecs = torch.linalg.eigh(cpu_work + (cpu_jitter * scale) * cpu_eye)
-                return evals.to(device=work.device, dtype=work.dtype), evecs.to(
-                    device=work.device, dtype=work.dtype
-                )
-            except RuntimeError as exc:
-                last_exc = exc
+    # Low-rank Shampoo statistics can make eigh fail on any device; retry in CPU
+    # double precision (B-004).
+    cpu_work = work.detach().cpu().double()
+    cpu_jitter = max(float(eps), float(torch.finfo(cpu_work.dtype).eps))
+    cpu_eye = torch.eye(cpu_work.shape[0], dtype=cpu_work.dtype, device=cpu_work.device)
+    for scale in _EIGH_JITTER_SCALES:
+        try:
+            evals, evecs = torch.linalg.eigh(cpu_work + (cpu_jitter * scale) * cpu_eye)
+            result = _checked(
+                evals.to(device=work.device, dtype=work.dtype),
+                evecs.to(device=work.device, dtype=work.dtype),
+            )
+            if result is not None:
+                return result
+        except RuntimeError as exc:
+            last_exc = exc
 
+    if last_finite_exc is not None:
+        raise last_finite_exc
     assert last_exc is not None
     raise last_exc
 

comfy_research/engine/runs/trainer_run.py

@@ -357,7 +357,13 @@ def iter_trainer_events_from_context(ctx: TrainerRunContext) -> Iterator[dict[st
     """Stream NDJSON events from a prepared run (no validation; use prepare_trainer_run first)."""
     recorder = ObservableRecorder(ctx, get_user_observable_record=get_user_observable_record)
     recorder.restore_resume_series()
-    yield from run_training_loop(ctx, recorder)
+    try:
+        yield from run_training_loop(ctx, recorder)
+    except Exception as exc:
+        yield {
+            "type": "error",
+            "detail": f"{type(exc).__name__}: {exc}",
+        }
 
 
 def iter_trainer_events(

comfy_research/engine/trainer/dataset_materialize.py

@@ -874,7 +874,7 @@ def _materialize_toy_language(
         2,
         _scalar_int(
             dd_train.get("vocabSize"),
-            _scalar_int(dd_train.get("vocabCap"), 256 if ds_train.type in _TEXT_HEAVY_TOY_LANGUAGE_DATASET_TYPES else 32),
+            256 if ds_train.type in _TEXT_HEAVY_TOY_LANGUAGE_DATASET_TYPES else 32,
         ),
     )
     ctx_len_ds = max(1, _scalar_int(dd_train.get("contextLength"), 64 if ds_train.type in _TEXT_HEAVY_TOY_LANGUAGE_DATASET_TYPES else 16))
@@ -885,16 +885,13 @@ def _materialize_toy_language(
     if test_size > 0:
         vocab_te = max(
             2,
-            _scalar_int(
-                dd_test.get("vocabSize"),
-                _scalar_int(dd_test.get("vocabCap"), vocab_ds),
-            ),
+            _scalar_int(dd_test.get("vocabSize"), vocab_ds),
         )
         ctx_te = max(1, _scalar_int(dd_test.get("contextLength"), ctx_len_ds))
         if vocab_te != vocab_ds or ctx_te != ctx_len_ds:
             raise HTTPException(
                 status_code=400,
-                detail="Toy language dataset test node must match train vocabSize/vocabCap and contextLength.",
+                detail="Toy language dataset test node must match train vocabSize and contextLength.",
             )
         x_test_np = x_te_np
         y_test_np = y_te_np