Blame for code/graph/config.py

"""Configuration objects for reproducible GRM model construction."""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any, Literal


@dataclass(frozen=True)
class PATConfig:
    enabled: bool = True
    hidden_dim: int = 128
    num_layers: int = 2
    num_heads: int = 4
    max_path_length: int = 8
    num_token_types: int = 10
    dropout: float = 0.1
    relation_aware_attention: bool = True
    query_conditioning: bool = True
    type_vocab_size: int = 1024
    attribute_vocab_size: int = 2048

    def __post_init__(self) -> None:
        if self.hidden_dim <= 0 or self.num_layers <= 0 or self.num_heads <= 0:
            raise ValueError("PAT dimensions and layer counts must be positive")
        if (
            self.max_path_length <= 0
            or self.num_token_types <= 0
            or self.type_vocab_size <= 0
            or self.attribute_vocab_size <= 0
        ):
            raise ValueError("PAT path length and token vocabulary must be positive")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("PAT dropout must be in [0, 1)")


@dataclass(frozen=True)
class LossConfig:
    graph: float = 0.1
    link: float = 0.2
    next_hop: float = 1.0
    relation: float = 0.5
    path: float = 1.0
    pat: float = 1.0
    answer: float = 1.0
    verification: float = 0.5
    contradiction: float = 0.5
    calibration: float = 0.1
    count: float = 0.3
    #: Set-level selection objective.  Weighted above the per-path term because
    #: committing to the right number of paths is what exact match measures.
    selection_set: float = 1.5
    #: Listwise ranking objective over the candidate set.  ``count_guided_selection``
    #: commits the top-k candidates by score, so the ranking -- not the threshold --
    #: is what inference consumes; optimising only the thresholded objective left a
    #: train/inference mismatch of the same kind the on-policy experiment targeted.
    selection_rank: float = 1.0
    #: Weight on the *unsupported* class in the verification objective.  About 85% of
    #: curriculum examples are answerable, and with a per-example batch the unweighted
    #: objective is minimised well by always predicting "supported": the shipped model
    #: abstained on 25% of answerable examples and 12.5% of unanswerable ones, worse
    #: than chance in the direction that matters.  Set from the measured base rate.
    unsupported_weight: float = 5.0
    #: Penalty on the mean pairwise cosine of a graph's final node states.  Long
    #: runs collapsed every node to one direction (mean cosine 0.03 -> 0.98 over
    #: ~100K steps, artifacts/retrain-20260928-full-epoch/freeze.json), which
    #: leaves answer ranking and unseen-ontology transfer nothing to work with.
    #: 0 disables it and keeps older configs' objective unchanged.
    node_spread: float = 0.0


@dataclass(frozen=True)
class TrainingConfig:
    precision: Literal["fp32", "fp16", "bf16"] = "fp32"
    link_head: Literal["cosine", "mlp"] = "cosine"
    gradient_checkpointing: bool = False
    batch_size: int = 8
    gradient_accumulation: int = 1
    learning_rate: float = 3e-4
    warmup_ratio: float = 0.03
    weight_decay: float = 0.01
    max_steps: int = 1000
    seed: int = 7
    distributed_backend: str = "nccl"
    checkpoint_every: int = 0
    evaluation_every: int = 0
    use_fsdp: bool = False
    deterministic: bool = True
    grad_clip: float = 1.0
    #: Examples backpropagated together in one autograd invocation.  Larger is
    #: faster but holds more activations; 1 reproduces per-example backward.
    backward_chunk_size: int = 8
    #: When > 0, each training example gets a freshly renamed ontology with this
    #: probability.  This forces the model to learn relational semantics from
    #: type signatures and structure rather than memorizing relation names.
    #: 0=no renaming, 1=always rename.  Default 0.0 gives the original behavior.
    rename_probability: float = 0.0
    #: Run each backward chunk through ``GraphModel.forward_batch`` (one pass of
    #: every network over the chunk) instead of one forward per example.  The
    #: objective is the same; per-example forwards left the GPU ~85% idle.
    batched_forward: bool = True
    #: Learning-rate multiplier for the query encoder alone.  Its first
    #: self-attention layer was the source of the gradient blow-up that stopped
    #: the v3 full-epoch run at step 183,000 while the rest of the model was
    #: stable.  1.0 keeps the optimizer exactly as before.
    query_encoder_lr_scale: float = 1.0
    #: Stop training the query encoder.  At 0.1x learning rate its gradients
    #: still grew (median 88 -> 234 over 15K steps after step 183,500), always
    #: concentrated in its input projection and first attention layer.
    freeze_query_encoder: bool = False

    def __post_init__(self) -> None:
        # Literal types are not enforced at runtime; a typo such as "BF16" or
        # "gin" would otherwise be accepted and silently mapped to a default.
        if self.precision not in ("fp32", "fp16", "bf16"):
            raise ValueError(
                f"precision must be fp32, fp16 or bf16: {self.precision!r}"
            )
        if self.distributed_backend not in ("nccl", "gloo"):
            raise ValueError(
                f"distributed_backend must be nccl or gloo: {self.distributed_backend!r}"
            )
        if self.use_fsdp:
            # Refused here, not only when distributed: a single-process run would
            # otherwise train without FSDP and say nothing.
            raise ValueError("use_fsdp is not implemented; set it to false")
        if (
            self.batch_size <= 0
            or self.gradient_accumulation <= 0
            or self.max_steps <= 0
        ):
            raise ValueError(
                "batch_size, gradient_accumulation, and max_steps must be positive"
            )
        if (
            self.learning_rate <= 0
            or self.weight_decay < 0
            or not 0.0 <= self.warmup_ratio <= 1.0
        ):
            raise ValueError("invalid optimizer or warmup configuration")
        if self.checkpoint_every < 0 or self.evaluation_every < 0:
            raise ValueError("checkpoint/evaluation intervals cannot be negative")
        if self.grad_clip <= 0:
            raise ValueError("grad_clip must be positive")
        if self.backward_chunk_size <= 0:
            raise ValueError("backward_chunk_size must be positive")


