From 455ec118b0ae22ede5835d6095aae1282727ecfc Mon Sep 17 00:00:00 2001 From: quqingyu <12210819@mail.sustech.edu.cn> Date: Tue, 11 Aug 2026 03:30:00 +0800 Subject: [PATCH 1/6] port(main): apply PR #251 node updates onto public main --- .../engine/analysis/representation_specs.py | 83 +++ .../datasets/toy_language_external_runtime.py | 2 +- .../engine/datasets/toy_language_inspect.py | 2 +- .../matrix_preconditioner_optimizers.py | 26 +- comfy_research/engine/runs/trainer_run.py | 8 +- .../engine/trainer/dataset_materialize.py | 9 +- .../engine/trainer/information_plane.py | 161 +++++- comfy_research/engine/trainer/loss_terms.py | 12 +- .../engine/trainer/observable_viz.py | 7 + comfy_research/engine/trainer/recorder.py | 7 +- .../engine/trainer/training_loop.py | 6 + comfy_research/generated/node_kind.py | 1 + comfy_research/generated/node_manifest.json | 59 ++- comfy_research/generated/node_params.py | 11 +- .../datasets/toy_language_phi1_style.py | 1 - .../datasets/toy_language_tinystories.py | 1 - .../definitions/frontend/metric_compare.py | 25 +- .../definitions/frontend/moving_statistics.py | 33 ++ .../observables/fourier_component.py | 11 +- .../observables/information_plane.py | 12 +- .../observables/representation_alpha_req.py | 20 +- .../observables/representation_rankme.py | 19 +- .../observables/weight_product_sv.py | 10 +- .../tests/test_frontend_channel_endstate.py | 2 +- .../tests/test_frontend_channel_ledger.py | 6 +- .../test_representation_resolve_smoke.py | 89 ++++ .../scripts/generate-connection-golden.ts | 38 +- frontend/src/components/NodesLibraryPanel.tsx | 143 +++-- .../nodes/GenericObservableNode.tsx | 136 ++++- .../nodes/LayerSpectralNormVizNode.tsx | 2 +- .../components/nodes/MetricCompareNode.tsx | 127 ++++- .../components/nodes/MovingStatisticsNode.tsx | 495 ++++++++++++++++++ .../ObservableVizHessianEigenvaluesNode.tsx | 2 +- .../nodes/ObservableVizKanRegNode.tsx | 2 +- .../nodes/ObservableVizReluNonlinearNode.tsx | 2 +- .../nodes/ObservableVizUserNode.tsx | 6 + .../nodes/ObservableVizWeightL1Node.tsx | 2 +- .../components/nodes/TensorSelectorNode.tsx | 14 + .../src/components/nodes/TensorVizNode.tsx | 132 ++++- .../nodes/TrainingVisualizationNode.tsx | 2 +- .../nodes/movingStatisticsDefaults.ts | 52 ++ .../components/nodes/movingStatisticsMath.ts | 39 ++ .../nodes/observableVizUserDefaults.ts | 4 +- .../nodes/toyLanguageDatasetDefaults.ts | 3 - frontend/src/generated/generatedNodeSpecs.ts | 18 +- .../__snapshots__/codegenGolden.sha256.json | 4 +- .../__snapshots__/connectionGolden.json | 53 +- .../__tests__/__snapshots__/hintGolden.json | 1 + .../__tests__/canvasConnectivity.test.ts | 94 ++++ .../__tests__/generatedSpecs.seam.test.ts | 7 +- ...cObservableRepresentationDropdown.test.tsx | 103 ++++ .../graph/__tests__/metricComparePlot.test.ts | 56 +- .../metricCompareSourceResolution.test.ts | 98 ++++ .../__tests__/movingStatisticsMath.test.ts | 39 ++ .../__tests__/movingStatisticsWiring.test.ts | 38 ++ .../observableVizCompareHandles.test.tsx | 66 +++ .../tensorSelectorSeriesSources.test.ts | 183 +++++++ frontend/src/graph/canvasConnectivity.ts | 145 +++++ frontend/src/graph/connectionRules.ts | 103 +++- .../src/graph/curveSeriesVizResolution.ts | 17 - .../src/graph/fakeTensorShapePropagation.ts | 17 + frontend/src/graph/metricComparePlot.ts | 54 +- .../graph/metricCompareSourceResolution.ts | 117 +++++ frontend/src/graph/nodeRegistry.ts | 2 + frontend/src/graph/nodeSockets.ts | 224 ++++++++ .../graph/observableVizCompareResolution.ts | 33 +- frontend/src/graph/observableVizVariant.ts | 21 + frontend/src/graph/resolveUpstreamTensor.ts | 258 +++++++-- .../specCode/toyLanguageDatasetSpecCode.ts | 2 - .../toyLanguageExternalNotebookEmbed.ts | 2 +- ...LanguageRuntimeNotebookEmbeds.generated.ts | 2 +- .../src/graph/tensorViz1dCompareResolution.ts | 24 - .../src/graph/tensorVizCompareResolution.ts | 53 ++ frontend/src/graph/trainerVizPayload.ts | 7 +- frontend/src/graph/useHydratedResolved.ts | 13 +- frontend/src/index.css | 36 ++ 76 files changed, 3347 insertions(+), 367 deletions(-) create mode 100644 comfy_research/nodes/definitions/frontend/moving_statistics.py create mode 100644 comfy_research/tests/test_representation_resolve_smoke.py create mode 100644 frontend/src/components/nodes/MovingStatisticsNode.tsx create mode 100644 frontend/src/components/nodes/movingStatisticsDefaults.ts create mode 100644 frontend/src/components/nodes/movingStatisticsMath.ts create mode 100644 frontend/src/graph/__tests__/canvasConnectivity.test.ts create mode 100644 frontend/src/graph/__tests__/genericObservableRepresentationDropdown.test.tsx create mode 100644 frontend/src/graph/__tests__/metricCompareSourceResolution.test.ts create mode 100644 frontend/src/graph/__tests__/movingStatisticsMath.test.ts create mode 100644 frontend/src/graph/__tests__/movingStatisticsWiring.test.ts create mode 100644 frontend/src/graph/__tests__/observableVizCompareHandles.test.tsx create mode 100644 frontend/src/graph/__tests__/tensorSelectorSeriesSources.test.ts create mode 100644 frontend/src/graph/canvasConnectivity.ts create mode 100644 frontend/src/graph/metricCompareSourceResolution.ts create mode 100644 frontend/src/graph/nodeSockets.ts delete mode 100644 frontend/src/graph/tensorViz1dCompareResolution.ts create mode 100644 frontend/src/graph/tensorVizCompareResolution.ts 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/optimizers/matrix_preconditioner_optimizers.py b/comfy_research/engine/optimizers/matrix_preconditioner_optimizers.py index 7f2504a..2d50a8f 100644 --- a/comfy_research/engine/optimizers/matrix_preconditioner_optimizers.py +++ b/comfy_research/engine/optimizers/matrix_preconditioner_optimizers.py @@ -66,19 +66,19 @@ def _eigh_psd_with_fallback(mat: torch.Tensor, eps: float) -> tuple[torch.Tensor 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) + return evals.to(device=work.device, dtype=work.dtype), evecs.to( + device=work.device, dtype=work.dtype + ) + except RuntimeError as exc: + last_exc = 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..719613d 100644 --- a/comfy_research/engine/trainer/information_plane.py +++ b/comfy_research/engine/trainer/information_plane.py @@ -106,6 +106,77 @@ 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": + return np.floor(v / 0.07).astype(np.int32) + 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 +186,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 +207,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 +250,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..baa3aaa 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,16 @@ 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)}." + ), + ) + 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..381ffc1 --- /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..b6ebef4 100644 --- a/comfy_research/nodes/definitions/observables/information_plane.py +++ b/comfy_research/nodes/definitions/observables/information_plane.py @@ -89,12 +89,20 @@ 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"), - )) + ) + 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: 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_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 ( <>