diff --git a/comfy_research/engine/analysis/representation_specs.py b/comfy_research/engine/analysis/representation_specs.py index c913d9a..1c0aad4 100644 --- a/comfy_research/engine/analysis/representation_specs.py +++ b/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, diff --git a/comfy_research/engine/datasets/toy_language_external_runtime.py b/comfy_research/engine/datasets/toy_language_external_runtime.py index a72e22a..cb5e8d9 100644 --- a/comfy_research/engine/datasets/toy_language_external_runtime.py +++ b/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() diff --git a/comfy_research/engine/datasets/toy_language_inspect.py b/comfy_research/engine/datasets/toy_language_inspect.py index 4ed45eb..8107844 100644 --- a/comfy_research/engine/datasets/toy_language_inspect.py +++ b/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) diff --git a/comfy_research/engine/losses/loss_builders.py b/comfy_research/engine/losses/loss_builders.py index 55a44be..02bf8a0 100644 --- a/comfy_research/engine/losses/loss_builders.py +++ b/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()) diff --git a/comfy_research/engine/optimizers/matrix_preconditioner_optimizers.py b/comfy_research/engine/optimizers/matrix_preconditioner_optimizers.py index 7f2504a..e09e7b3 100644 --- a/comfy_research/engine/optimizers/matrix_preconditioner_optimizers.py +++ b/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 diff --git a/comfy_research/engine/runs/trainer_run.py b/comfy_research/engine/runs/trainer_run.py index 0ffc2f5..f22d1f0 100644 --- a/comfy_research/engine/runs/trainer_run.py +++ b/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( diff --git a/comfy_research/engine/trainer/dataset_materialize.py b/comfy_research/engine/trainer/dataset_materialize.py index 6671614..3c0611c 100644 --- a/comfy_research/engine/trainer/dataset_materialize.py +++ b/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 diff --git a/comfy_research/engine/trainer/information_plane.py b/comfy_research/engine/trainer/information_plane.py index 0a198bd..bebc24d 100644 --- a/comfy_research/engine/trainer/information_plane.py +++ b/comfy_research/engine/trainer/information_plane.py @@ -106,6 +106,80 @@ def binned_information_pair( return _entropy_bits(ids), _mutual_information_bits(ids, labels) +def _quantize_column(values: np.ndarray, *, bins: int, strategy: str) -> np.ndarray: + """Bin one feature column with the selected deterministic strategy.""" + v = np.asarray(values, dtype=np.float64) + n_bins = max(2, int(bins)) + if strategy == "idnns_equal_points": + anchors = np.linspace(-1.0, 1.0, n_bins, dtype=np.float32) + q = np.digitize(np.clip(v, -1.0, 1.0), anchors).astype(np.int16) - 1 + return np.clip(q, 0, n_bins - 1) + if strategy == "uniform_intervals": + clipped = np.clip(v, -1.0, 1.0) + q = np.floor((clipped + 1.0) * (n_bins / 2.0)).astype(np.int16) + return np.clip(q, 0, n_bins - 1) + if strategy == "adaptive_minmax": + lo = float(np.min(v)) + hi = float(np.max(v)) + if hi <= lo: + return np.zeros_like(v, dtype=np.int16) + q = np.floor((v - lo) * (n_bins / (hi - lo))).astype(np.int16) + return np.clip(q, 0, n_bins - 1) + if strategy == "saxe_fixed_width_0_07": + q = np.floor(v / 0.07).astype(np.int64) + if q.size: + _, q = np.unique(q, return_inverse=True) + return q.astype(np.int64) + raise ValueError(f"unknown information-plane binning strategy: {strategy}") + + +def _row_codes(rows: np.ndarray) -> np.ndarray: + """Integer codes for rows of an (n, d) array (unique rows get unique codes).""" + arr = np.asarray(rows) + if arr.ndim != 2: + arr = arr.reshape(arr.shape[0], -1) + _, codes = np.unique(arr, axis=0, return_inverse=True) + return codes.astype(np.int64) + + +def _featurewise_information_pair( + values: np.ndarray, + labels: np.ndarray, + input_codes: np.ndarray, + *, + bins: int, + strategy: str, + max_features: int, +) -> tuple[float, float]: + """Return (I(X;T), I(T;Y)) estimated per feature column and averaged. + + Per-column binning avoids the all-rows-unique row-code collapse that made + both coordinates constant on high-dimensional token outputs (B-007). + """ + values = np.asarray(values, dtype=np.float64) + if values.ndim != 2 or values.shape[0] == 0: + return float("nan"), float("nan") + if not np.isfinite(values).all(): + return float("nan"), float("nan") + n = values.shape[0] + label_ids = np.asarray(labels).reshape(-1) + if label_ids.size != n or input_codes.size != n: + return float("nan"), float("nan") + features = values[:, :max_features] if values.shape[1] > max_features else values + ix_vals: list[float] = [] + iy_vals: list[float] = [] + for col in range(features.shape[1]): + quantized = _quantize_column(features[:, col], bins=bins, strategy=strategy) + ix = _mutual_information_bits(quantized, input_codes) + iy = _mutual_information_bits(quantized, label_ids) + if math.isfinite(ix) and math.isfinite(iy): + ix_vals.append(ix) + iy_vals.append(iy) + if not ix_vals: + return float("nan"), float("nan") + return float(np.mean(ix_vals)), float(np.mean(iy_vals)) + + def information_plane_for_model( model: nn.Module, x: torch.Tensor, @@ -115,12 +189,18 @@ def information_plane_for_model( include_output: bool = True, binning: str = "uniform_intervals", output_mapping: str = "tanh", + max_features: int = 8, ) -> list[list[float]]: - """Return one binned information point per activation (and optional output). + """Return one binned information point per captured layer (and optional output). - The forward hooks are temporary and model training mode is restored even if - the forward pass fails. Conversion to CPU happens only after all bounded - samples have been collected by the caller. + The first coordinate is a real ``I(X;T)`` estimate (mutual information + between the binned input encoding and the representation); the second is + ``I(T;Y)`` (B-007). Hooks prefer ``nn`` activation modules; models that + only use functional activations (RWKV/hyena/transformer) fall back to + capturing every module output. High-dimensional activations are measured + per feature column and aggregated so both coordinates vary with training + instead of degenerating to constants. Training mode is restored even if + the forward pass fails. """ captured: list[torch.Tensor] = [] handles: list[torch.utils.hooks.RemovableHandle] = [] @@ -130,18 +210,32 @@ def hook(_module: nn.Module, _args: tuple[object, ...], output: object) -> None: if isinstance(output, torch.Tensor) and output.ndim >= 1 and output.shape[0] == x.shape[0]: captured.append(output.detach()) - for module in model.modules(): - if isinstance(module, activation_types): - handles.append(module.register_forward_hook(hook)) was_training = model.training - try: - model.eval() - with torch.no_grad(): - output = _forward_reg(model, x) - finally: - for handle in handles: - handle.remove() - model.train(was_training) + + def _run() -> torch.Tensor | None: + try: + model.eval() + with torch.no_grad(): + return _forward_reg(model, x) + finally: + for handle in handles: + handle.remove() + model.train(was_training) + + # Preferred: nn activation modules define the semantic "layers". + for name, module in model.named_modules(): + if name and isinstance(module, activation_types): + handles.append(module.register_forward_hook(hook)) + output = _run() + # B-007: functional-activation models (RWKV/hyena) expose no nn activation + # modules ? capture every module output instead. + if not captured: + handles = [] + for name, module in model.named_modules(): + if name: + handles.append(module.register_forward_hook(hook)) + output = _run() + if include_output and isinstance(output, torch.Tensor) and output.ndim >= 1: if output_mapping == "probability": if output.ndim >= 2 and int(output.shape[-1]) > 1: @@ -159,11 +253,47 @@ def hook(_module: nn.Module, _args: tuple[object, ...], output: object) -> None: raise ValueError(f"unknown information-plane output mapping: {output_mapping}") captured.append(output.detach()) + x_np = x.detach().cpu().numpy() labels_np = labels.detach().cpu().numpy() points: list[list[float]] = [] for value in captured: - ix, iy = binned_information_pair( - value.detach().cpu().numpy(), labels_np, bins=bins, strategy=binning + arr = value.detach().cpu().numpy() + per_position = arr.ndim >= 3 + if per_position: + b, t = int(arr.shape[0]), int(arr.shape[1]) + arr = arr.reshape(b * t, -1) + elif arr.ndim == 1: + arr = arr[:, None] + elif arr.ndim != 2: + continue + n = int(arr.shape[0]) + if n == 0: + continue + if per_position: + if x_np.ndim >= 2 and x_np.shape[0] == b and x_np.shape[1] == t: + inputs_np = x_np.reshape(n, -1) + elif x_np.shape[0] == b: + inputs_np = np.repeat(x_np.reshape(b, -1), t, axis=0) + else: + inputs_np = x_np.reshape(n, -1) + if labels_np.size == n: + labels_n = labels_np.reshape(-1) + elif labels_np.size == b: + labels_n = np.repeat(labels_np.reshape(-1), t) + else: + continue + else: + inputs_np = x_np + labels_n = labels_np.reshape(-1) + if labels_n.size != n: + continue + ix, iy = _featurewise_information_pair( + arr, + labels_n, + _row_codes(inputs_np), + bins=bins, + strategy=binning, + max_features=max_features, ) if math.isfinite(ix) and math.isfinite(iy): points.append([float(ix), float(iy)]) diff --git a/comfy_research/engine/trainer/loss_terms.py b/comfy_research/engine/trainer/loss_terms.py index 889057c..c13c686 100644 --- a/comfy_research/engine/trainer/loss_terms.py +++ b/comfy_research/engine/trainer/loss_terms.py @@ -6,7 +6,7 @@ import torch.nn as nn from fastapi import HTTPException -from comfy_research.engine.losses.loss_builders import TrainerTask +from comfy_research.engine.losses.loss_builders import MultiSlotCrossEntropyLoss, TrainerTask from comfy_research.engine.node_builder_registry import registered_builder_node_types_for from comfy_research.engine.trainer.graph import _incoming_all from comfy_research.engine.trainer.model_helpers import _flatten_features_for_mse @@ -36,6 +36,25 @@ def _trainer_primary_loss_tensor( if trainer_task == "mse_regression": return criterion(_flatten_features_for_mse(pred), _flatten_features_for_mse(y)) * loss_scale if trainer_task in ("cross_entropy_dense", "token_classification", "vision_classification"): + if isinstance(criterion, MultiSlotCrossEntropyLoss): + if pred.dim() != 3 or y.dim() != 2: + raise HTTPException( + status_code=400, + detail=( + "MultiSlot cross-entropy expects pred [batch, K, vocab] and targets [batch, K]; " + f"got pred {tuple(int(d) for d in pred.shape)} and targets {tuple(int(d) for d in y.shape)}." + ), + ) + if tuple(pred.shape[:2]) != tuple(y.shape): + raise HTTPException( + status_code=400, + detail=( + "MultiSlot cross-entropy target [batch, K] must match pred [batch, K, vocab]; " + f"got pred {tuple(int(d) for d in pred.shape)} and targets " + f"{tuple(int(d) for d in y.shape)}." + ), + ) + return criterion(pred, y) * loss_scale if trainer_task == "token_classification" and pred.dim() == 3 and y.dim() == 2: if int(pred.shape[0]) != int(y.shape[0]) or int(pred.shape[1]) != int(y.shape[1]): raise HTTPException( diff --git a/comfy_research/engine/trainer/observable_viz.py b/comfy_research/engine/trainer/observable_viz.py index fa3a708..2b45436 100644 --- a/comfy_research/engine/trainer/observable_viz.py +++ b/comfy_research/engine/trainer/observable_viz.py @@ -120,6 +120,7 @@ def observable_viz_metric_updates( metric_histories: dict[str, list[float]], embedding_histories: dict[str, list[list[list[float]]]], attention_slice_histories: dict[str, list[dict[str, Any]]] | None = None, + warnings: dict[str, str] | None = None, ) -> list[dict[str, Any]]: """Build per-viz payloads for weight norm dynamics (same step_ticks as loss).""" has_viz_target = any( @@ -166,6 +167,7 @@ def observable_viz_metric_updates( metric_histories, embedding_histories, attention_slice_histories, + warnings, ) out: list[dict[str, Any]] = [] @@ -479,4 +481,9 @@ def observable_viz_metric_updates( if not hist: continue out.append(_obs_viz_stream_row({"node_id": viz.id, "value_history": hist}, paired)) + if warnings: + for row in out: + paired_id = row.get("paired_observable_id") + if paired_id in warnings: + row["warning"] = warnings[paired_id] return out diff --git a/comfy_research/engine/trainer/recorder.py b/comfy_research/engine/trainer/recorder.py index 3bb3e8c..a4f1fb5 100644 --- a/comfy_research/engine/trainer/recorder.py +++ b/comfy_research/engine/trainer/recorder.py @@ -101,10 +101,15 @@ def __init__(self, ctx: TrainerRunContext, *, get_user_observable_record: Any) - self._rep_tensors_cache = None self._attn_layers_cache = False self.observable_rng_generators: dict[str, torch.Generator] = {} + # Per-observable human-readable warnings (e.g. representation fallbacks). + self.observable_warnings: dict[str, str] = {} # Observable state initialized before the training loop starts. self.attention_arrays_init: dict[str, np.ndarray] | None = None if isinstance(self.model, _EMBEDDING_OBSERVABLE_MODEL_TYPES): - self.attention_arrays_init = self.model.observable_numpy_arrays() + self.attention_arrays_init = { + k: (v.copy() if hasattr(v, "copy") else v) + for k, v in self.model.observable_numpy_arrays().items() + } self.kan_regs = [on for on in self.observable_nodes if on.type == NodeKind.kan_reg] diff --git a/comfy_research/engine/trainer/training_loop.py b/comfy_research/engine/trainer/training_loop.py index c412325..a2e4953 100644 --- a/comfy_research/engine/trainer/training_loop.py +++ b/comfy_research/engine/trainer/training_loop.py @@ -560,6 +560,7 @@ def _lr_for_step(s: int) -> float | None: observable_metric_histories, observable_embedding_histories, recorder.observable_attention_slice_histories, + recorder.observable_warnings, ) yield { "type": "metrics", @@ -571,6 +572,7 @@ def _lr_for_step(s: int) -> float | None: "epoch_ticks": list(epoch_ticks), "observable_viz_updates": observable_updates, "observable_metric_histories": dict(observable_metric_histories), + "observable_warnings": dict(recorder.observable_warnings), } if cbs_epochs_complete: @@ -592,6 +594,7 @@ def _lr_for_step(s: int) -> float | None: observable_metric_histories, observable_embedding_histories, recorder.observable_attention_slice_histories, + recorder.observable_warnings, ) yield { "type": "paused", @@ -608,6 +611,7 @@ def _lr_for_step(s: int) -> float | None: "observable_metric_histories": observable_metric_histories, "observable_embedding_histories": observable_embedding_histories, "observable_attention_slice_histories": recorder.observable_attention_slice_histories, + "observable_warnings": dict(recorder.observable_warnings), "train_loop_seconds": train_loop_compute_seconds, } return @@ -624,6 +628,7 @@ def _lr_for_step(s: int) -> float | None: observable_metric_histories, observable_embedding_histories, recorder.observable_attention_slice_histories, + recorder.observable_warnings, ) yield { "type": "complete", @@ -639,6 +644,7 @@ def _lr_for_step(s: int) -> float | None: "observable_metric_histories": observable_metric_histories, "observable_embedding_histories": observable_embedding_histories, "observable_attention_slice_histories": recorder.observable_attention_slice_histories, + "observable_warnings": dict(recorder.observable_warnings), "train_loop_seconds": train_loop_compute_seconds, } finally: diff --git a/comfy_research/generated/node_kind.py b/comfy_research/generated/node_kind.py index 69118fa..a61704f 100644 --- a/comfy_research/generated/node_kind.py +++ b/comfy_research/generated/node_kind.py @@ -98,6 +98,7 @@ class GeneratedNodeKind(str, Enum): modular_addition_dataset = "modular_addition_dataset" moe_mlp_model = "moe_mlp_model" moe_mlp_token_model = "moe_mlp_token_model" + moving_statistics = "moving_statistics" mpp_spatiotemporal_model = "mpp_spatiotemporal_model" mse_loss = "mse_loss" multi_hop_fact_chain_dataset = "multi_hop_fact_chain_dataset" diff --git a/comfy_research/generated/node_manifest.json b/comfy_research/generated/node_manifest.json index 41206d5..6f38eb4 100644 --- a/comfy_research/generated/node_manifest.json +++ b/comfy_research/generated/node_manifest.json @@ -4529,6 +4529,53 @@ } ] }, + { + "type": "moving_statistics", + "label": "Moving statistics", + "category": "analysis", + "hint": "Rolling-window mean or standard deviation over a rank-1 (1D) tensor with configurable window, stride, and min periods; pass the result downstream as tensor.", + "resizable": true, + "hasDefaults": true, + "hasCodegen": false, + "fields": [ + { + "kind": "enum", + "key": "statistic", + "label": "Statistic", + "defaultValue": "mean" + }, + { + "kind": "int", + "key": "window", + "label": "Window", + "defaultValue": 10 + }, + { + "kind": "int", + "key": "stride", + "label": "Stride", + "defaultValue": 1 + }, + { + "kind": "int", + "key": "minPeriods", + "label": "Min Periods", + "defaultValue": 1 + }, + { + "kind": "boolean", + "key": "logScaleX", + "label": "Log Scale X", + "defaultValue": false + }, + { + "kind": "boolean", + "key": "logScaleY", + "label": "Log Scale Y", + "defaultValue": false + } + ] + }, { "type": "mpp_spatiotemporal_model", "label": "MPP-style spatiotemporal ViT", @@ -6549,12 +6596,6 @@ "label": "Inspect Format", "defaultValue": "id" }, - { - "kind": "float", - "key": "vocabCap", - "label": "Vocab Cap", - "defaultValue": 256 - }, { "kind": "enum", "key": "tokenizerMode", @@ -8747,12 +8788,6 @@ "label": "Inspect Format", "defaultValue": "id" }, - { - "kind": "float", - "key": "vocabCap", - "label": "Vocab Cap", - "defaultValue": 256 - }, { "kind": "enum", "key": "tokenizerMode", diff --git a/comfy_research/generated/node_params.py b/comfy_research/generated/node_params.py index 4fda3dd..bfed646 100644 --- a/comfy_research/generated/node_params.py +++ b/comfy_research/generated/node_params.py @@ -736,6 +736,14 @@ class MoeMlpTokenModelParams(NodeParamsBase): tieWeights: str | list[str] = "yes" seed: int | list[int] = 0 +class MovingStatisticsParams(NodeParamsBase): + statistic: str | list[str] = "mean" + window: int | list[int] = 10 + stride: int | list[int] = 1 + minPeriods: int | list[int] = 1 + logScaleX: bool | list[bool] = False + logScaleY: bool | list[bool] = False + class MppSpatiotemporalModelParams(NodeParamsBase): contextFrames: int | list[int] = 4 channels: int | list[int] = 1 @@ -1045,7 +1053,6 @@ class Phi1StyleDatasetParams(NodeParamsBase): dataSource: str | list[str] = "synthetic" cacheDir: str | list[str] = "" inspectFormat: str | list[str] = "id" - vocabCap: float | list[float] = 256 tokenizerMode: str | list[str] = "char" seqLen: int | list[int] = 600 stride: int | list[int] = 96 @@ -1404,7 +1411,6 @@ class TinystoriesDatasetParams(NodeParamsBase): dataSource: str | list[str] = "synthetic" cacheDir: str | list[str] = "" inspectFormat: str | list[str] = "id" - vocabCap: float | list[float] = 256 tokenizerMode: str | list[str] = "char" seqLen: int | list[int] = 512 stride: int | list[int] = 64 @@ -1638,6 +1644,7 @@ class VitModelParams(NodeParamsBase): "modular_addition_dataset": ModularAdditionDatasetParams, "moe_mlp_model": MoeMlpModelParams, "moe_mlp_token_model": MoeMlpTokenModelParams, + "moving_statistics": MovingStatisticsParams, "mpp_spatiotemporal_model": MppSpatiotemporalModelParams, "mse_loss": MseLossParams, "multi_hop_fact_chain_dataset": MultiHopFactChainDatasetParams, diff --git a/comfy_research/nodes/definitions/datasets/toy_language_phi1_style.py b/comfy_research/nodes/definitions/datasets/toy_language_phi1_style.py index 718fef5..0583295 100644 --- a/comfy_research/nodes/definitions/datasets/toy_language_phi1_style.py +++ b/comfy_research/nodes/definitions/datasets/toy_language_phi1_style.py @@ -27,7 +27,6 @@ EnumField(key="dataSource", label="Data Source", default="synthetic", options=('synthetic', 'download')), EnumField(key="cacheDir", label="Cache Dir", default=""), EnumField(key="inspectFormat", label="Inspect Format", default="id"), - FloatField(key="vocabCap", label="Vocab Cap", default=256, sweepable=False), EnumField(key="tokenizerMode", label="Tokenizer Mode", default="char", options=('char', 'word')), IntField(key="seqLen", label="Seq Len", default=600), IntField(key="stride", label="Stride", default=96, min=1), diff --git a/comfy_research/nodes/definitions/datasets/toy_language_tinystories.py b/comfy_research/nodes/definitions/datasets/toy_language_tinystories.py index 79a4ad6..3883805 100644 --- a/comfy_research/nodes/definitions/datasets/toy_language_tinystories.py +++ b/comfy_research/nodes/definitions/datasets/toy_language_tinystories.py @@ -27,7 +27,6 @@ EnumField(key="dataSource", label="Data Source", default="synthetic", options=('synthetic', 'download')), EnumField(key="cacheDir", label="Cache Dir", default=""), EnumField(key="inspectFormat", label="Inspect Format", default="id"), - FloatField(key="vocabCap", label="Vocab Cap", default=256, sweepable=False), EnumField(key="tokenizerMode", label="Tokenizer Mode", default="char", options=('char', 'word')), IntField(key="seqLen", label="Seq Len", default=512), IntField(key="stride", label="Stride", default=64, min=1), diff --git a/comfy_research/nodes/definitions/frontend/metric_compare.py b/comfy_research/nodes/definitions/frontend/metric_compare.py index 0a4bb1a..ab099bb 100644 --- a/comfy_research/nodes/definitions/frontend/metric_compare.py +++ b/comfy_research/nodes/definitions/frontend/metric_compare.py @@ -4,6 +4,19 @@ from comfy_research.nodes.registry import frontend_node_def from comfy_research.nodes.schema import EnumField, FrontendNodeDef, FrontendSpec, InPort, PortAccept +# 1D-series-capable viz sources: curve series viz, observable mirrors +# (observable_viz / observable_accuracy), training viz, and the tensor viz +# display family (tensor_viz_0d/1d/2d/general). Sources that don't resolve to a +# comparable 1D series surface a warning on the Metric compare node instead of +# failing at the port level, so any 1D tensor visualization can be wired here. +_COMPARE_SOURCES = ( + PortAccept(handles=("compare",), source_type="curve_series_viz"), + PortAccept(handles=("compare",), source_type="observable_viz"), + PortAccept(handles=("compare",), source_type="observable_accuracy"), + PortAccept(handles=("compare",), source_type="training_visualization"), + PortAccept(handles=("compare",), source_family="observable_user_tensor_viz_display"), +) + DEF = frontend_node_def( FrontendNodeDef( @@ -19,16 +32,8 @@ ("logScaleY", False), ), ports=( - InPort(id="left", accepts=( - PortAccept(handles=("compare",), source_type="curve_series_viz"), - PortAccept(handles=("compare",), source_type="observable_viz"), - PortAccept(handles=("compare",), source_type="tensor_viz_1d"), - )), - InPort(id="right", accepts=( - PortAccept(handles=("compare",), source_type="curve_series_viz"), - PortAccept(handles=("compare",), source_type="observable_viz"), - PortAccept(handles=("compare",), source_type="tensor_viz_1d"), - )), + InPort(id="left", accepts=_COMPARE_SOURCES), + InPort(id="right", accepts=_COMPARE_SOURCES), ), frontend=FrontendSpec(component_key="MetricCompareNode"), ) diff --git a/comfy_research/nodes/definitions/frontend/moving_statistics.py b/comfy_research/nodes/definitions/frontend/moving_statistics.py new file mode 100644 index 0000000..32cf3cc --- /dev/null +++ b/comfy_research/nodes/definitions/frontend/moving_statistics.py @@ -0,0 +1,33 @@ +"""moving_statistics ? FrontendNodeDef-channel definition.""" +from __future__ import annotations + +from comfy_research.nodes.registry import frontend_node_def +from comfy_research.nodes.schema import BoolField, EnumField, IntField, FrontendNodeDef, FrontendSpec + +DEF = frontend_node_def( + FrontendNodeDef( + type="moving_statistics", + label="Moving statistics", + category="analysis", + hint="Rolling-window mean or standard deviation over a rank-1 (1D) tensor with configurable window, stride, and min periods; pass the result downstream as tensor.", + fields=( + EnumField(key="statistic", label="Statistic", default="mean"), + IntField(key="window", label="Window", default=10, min=1, max=1024), + IntField(key="stride", label="Stride", default=1, min=1, max=1024), + IntField(key="minPeriods", label="Min Periods", default=1, min=1, max=1024), + BoolField(key="logScaleX", label="Log Scale X", default=False), + BoolField(key="logScaleY", label="Log Scale Y", default=False), + ), + defaults=( + ("statistic", "mean"), + ("window", 10), + ("stride", 1), + ("minPeriods", 1), + ("logScaleX", False), + ("logScaleY", False), + ("outputTensor", None), + ("lastError", None), + ), + frontend=FrontendSpec(component_key="MovingStatisticsNode"), + ) +) diff --git a/comfy_research/nodes/definitions/observables/fourier_component.py b/comfy_research/nodes/definitions/observables/fourier_component.py index 1bb5686..9f80f2e 100644 --- a/comfy_research/nodes/definitions/observables/fourier_component.py +++ b/comfy_research/nodes/definitions/observables/fourier_component.py @@ -50,6 +50,12 @@ def record(rec: "ObservableRecorder", on: "Node") -> None: histories = rec.observable_metric_histories if rec.trainer_task != "mse_regression": + warnings = getattr(rec, "observable_warnings", None) + if warnings is not None: + warnings[on.id] = ( + "observable_fourier_component only supports mse_regression tasks; " + "recording NaN on token/other tasks." + ) histories[on.id].append(float("nan")) return data: dict[str, Any] = on.data or {} @@ -63,6 +69,9 @@ def record(rec: "ObservableRecorder", on: "Node") -> None: input_axis=_scalar_int(data.get("inputAxis"), 0), output_index=_scalar_int(data.get("outputIndex"), 0), ) - except Exception: + except Exception as exc: + warnings = getattr(rec, "observable_warnings", None) + if warnings is not None: + warnings[on.id] = f"fourier component measurement failed: {exc}" value = float("nan") histories[on.id].append(float(value)) diff --git a/comfy_research/nodes/definitions/observables/information_plane.py b/comfy_research/nodes/definitions/observables/information_plane.py index ee830a2..bcd10dd 100644 --- a/comfy_research/nodes/definitions/observables/information_plane.py +++ b/comfy_research/nodes/definitions/observables/information_plane.py @@ -89,12 +89,25 @@ def record(rec: "ObservableRecorder", on: "Node") -> None: if count < int(x_eval.shape[0]): indices = torch.linspace(0, int(x_eval.shape[0]) - 1, count, device=x_eval.device).long() x_eval, y_eval = x_eval.index_select(0, indices), y_eval.index_select(0, indices) - history.append(information_plane_for_model( + points = information_plane_for_model( rec.model, x_eval, y_eval, bins=max(2, _scalar_int(data.get("bins"), 30)), include_output=_scalar_bool(data.get("includeOutput"), True), binning=str(data.get("binning") or "uniform_intervals"), output_mapping=str(data.get("outputMapping") or "tanh"), - )) - except Exception: + ) + if not points: + warnings = getattr(rec, "observable_warnings", None) + if warnings is not None: + warnings[on.id] = ( + "no activation tensors captured for information plane on " + f"{type(rec.model).__name__}; appending empty frame." + ) + history.append(points) + except Exception as exc: + warnings = getattr(rec, "observable_warnings", None) + if warnings is not None: + warnings[on.id] = ( + f"information plane measurement failed: {type(exc).__name__}: {exc}; appending empty frame." + ) history.append([]) diff --git a/comfy_research/nodes/definitions/observables/representation_alpha_req.py b/comfy_research/nodes/definitions/observables/representation_alpha_req.py index a2f8c09..7d0fb97 100644 --- a/comfy_research/nodes/definitions/observables/representation_alpha_req.py +++ b/comfy_research/nodes/definitions/observables/representation_alpha_req.py @@ -40,12 +40,30 @@ @recorder_for(REPRESENTATION_ALPHA_REQ) def record(rec: "ObservableRecorder", on: "Node") -> None: + from comfy_research.engine.analysis.representation_specs import resolve_representation_for_log from comfy_research.engine.trainer.scalar import _scalar_bool, _scalar_str data: dict[str, Any] = on.data or {} representation_id = _scalar_str(data.get("representationId"), "0::output").strip() + representation_tensors = rec._representation_tensors_for_log() + representation, resolved_id = resolve_representation_for_log( + rec.model, + representation_tensors, + representation_id, + ) + warnings = getattr(rec, "observable_warnings", None) + if resolved_id is not None and resolved_id != representation_id and warnings is not None: + warnings[on.id] = ( + f"representationId {representation_id!r} not found on model; " + f"using {resolved_id!r}." + ) + if representation is None and warnings is not None: + warnings[on.id] = ( + f"representationId {representation_id!r} could not be resolved on " + f"{type(rec.model).__name__} (no representation tensors collected); recording NaN." + ) _, eigenvalues = _rankme_and_eigenvalues( - rec._representation_tensors_for_log().get(representation_id), + representation, token_positions_as_samples=_scalar_bool(data.get("tokenPositionsAsSamples"), False), ) positive = [float(x) for x in eigenvalues if math.isfinite(x) and x > 1e-12] diff --git a/comfy_research/nodes/definitions/observables/representation_rankme.py b/comfy_research/nodes/definitions/observables/representation_rankme.py index 67f5710..569e7f9 100644 --- a/comfy_research/nodes/definitions/observables/representation_rankme.py +++ b/comfy_research/nodes/definitions/observables/representation_rankme.py @@ -94,12 +94,29 @@ def record(rec: "ObservableRecorder", on: "Node") -> None: """Log centered-covariance RankMe and the two leading eigenvalues.""" import torch.nn as nn + from comfy_research.engine.analysis.representation_specs import resolve_representation_for_log from comfy_research.engine.trainer.scalar import _scalar_bool, _scalar_str histories = rec.observable_metric_histories data: dict[str, Any] = on.data or {} representation_id = _scalar_str(data.get("representationId"), "0::output").strip() - representation = rec._representation_tensors_for_log().get(representation_id) + representation_tensors = rec._representation_tensors_for_log() + representation, resolved_id = resolve_representation_for_log( + rec.model, + representation_tensors, + representation_id, + ) + warnings = getattr(rec, "observable_warnings", None) + if resolved_id is not None and resolved_id != representation_id and warnings is not None: + warnings[on.id] = ( + f"representationId {representation_id!r} not found on model; " + f"using {resolved_id!r}." + ) + if representation is None and warnings is not None: + warnings[on.id] = ( + f"representationId {representation_id!r} could not be resolved on " + f"{type(rec.model).__name__} (no representation tensors collected); recording NaN." + ) rankme, eigenvalues = _rankme_and_eigenvalues( representation, token_positions_as_samples=_scalar_bool(data.get("tokenPositionsAsSamples"), False), diff --git a/comfy_research/nodes/definitions/observables/weight_product_sv.py b/comfy_research/nodes/definitions/observables/weight_product_sv.py index 933a7e4..87b6a6d 100644 --- a/comfy_research/nodes/definitions/observables/weight_product_sv.py +++ b/comfy_research/nodes/definitions/observables/weight_product_sv.py @@ -52,15 +52,23 @@ def record(rec: "ObservableRecorder", on: "Node") -> None: for name, p in model.named_parameters() if name.endswith(".weight") ] + warnings = getattr(rec, "observable_warnings", None) if weight_matrices: try: W_eff = weight_matrices[0] for W in weight_matrices[1:]: W_eff = W @ W_eff sv_sorted = torch.linalg.svdvals(W_eff).sort(descending=True).values - except Exception: + except Exception as exc: + if warnings is not None: + warnings[on.id] = ( + f"weight_product_sv: .weight tensors of {type(model).__name__} " + "cannot form a single W_eff chain (embedding/attention/head dims); recording NaN." + ) sv_sorted = torch.zeros(0) else: + if warnings is not None: + warnings[on.id] = "weight_product_sv: model has no .weight parameters; recording NaN." sv_sorted = torch.zeros(0) primary_sv = float(sv_sorted[0].item()) if sv_sorted.numel() > 0 else float("nan") observable_metric_histories[on.id].append(primary_sv) diff --git a/comfy_research/tests/test_frontend_channel_endstate.py b/comfy_research/tests/test_frontend_channel_endstate.py index 2f01ae5..a609f21 100644 --- a/comfy_research/tests/test_frontend_channel_endstate.py +++ b/comfy_research/tests/test_frontend_channel_endstate.py @@ -17,7 +17,7 @@ def test_frontend_channel_holds_all_registered_defs() -> None: from comfy_research.tests.test_frontend_channel_ledger import REGISTERED_FRONTEND_TYPES assert registry.frontend_def_types() == frozenset(REGISTERED_FRONTEND_TYPES) - assert len(REGISTERED_FRONTEND_TYPES) == 51 + assert len(REGISTERED_FRONTEND_TYPES) == 52 def test_eight_registry_symmetry_holds() -> None: diff --git a/comfy_research/tests/test_frontend_channel_ledger.py b/comfy_research/tests/test_frontend_channel_ledger.py index 79e0a8e..5b6853e 100644 --- a/comfy_research/tests/test_frontend_channel_ledger.py +++ b/comfy_research/tests/test_frontend_channel_ledger.py @@ -16,7 +16,7 @@ CORE_FRONTEND_NODES = ("comment", "graph_assist_failure_overlay", "hypothesis", "url_node") ANALYSIS_FRONTEND_NODES = ( "basic_calculator", "derivative_curve", "effective_rank", - "model_weight_tensors", "pca", "prediction", "series_endpoint_gap", "shape_checker", + "model_weight_tensors", "moving_statistics", "pca", "prediction", "series_endpoint_gap", "shape_checker", "smoothing_curve", "statistics", "statistics2", "svd", "tensor_reader", ) TENSOR_TABLE_FRONTEND_NODES = ( @@ -68,8 +68,8 @@ def test_frontend_registry_count() -> None: """The registry exposes the complete supported frontend node set.""" from comfy_research.nodes import registry - assert len(registry.frontend_def_types()) == 51 - assert len(REGISTERED_FRONTEND_TYPES) == 51 + assert len(registry.frontend_def_types()) == 52 + assert len(REGISTERED_FRONTEND_TYPES) == 52 def test_frontend_registry_types() -> None: diff --git a/comfy_research/tests/test_information_plane_observable.py b/comfy_research/tests/test_information_plane_observable.py index bb5d582..b2012c9 100644 --- a/comfy_research/tests/test_information_plane_observable.py +++ b/comfy_research/tests/test_information_plane_observable.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math from collections import defaultdict from types import SimpleNamespace from unittest.mock import patch @@ -104,3 +105,21 @@ def test_recorder_degrades_missing_hook_shape_and_calculation_failures_to_empty_ exploding = _rec(_Exploding(), x, y) record(exploding, _node()) assert exploding.observable_embedding_histories["plane"] == [[]] + + +def test_saxe_fixed_width_binning_tolerates_negative_activations() -> None: + model = _TinyNet() + x = -torch.ones(32, 2) + y = torch.zeros(32, dtype=torch.long) + points = information_plane_for_model(model, x, y, bins=8, binning="saxe_fixed_width_0_07") + assert len(points) == 2 + assert all(len(point) == 2 and all(math.isfinite(value) for value in point) for point in points) + + +def test_recorder_warns_when_measurement_raises() -> None: + x, y = _arrays() + rec = _rec(_Exploding(), x, y) + rec.observable_warnings = {} + record(rec, _node()) + assert rec.observable_embedding_histories["plane"] == [[]] + assert "information plane measurement failed" in rec.observable_warnings["plane"] diff --git a/comfy_research/tests/test_loss_terms.py b/comfy_research/tests/test_loss_terms.py new file mode 100644 index 0000000..e55d5de --- /dev/null +++ b/comfy_research/tests/test_loss_terms.py @@ -0,0 +1,83 @@ +"""Primary loss tensor regression: MultiSlot keeps [B, K, V] while token CE flattens.""" + +import pytest + +torch = pytest.importorskip("torch") + +from fastapi import HTTPException # noqa: E402 + +from comfy_research.engine.losses.loss_builders import MultiSlotCrossEntropyLoss # noqa: E402 +from comfy_research.engine.trainer.loss_terms import _trainer_primary_loss_tensor # noqa: E402 + + +def test_multi_slot_primary_loss_accepts_bkv_and_backprops() -> None: + torch.manual_seed(0) + pred = torch.randn(4, 3, 7, requires_grad=True) # [B, K, V] + targets = torch.randint(0, 7, (4, 3)) # [B, K] + criterion = MultiSlotCrossEntropyLoss(label_smoothing=0.0) + loss = _trainer_primary_loss_tensor( + pred, + targets, + trainer_task="token_classification", + criterion=criterion, + loss_scale=1.0, + ) + assert loss.ndim == 0 + assert torch.isfinite(loss) + assert loss.item() > 0.0 + loss.backward() + assert pred.grad is not None and torch.isfinite(pred.grad).all() + + +def test_token_ce_primary_loss_still_flattens_before_criterion() -> None: + torch.manual_seed(1) + pred = torch.randn(4, 3, 7) # [B, T, V] + targets = torch.randint(0, 7, (4, 3)) # [B, T] + criterion = torch.nn.CrossEntropyLoss(label_smoothing=0.0) + loss = _trainer_primary_loss_tensor( + pred, + targets, + trainer_task="token_classification", + criterion=criterion, + loss_scale=1.0, + ) + expected = criterion(pred.reshape(-1, 7), targets.reshape(-1).long()) + assert torch.allclose(loss, expected) + + +def test_multi_slot_primary_loss_rejects_flat_shapes() -> None: + pred = torch.randn(12, 7) # flattened [B*K, V] - the old buggy path + targets = torch.randint(0, 7, (12,)) + criterion = MultiSlotCrossEntropyLoss() + with pytest.raises(HTTPException) as exc_info: + _trainer_primary_loss_tensor( + pred, + targets, + trainer_task="token_classification", + criterion=criterion, + loss_scale=1.0, + ) + assert "MultiSlot" in str(exc_info.value.detail) + + +def test_multi_slot_primary_loss_rejects_transposed_target() -> None: + pred = torch.randn(2, 3, 5) # [B=2, K=3, V=5] + targets = torch.randint(0, 5, (3, 2)) # transposed [K, B]: same six slots, wrong pairing + criterion = MultiSlotCrossEntropyLoss() + with pytest.raises(HTTPException) as exc_info: + _trainer_primary_loss_tensor( + pred, + targets, + trainer_task="token_classification", + criterion=criterion, + loss_scale=1.0, + ) + assert "must match" in str(exc_info.value.detail) + + +def test_multi_slot_criterion_rejects_transposed_target() -> None: + pred = torch.randn(2, 3, 5) + targets = torch.randint(0, 5, (3, 2)) + criterion = MultiSlotCrossEntropyLoss() + with pytest.raises(ValueError, match="must match"): + criterion(pred, targets) diff --git a/comfy_research/tests/test_matrix_preconditioner_eigh.py b/comfy_research/tests/test_matrix_preconditioner_eigh.py new file mode 100644 index 0000000..21b9b20 --- /dev/null +++ b/comfy_research/tests/test_matrix_preconditioner_eigh.py @@ -0,0 +1,65 @@ +"""B-004: preconditioner eigensystem validation and CPU-double fallback.""" + +import pytest + +torch = pytest.importorskip("torch") + +from comfy_research.engine.optimizers.matrix_preconditioner_optimizers import ( # noqa: E402 + _eigh_psd_with_fallback, +) + + +def test_eigh_rejects_non_finite_input_matrix() -> None: + bad = torch.tensor([[1.0, float("nan")], [float("nan"), 1.0]]) + with pytest.raises(ValueError, match="non-finite"): + _eigh_psd_with_fallback(bad, eps=1e-8) + + +def test_eigh_rejects_non_finite_result(monkeypatch: pytest.MonkeyPatch) -> None: + real_eigh = torch.linalg.eigh + + def nan_eigh(mat: torch.Tensor, *args: object, **kwargs: object): + evals, evecs = real_eigh(mat, *args, **kwargs) + return evals * float("nan"), evecs + + monkeypatch.setattr(torch.linalg, "eigh", nan_eigh) + with pytest.raises(ValueError, match="non-finite"): + _eigh_psd_with_fallback(torch.tensor([[2.0, 0.0], [0.0, 1.0]]), eps=1e-8) + + +def test_eigh_retries_after_transient_non_finite_result(monkeypatch: pytest.MonkeyPatch) -> None: + real_eigh = torch.linalg.eigh + calls = 0 + + def transient_nan_eigh(mat: torch.Tensor, *args: object, **kwargs: object): + nonlocal calls + calls += 1 + evals, evecs = real_eigh(mat, *args, **kwargs) + if calls == 1: + return evals * float("nan"), evecs + return evals, evecs + + monkeypatch.setattr(torch.linalg, "eigh", transient_nan_eigh) + evals, evecs = _eigh_psd_with_fallback(torch.tensor([[2.0, 0.0], [0.0, 1.0]]), eps=1e-8) + assert calls >= 2 + assert torch.isfinite(evals).all() + assert torch.isfinite(evecs).all() + + +def test_eigh_cpu_double_fallback_succeeds(monkeypatch: pytest.MonkeyPatch) -> None: + real_eigh = torch.linalg.eigh + + def device_eigh_fails(mat: torch.Tensor, *args: object, **kwargs: object): + if mat.dtype != torch.float64: + raise RuntimeError("simulated device eigh failure") + return real_eigh(mat, *args, **kwargs) + + monkeypatch.setattr(torch.linalg, "eigh", device_eigh_fails) + mat = torch.tensor([[2.0, 0.5], [0.5, 1.0]]) + evals, evecs = _eigh_psd_with_fallback(mat, eps=1e-8) + assert torch.isfinite(evals).all() + assert torch.isfinite(evecs).all() + assert evals.dtype == mat.dtype + assert evecs.dtype == mat.dtype + reconstruct = (evecs * evals.unsqueeze(0)) @ evecs.T + assert torch.allclose(reconstruct, (mat + mat.T) / 2, atol=1e-4) diff --git a/comfy_research/tests/test_representation_resolve_smoke.py b/comfy_research/tests/test_representation_resolve_smoke.py new file mode 100644 index 0000000..2aa552c --- /dev/null +++ b/comfy_research/tests/test_representation_resolve_smoke.py @@ -0,0 +1,89 @@ +"""Runtime representation resolution fallback for rankme/alphaReQ observables. + +The lottery draw and manually wired observables can carry a stale default +representation id (``0::output``) that does not exist on most models. The +recorders resolve ids against the real model and fall back to the final hidden +representation; these tests pin that behaviour for the models the draw produces. +""" + +from __future__ import annotations + +import math + +import pytest +import torch +import torch.nn as nn + +from comfy_research.engine.analysis.representation_specs import ( + collect_representation_tensors, + resolve_representation_for_log, +) +from comfy_research.engine.models.model_builders import ( + _build_gated_mlp, + _build_mlp, + _build_moe_mlp, +) +from comfy_research.nodes.definitions.observables.representation_rankme import ( + _rankme_and_eigenvalues, +) + + +def _assert_resolves_to(model: nn.Module, expected_key: str, *, requested: str = "0::output") -> None: + reps = collect_representation_tensors(model, torch.randn(16, 11), 2) + assert requested not in reps, f"{requested} unexpectedly present: {sorted(reps)}" + tensor, resolved = resolve_representation_for_log(model, reps, requested) + assert resolved == expected_key + assert tensor is not None + rankme, eigenvalues = _rankme_and_eigenvalues(tensor, token_positions_as_samples=False) + assert math.isfinite(rankme) + assert eigenvalues and all(math.isfinite(v) for v in eigenvalues) + + +def test_mlp_default_is_valid_and_exact_match_wins() -> None: + # Plain MLP is a top-level Sequential, so the "0::output" default is a real + # module (the first Linear) and must resolve exactly, without a fallback. + model = _build_mlp(11, 1, 2, 48, "silu") + reps = collect_representation_tensors(model, torch.randn(16, 11), 2) + assert "0::output" in reps + tensor, resolved = resolve_representation_for_log(model, reps, "0::output") + assert resolved == "0::output" + assert tensor is not None + rankme, _ = _rankme_and_eigenvalues(tensor, token_positions_as_samples=False) + assert math.isfinite(rankme) + + +def test_gated_mlp_falls_back_to_out_input() -> None: + _assert_resolves_to(_build_gated_mlp(11, 1, 2, 48, "silu"), expected_key="out::input") + + +def test_moe_mlp_falls_back_to_last_expert_linear_input() -> None: + _assert_resolves_to(_build_moe_mlp(11, 1, 2, 32, 2, "silu"), expected_key="experts.1.4::input") + + +def test_exact_match_wins() -> None: + model = _build_gated_mlp(11, 1, 2, 48, "silu") + reps = collect_representation_tensors(model, torch.randn(16, 11), 2) + tensor, resolved = resolve_representation_for_log(model, reps, "gates.0::output") + assert resolved == "gates.0::output" + assert tensor is not None + + +def test_empty_tensors_resolve_to_none() -> None: + tensor, resolved = resolve_representation_for_log(nn.Sequential(), {}, "0::output") + assert tensor is None + assert resolved is None + + +def test_kan_falls_back_to_last_captured_output() -> None: + kan = pytest.importorskip("kan") + from comfy_research.engine.models.kan_model_build import build_kan_regression + + model = build_kan_regression(11, 1, 2, 5, 3, 3, 0, "silu", fast_training=True) + model.eval() + reps = collect_representation_tensors(model, torch.randn(16, 11), 2) + assert "0::output" not in reps + tensor, resolved = resolve_representation_for_log(model, reps, "0::output") + assert resolved == "act_fun.0.base_fun::output" + assert tensor is not None + rankme, _ = _rankme_and_eigenvalues(tensor, token_positions_as_samples=False) + assert math.isfinite(rankme) diff --git a/frontend/scripts/generate-connection-golden.ts b/frontend/scripts/generate-connection-golden.ts index 8fa4c1d..ce7caf5 100644 --- a/frontend/scripts/generate-connection-golden.ts +++ b/frontend/scripts/generate-connection-golden.ts @@ -14,6 +14,7 @@ import type { Connection, Edge, Node } from "@xyflow/react"; import { applyCanvasConnection, + CONNECTION_PROBE_HANDLE_PAIRS as HANDLE_PAIRS, isValidCanvasConnection, planAutoConnectCanvas, } from "../src/graph/connectionRules"; @@ -22,43 +23,6 @@ import { NODE_SPEC_REGISTRY } from "../src/graph/nodeRegistrySpec"; // bundle 后 import.meta.url 指向 node_modules/.cache——生成器恒从 frontend/ 运行,用 cwd。 export const GOLDEN_PATH = resolve(process.cwd(), "src/graph/__tests__/__snapshots__/connectionGolden.json"); -/** (sh, th) 组合集:cascade 分支中实际出现的搭配 + hostile 未知口(recon 词汇表)。 */ -const HANDLE_PAIRS: ReadonlyArray = [ - ["tensor", "tensor"], ["tensor", "tensor_in"], ["tensor", "tensor_1"], ["tensor", "tensor_2"], - ["tensor", "tensor_return"], ["tensor", "tensor_list"], ["tensor", "stream"], - ["tensor_out", "tensor"], ["tensor_out", "tensor_in"], ["tensor_out", "tensor_1"], - ["out_tensor", "tensor"], ["out_tensor", "stream"], ["out_tensor", "observable_in"], - ["observable", "observable_in"], ["tensor", "distribution"], ["sample_tensor", "distribution"], - ["out_tensor_list", "tensor_list"], ["tensor_boundary", "tensor_in"], ["tensor_boundary", "tensor"], - ["tensor_list", "tensor_list"], ["selected_tensor", "tensor"], - ["dataset", "dataset"], ["train_dataset", "dataset"], ["test_dataset", "dataset"], - ["dataset", "dataset_a"], ["dataset", "dataset_b"], ["dataset", "train_dataset"], - ["train_dataset", "dataset_a"], ["test_dataset", "dataset_a"], - ["train_dataset", "dataset_b"], ["test_dataset", "dataset_b"], - ["model", "model"], ["model", "tensor"], ["tensor", "model"], - ["loss", "loss"], ["optimizer", "optimizer"], ["observables", "observables"], - ["observable", "observables"], ["observable_results", "tensor"], ["loss_results", "tensor_list"], - ["checkpoint", "model_checkpoint"], ["model_checkpoint", "model"], - ["initialization", "initialization"], ["init", "initialization"], - ["lr_schedule", "lr_schedule"], ["mup_lr_schedule", "mup_lr_schedule"], - ["lr_schedule", "optimizer_lr_schedule"], ["mup_lr_schedule", "optimizer_lr_schedule"], - ["mup_lr_schedule", "lr_schedule"], - ["sample_tensor", "train_input"], ["sample_tensor", "test_input"], - ["input_distribution", "distribution"], ["distribution", "input_distribution"], - ["transformed_tensor", "tensor"], ["principal_components", "tensor"], ["u", "tensor"], ["s", "tensor"], ["v", "tensor"], - ["table", "table"], ["table", "tensor_list"], ["env", "env"], - ["coords", "coords"], ["coords", "pred_coords"], ["coords", "true_coords"], - ["comment", "comment"], ["", ""], ["tensor", ""], - ["__unknown_src__", "tensor"], ["tensor", "__unknown_tgt__"], ["__unknown_src__", "__unknown_tgt__"], - // 扩针:strip 源 handle 接 model 口 - //(fullModel/combined tensor-io 源)+ curve_annotator 入口。 - ["tensor_out", "model"], ["annotator", "from_viz"], - // 扩针(旧规则基线):paper-repro 新口——trainer batch_schedule、 - // 参数路径采样器双 checkpoint 口、curve-series 链(stream/series/curves)。 - ["batch_schedule", "batch_schedule"], ["model", "checkpoint_sb"], ["model", "checkpoint_lb"], - ["stream", "stream"], ["out_tensor_list", "stream"], ["series", "curves"], -]; - /** ioMode 敏感型(model/layer/combined 系)才展开 ioMode 变体。 */ const IO_MODE_VALUES = ["model", "input-output"] as const; const IO_SENSITIVE_PREFIX = /(_model|_layer|combined_model|reshape|flatten|softmax|causal_mask|tensor_splitter|einsum|elementwise_transform|tensor_slicing)$/; diff --git a/frontend/src/components/NodesLibraryPanel.tsx b/frontend/src/components/NodesLibraryPanel.tsx index f9cc3b1..718fed2 100644 --- a/frontend/src/components/NodesLibraryPanel.tsx +++ b/frontend/src/components/NodesLibraryPanel.tsx @@ -10,7 +10,7 @@ import { type ReactNode, type CSSProperties, } from "react"; -import { useReactFlow } from "@xyflow/react"; +import { useNodes, useReactFlow, type Node } from "@xyflow/react"; import { DeletingBusyOverlay } from "./DeletingBusyOverlay"; import type { ResearchGraphActions } from "../context/ResearchGraphContext"; import { useResearchGraph } from "../context/ResearchGraphContext"; @@ -32,6 +32,8 @@ import { useCombinedModelLibraryTemplates } from "../hooks/useCombinedModelLibra import { useNodeCategories, type CatalogNodeChild, type CatalogNodeCategory } from "../hooks/useNodeCategories"; import { GENERATED_NODE_SPECS } from "../generated/generatedNodeSpecs"; import { nodeRegistryDefaults } from "../graph/nodeRegistrySpec"; +import { canNodeTypeConnectToAny } from "../graph/canvasConnectivity"; +import { readNodeCanvasIoMode } from "../graph/nodeCanvasIoMode"; import { appendResearchNode } from "../graph/nodeInstanceTitle"; import { beginLibraryNodeDrag, @@ -570,6 +572,13 @@ export function NodesLibraryPanel() { const categories = useNodeCategories(); const combinedModelEntries = useCombinedModelLibraryTemplates(); const { screenToFlowPosition } = useReactFlow(); + /** Live canvas nodes; library drag drafts are excluded so the classification + * does not churn while a node is being dragged. */ + const canvasNodes = useNodes().filter( + (n) => !String(n.className ?? "").includes("cr-library-drag-draft"), + ); + const canvasNodesRef = useRef(canvasNodes); + canvasNodesRef.current = canvasNodes; useEffect( () => () => endLibraryNodeDrag(), @@ -822,6 +831,86 @@ export function NodesLibraryPanel() { const selectedLibraryCategory = libraryCategoryId === null ? null : categories.find((c) => c.id === libraryCategoryId) ?? null; + /** Canvas composition signature (ids + types + ioModes), ignoring position + * and drag churn so the classification only recomputes when it can change. */ + const canvasComposition = useMemo( + () => + canvasNodes + .map( + (n) => + `${n.id}:${String(n.type ?? "")}:${readNodeCanvasIoMode( + (n.data ?? {}) as Record | undefined, + )}`, + ) + .sort() + .join("|"), + [canvasNodes], + ); + + /** Connectable catalog types given the nodes currently on the canvas; null + * when the canvas is empty (every type can start a graph). */ + const connectableTypes = useMemo(() => { + const nodes = canvasNodesRef.current; + if (nodes.length === 0) return null; + const out = new Set(); + for (const nodeType of Object.keys(GENERATED_NODE_SPECS)) { + if (canNodeTypeConnectToAny(nodeType, nodes)) out.add(nodeType); + } + return out; + }, [canvasComposition]); + + const renderChildRow = (ch: NodeChild, c: NodeCategory) => ( + + ); + + /** Split a category's children into connectable / not-connectable groups + * relative to the current canvas, rendered as two labelled sub-lists. */ + const renderChildrenList = (c: NodeCategory, children: NodeChild[]) => { + if (connectableTypes == null) return children.map((ch) => renderChildRow(ch, c)); + const connectable: NodeChild[] = []; + const notConnectable: NodeChild[] = []; + for (const ch of children) { + (connectableTypes.has(ch.id) ? connectable : notConnectable).push(ch); + } + return ( + <> +
  • + Connectable · {connectable.length} +
    • + {connectable.map((ch) => renderChildRow(ch, c))} + {notConnectable.length > 0 && ( + <> +
  • + + Not connectable · {notConnectable.length} + +
    • + {notConnectable.map((ch) => renderChildRow(ch, c))} + + )} + + ); + }; + return ( <>