@dataclass(frozen=True)
class GraphConfig:
    model_version: str = "graph-smoke-v1"
    target_parameters: int = 100_000_000
    hidden_dim: int = 128
    graph_input_dim: int = 8
    edge_input_dim: int = 8
    num_node_types: int = 64
    num_relations: int = 128
    num_relation_types: int = 128
    vocab_size: int = 16_384
    max_query_tokens: int = 128
    graph_layers: int = 3
    graph_encoder: Literal["sage", "gat"] = "sage"
    use_gat: bool = False
    #: Identity skip connection around each message-passing layer.  Without it
    #: eight layers of mean aggregation oversmooth: node states converged to a
    #: shared direction on long runs.  Off by default so existing checkpoints
    #: keep computing what they were trained to compute.
    graph_residual: bool = False
    #: Mark the query's start entities in the graph encoder's input: a learned
    #: start/other embedding plus the query state on start nodes, so node states
    #: are computed relative to where the question begins.  Synthetic entity
    #: names are random identifiers, so without a marker nodes that play the same
    #: structural role are indistinguishable and long runs collapsed them together.
    start_marking: bool = False
    #: Soft cap on attention scores in PAT and the reasoning core's graph
    #: cross-attention (``graph.utils.soft_cap``); 0 disables it.
    attention_logit_cap: float = 0.0
    #: Normalise attention queries and keys in PAT and the graph
    #: cross-attention (``graph.utils.qk_normalize``).
    attention_qk_norm: bool = False
    num_reasoning_layers: int = 2
    num_reasoning_heads: int = 4
    #: FFN expansion inside the reasoning core.  Exposed because it is the main
    #: lever for hitting the section 6 reasoning-transformer budget without
    #: changing depth or width.
    reasoning_ffn_multiplier: int = 4
    dropout: float = 0.1
    max_hops: int = 8
    beam_width: int = 8
    candidate_width: int = 32
    relation_top_k: int = 16
    answer_top_k: int = 8
    #: Largest count the aggregate head can emit.
    max_count: int = 32
    #: Cap on supervised positive/negative paths per example, bounding the PAT
    #: and selection batch built for one training step.
    max_supervised_paths: int = 8
    #: How many reachable paths to scan when topping up selection negatives from
    #: the distribution the beam actually produces.
    reachable_negative_pool: int = 256
    #: Negatives per positive in the selection batch.  This sets the *base rate*
    #: the threshold is fitted against, and it has to match inference: the beam
    #: surfaces roughly one gold path in six, so training on a balanced set
    #: produced a threshold that kept ten candidates when the gold set held two.
    selection_negative_ratio: int = 6
    #: Train path selection on the candidates the beam actually produces.
    #:
    #: Measured WORSE and therefore off by default.  At 200 steps on a 128-dim
    #: probe it lost on every metric (F1 0.254 vs 0.310, selection precision 0.149
    #: vs 0.195, renamed-ontology PAT 0.429 vs 0.741).  The cause looks like cold
    #: start: early in training the beam is near-random, so its candidate set holds
    #: almost no gold path and the objective sees a degenerate, low-diversity
    #: signal.  Kept because it is the right idea once a competent policy exists to
    #: bootstrap from -- worth revisiting as a late-training phase, not from step 0.
    on_policy_selection: bool = False
    #: Size the selected set from the count head's prediction instead of a
    #: threshold.  The threshold is uncalibrated against the inference candidate
    #: distribution, so it kept ~9 paths where the gold set held ~2.
    count_guided_selection: bool = True
    #: Semantic text features for entity, relation and query strings.  ``None`` keeps
    #: the original behaviour: every string becomes a one-hot at a hashed index, so two
    #: relations with the same meaning and different names get *orthogonal* features and
    #: an unseen relation string carries no meaning at all.  That is the measured cause
    #: of the model's transfer failure -- unseen-ontology exact match 0.150 with PAT at
    #: 0.557 against a 0.5 floor, and losing to a released 34M model on real typed KGs
    #: whose one advantage is reading strings through a pretrained encoder.
    #: Either "hashed-bag" (dependency-free, shares substrings) or
    #: "sentence-transformer" (frozen pretrained, shares meaning).
    #: Scalar evidence signals fed to the verification trunk.  The abstention decision
    #: was measured worse than chance -- abstaining on 25% of answerable examples and
    #: 12.5% of unanswerable ones -- because the pooled evidence the support head reads
    #: is built from gold positives during training and from the beam's own candidates
    #: at inference.  These signals (max and mean PAT query-compatibility, and the
    #: candidate count) are computed identically in both regimes, and PAT compatibility
    #: is already a well-trained discriminator (0.959).  Set to 3 to enable; it adds
    #: weights, so a checkpoint trained without it cannot be loaded with it.
    verification_signals: int = 0
    #: Soft bound on the verification and contradiction logits: cap * tanh(logit / cap).
    #: 0 disables it. Unbounded logits let one confidently wrong example produce a
    #: binary cross-entropy of 40 to 116, which destabilizes training late in the run.
    verification_logit_cap: float = 0.0
    #: Subtract a running mean of the verifier state before the heads read it. The
    #: state shares a large common direction across examples (mean pairwise cosine
    #: 0.69 at the collapse), which leaves the support head with no per-example signal.
    #: Uses running statistics at inference, so it has no train/test leak.
    verification_center: bool = False
    #: Report the trained confidence head's probability as the result confidence
    #: instead of sigmoid(best path score). The head is trained with a Brier loss
    #: against answer correctness; path scores are trained for ranking and are not
    #: probabilities, which is why path-score confidence calibrates poorly.
    confidence_from_head: bool = False
    #: Support probability below which the model abstains rather than answering.
    #:
    #: The evidence signals above make abstention work in-distribution, and they make it
    #: *distribution-sensitive*: they are built from PAT query-compatibility, which
    #: saturates low on relation vocabulary the model has never seen.  On real typed KGs
    #: the gate therefore fired on 100% of queries -- every real-KG metric went to
    #: 0.000, even when the gold relation path was supplied.  A caller working outside
    #: the training distribution should lower this, or pass ``abstain=False`` to
    #: ``reason`` and rely on the deterministic verifier alone.
    abstention_threshold: float = 0.5
    text_features: str | None = None
    text_feature_dim: int = 384
    text_feature_model: str | None = None
    #: Depth and width of the shared neural verification trunk (PLAN section 20).
    verifier_layers: int = 2
    verifier_width_multiplier: int = 4
    pat: PATConfig = field(default_factory=PATConfig)
    loss: LossConfig = field(default_factory=LossConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)

    def __post_init__(self) -> None:
        if self.graph_encoder not in ("sage", "gat"):
            raise ValueError(
                f"graph_encoder must be sage or gat: {self.graph_encoder!r}"
            )
        if (
            self.hidden_dim <= 0
            or self.graph_layers <= 0
            or self.graph_input_dim <= 0
            or self.edge_input_dim <= 0
        ):
            raise ValueError("model dimensions and graph layer count must be positive")
        if (
            self.num_node_types <= 0
            or self.num_relations <= 0
            or self.num_relation_types <= 0
        ):
            raise ValueError("node and relation vocabularies must be positive")
        if self.vocab_size <= 0 or self.max_query_tokens <= 0:
            raise ValueError("query vocabulary and token limit must be positive")
        if self.pat.hidden_dim != self.hidden_dim:
            raise ValueError(
                "PAT hidden_dim must match GRM hidden_dim for residual fusion"
            )
        if self.pat.num_heads <= 0 or self.hidden_dim % self.pat.num_heads:
            raise ValueError("hidden_dim must be divisible by PAT num_heads")
        if self.num_reasoning_heads <= 0 or self.hidden_dim % self.num_reasoning_heads:
            raise ValueError("hidden_dim must be divisible by reasoning heads")
        if self.max_hops <= 0 or self.beam_width <= 0:
            raise ValueError("max_hops and beam_width must be positive")
        if (
            self.candidate_width <= 0
            or self.relation_top_k <= 0
            or self.answer_top_k <= 0
        ):
            raise ValueError("reasoning widths must be positive")
        if self.reasoning_ffn_multiplier <= 0:
            raise ValueError("reasoning_ffn_multiplier must be positive")
        if self.max_count <= 0 or self.max_supervised_paths <= 0:
            raise ValueError("max_count and max_supervised_paths must be positive")
        if self.reachable_negative_pool <= 0:
            raise ValueError("reachable_negative_pool must be positive")
        if self.selection_negative_ratio <= 0:
            raise ValueError("selection_negative_ratio must be positive")
        if self.verifier_layers <= 0 or self.verifier_width_multiplier <= 0:
            raise ValueError("verifier depth and width multiplier must be positive")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("model dropout must be in [0, 1)")

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)