From f301fd98d281491868b96fa8470a5d4b3d1303c2 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 14 Aug 2026 10:35:46 +0200 Subject: [PATCH] Make HistoryEncoder a pluggable registry, like Router/Objective (gitea #35) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stage2Autoregressive.init_history_cache and .history_step both isinstance-checked self.history_encoder against AttentionHistory to decide whether to use its real incremental-cache methods or a no-op fallback, so a third history type couldn't be added without editing Stage2Autoregressive itself. The two-value "markov"/"attention" enum was also independently hardcoded in three places (Stage2Autoregressive's own validation, config.py's validate_config, and AutoregressiveConfig.from_dict's default). Mirrors the Router (giant/model/routers.py) and Objective (giant/model/objectives.py, gitea #32) pattern: HistoryEncoder now declares working O(1) init_cache/step defaults (init_cache -> None, step -> one forward() call), so every registered history type satisfies the incremental interface without opting in; AttentionHistory overrides both with its real KV-cache versions since its forward() needs the full prefix. Added HISTORY_REGISTRY/register_history/build_history, registered "markov" and "attention", and deleted both isinstance checks in models.py. Per user decision during planning, config.py's validate_config now imports HISTORY_REGISTRY and checks membership dynamically instead of keeping its own hardcoded tuple, making the registry the single source of truth end to end (verified no import cycle: config.py had no prior dependency on giant.model, and giant.model.history has none on giant.config). No config-schema change and no checkpoint impact — this is a pure internal-interface refactor. Co-Authored-By: Claude Opus 5 --- giant/config.py | 7 ++-- giant/model/history.py | 76 +++++++++++++++++++++++++++++++++--------- giant/model/models.py | 26 ++++++--------- giant/model/network.py | 13 +++++++- tests/test_network.py | 45 +++++++++++++++++++++++++ 5 files changed, 133 insertions(+), 34 deletions(-) diff --git a/giant/config.py b/giant/config.py index 0bf1fc8..1310c5e 100644 --- a/giant/config.py +++ b/giant/config.py @@ -15,6 +15,7 @@ import numpy as np import torch from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing +from giant.model.history import HISTORY_REGISTRY class Conditioning(str, Enum): @@ -1341,8 +1342,10 @@ def validate_config(cfg: dict) -> None: "AutoregressiveConfig.order's docstring)" ) history = _get_path(cfg, "stage2_model.autoregressive.history") - if history not in ("markov", "attention"): - raise ValueError(f"stage2_model.autoregressive.history = {history!r} — must be 'markov' or 'attention'") + if history not in HISTORY_REGISTRY: + raise ValueError( + f"stage2_model.autoregressive.history = {history!r} — must be one of {sorted(HISTORY_REGISTRY)}" + ) teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing") if teacher_forcing not in ("always", "scheduled", "never"): raise ValueError( diff --git a/giant/model/history.py b/giant/model/history.py index 80fd22b..a3cc914 100644 --- a/giant/model/history.py +++ b/giant/model/history.py @@ -1,5 +1,9 @@ """History encoders — stage-2 autoregressive only. Self-contained, no -dependency on any other `giant.model` submodule (issues.md Issue 8).""" +dependency on any other `giant.model` submodule (issues.md Issue 8), except +for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors +`giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35).""" + +import inspect import torch import torch.nn as nn @@ -9,19 +13,57 @@ class HistoryEncoder(nn.Module): """Interface for stage-2 autoregressive per-token history summaries: `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over a full (teacher-forced) token sequence — used by training. `MarkovHistory` - and `AttentionHistory` are the two implementations. Inference - (`giant/sample.py`) generates one token at a - time and cannot afford `forward`'s per-step cost to be O(K) (attention - would then be O(K^2) over a rollout's k_max loop); encoders that need - incremental state for that path additionally implement `init_cache`/ - `step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since - its per-step cost is already O(1) (it only ever looks at the previous - token, not the full prefix).""" + and `AttentionHistory` are the two registered implementations (see + `HISTORY_REGISTRY`/`build_history`). Inference (`giant/sample.py`) + generates one token at a time and cannot afford `forward`'s per-step cost + to be O(K) (attention would then be O(K^2) over a rollout's k_max loop), + so this interface also declares `init_cache`/`step` for that incremental + path, with working O(1) defaults here (`init_cache` -> `None`, `step` -> + one `forward` call ignoring `cache`) — correct for any encoder whose + per-step cost is already O(1) (i.e. it only ever looks at the previous + token, not the full prefix), which is what `MarkovHistory` relies on. + `AttentionHistory` overrides both with real incremental-cache versions, + since its `forward` genuinely needs the full prefix.""" def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: raise NotImplementedError + def init_cache(self) -> object: + return None + def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]: + return self.forward(feat, has_prev), cache + + +HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {} + + +def register_history(name: str): + def decorator(cls: type[HistoryEncoder]) -> type[HistoryEncoder]: + HISTORY_REGISTRY[name] = cls + return cls + + return decorator + + +def build_history(name: str, in_dim: int, out_dim: int, **kwargs) -> HistoryEncoder: + """Factory: look up a `HistoryEncoder` subclass by name from the registry. + + Every registered history type is fed the same `stage2_model.autoregressive` + kwargs; kwargs not declared by that type's constructor are silently + dropped, so per-type hyperparameters (e.g. `AttentionHistory`'s + `n_heads`/`n_layers`) can coexist in one config without special-casing — + mirrors `giant.model.routers.build_router`. + """ + if name not in HISTORY_REGISTRY: + raise ValueError(f"unknown history type {name!r}; available: {sorted(HISTORY_REGISTRY)}") + cls = HISTORY_REGISTRY[name] + accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "in_dim", "out_dim"} + filtered = {k: v for k, v in kwargs.items() if k in accepted} + return cls(in_dim, out_dim, **filtered) + + +@register_history("markov") class MarkovHistory(HistoryEncoder): """Summarizes the previous secondary's own `(energy_fraction, direction, type_representation)` through one small MLP — the "markov" history: @@ -90,6 +132,7 @@ class _CausalAttnBlock(nn.Module): return x, kv +@register_history("attention") class AttentionHistory(HistoryEncoder): """Causal self-attention over the emitted-token prefix — the more expressive alternative to `MarkovHistory`'s fixed previous-token-only @@ -137,16 +180,19 @@ class AttentionHistory(HistoryEncoder): def step( self, - token_feat: torch.Tensor, + feat: torch.Tensor, has_prev: torch.Tensor, - cache: list[torch.Tensor | None], - ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: - """`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest - token's own features (what would be `feat[:, k]` in `forward`). + cache: object, + ) -> tuple[torch.Tensor, object]: + """`feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest token's + own features (what would be `feat[:, k]` in `forward`). `cache`: the + `list[Tensor | None]` from `init_cache`/a previous `step` call (typed + `object` here to match `HistoryEncoder.step`'s base signature). Advances every block's cache by this position and returns this position's output (`(B, 1, out_dim)`, the correct history summary for the NEXT slot) plus the updated cache.""" - x = self._embed(token_feat, has_prev) + assert isinstance(cache, list) + x = self._embed(feat, has_prev) new_cache: list[torch.Tensor | None] = [] for block, kv in zip(self.blocks, cache): x, kv_new = block.step(x, kv) diff --git a/giant/model/models.py b/giant/model/models.py index 3d5d3b2..3896704 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -6,7 +6,7 @@ import torch.nn as nn from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM from giant.model.encoders import ConditionEncoder -from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory +from giant.model.history import HistoryEncoder, build_history from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding from giant.model.objectives import build_objective from giant.model.routers import Router @@ -358,8 +358,6 @@ class Stage2Autoregressive(nn.Module): cond_enc: ConditionEncoder | None = None, ) -> None: super().__init__() - if history not in ("markov", "attention"): - raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'") self.history_kind = history self.generator_kind = generator self.noise_dim = noise_dim @@ -384,10 +382,8 @@ class Stage2Autoregressive(nn.Module): # this, a reasonable default rather than a design-doc-specified value. history_dim = cond_out_dim hist_in_dim = CONT_SLOT_DIM + self.type_dim - self.history_encoder: HistoryEncoder = ( - AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers) - if history == "attention" - else MarkovHistory(hist_in_dim, history_dim) + self.history_encoder: HistoryEncoder = build_history( + history, hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers ) token_fuse_in = cond_out_dim + context_dim + history_dim + 2 # +2: remaining_frac, slot_idx self.token_fuse = nn.Sequential( @@ -462,12 +458,12 @@ class Stage2Autoregressive(nn.Module): def init_history_cache(self): """Inference-only incremental-decoding state for `self.history_encoder` - (`giant/sample.py`'s AR loop): `None` under `history="markov"` (its - per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or - `AttentionHistory.init_cache()` under `history="attention"`.""" - if isinstance(self.history_encoder, AttentionHistory): - return self.history_encoder.init_cache() - return None + (`giant/sample.py`'s AR loop) — whatever `self.history_encoder.init_cache()` + returns for the configured `history` type: `None` under `history="markov"` + (its per-step cost is already O(1) — see `HistoryEncoder`'s docstring), + or `AttentionHistory.init_cache()`'s real per-block KV cache under + `history="attention"`.""" + return self.history_encoder.init_cache() def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]: """One inference slot's worth of history encoding: advances `cache` @@ -479,9 +475,7 @@ class Stage2Autoregressive(nn.Module): every model call made for this slot), `new_cache` is what to pass into the *next* slot's `history_step`. Must be called exactly once per slot — see `AttentionHistory.step`'s docstring.""" - if isinstance(self.history_encoder, AttentionHistory): - return self.history_encoder.step(token_feat, has_prev, cache) - return self.history_encoder(token_feat, has_prev), cache + return self.history_encoder.step(token_feat, has_prev, cache) def forward( self, diff --git a/giant/model/network.py b/giant/model/network.py index fbdcb44..cb21523 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -10,7 +10,15 @@ import X` call site keeps working unchanged. from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_state_dict from giant.model.builders import build_critics, build_models from giant.model.encoders import ConditionEncoder, cat_col_layout -from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory, _CausalAttnBlock +from giant.model.history import ( + HISTORY_REGISTRY, + AttentionHistory, + HistoryEncoder, + MarkovHistory, + _CausalAttnBlock, + build_history, + register_history, +) from giant.model.layers import ( BLOCK_REGISTRY, AdaLNResBlock, @@ -78,6 +86,7 @@ __all__ = [ "ExpertTrunk", "FilmResBlock", "FlowObjective", + "HISTORY_REGISTRY", "HistoryEncoder", "MarkovHistory", "OBJECTIVE_REGISTRY", @@ -106,6 +115,7 @@ __all__ = [ "build_composed_router", "build_critics", "build_expert_body", + "build_history", "build_models", "build_objective", "build_router", @@ -113,6 +123,7 @@ __all__ = [ "cat_col_layout", "migrate_legacy_state_dict", "register_block", + "register_history", "register_objective", "register_router", "register_trunk", diff --git a/tests/test_network.py b/tests/test_network.py index bf2fed0..f3fd5f2 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -5,14 +5,17 @@ import torch from giant import config as gconfig from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM from giant.model.network import ( + HISTORY_REGISTRY, AttentionHistory, ConditionEncoder, + HistoryEncoder, MarkovHistory, SinusoidalEmbedding, Stage1Model, Stage2Autoregressive, Stage2OneShot, build_critics, + build_history, build_models, cat_col_layout, stage2_trunk_sec_dim, @@ -407,6 +410,48 @@ def test_attention_history_step_matches_forward(): assert torch.allclose(stepped, expected, atol=1e-5) +# --- HISTORY_REGISTRY / build_history (gitea #35) ---------------------------- + + +def test_history_registry_has_exactly_the_two_known_histories(): + assert set(HISTORY_REGISTRY) == {"markov", "attention"} + + +def test_build_history_returns_correct_concrete_type(): + assert isinstance(build_history("markov", 4, 6), MarkovHistory) + assert isinstance(build_history("attention", 4, 8), AttentionHistory) + + +def test_build_history_unknown_name_raises(): + with pytest.raises(ValueError): + build_history("bogus", 4, 6) + + +def test_build_history_filters_kwargs_by_signature(): + """Attention-only kwargs (n_heads/n_layers) must be silently dropped when + building a MarkovHistory, matching build_router's documented behavior for + per-type hyperparameters coexisting in one config.""" + hist = build_history("markov", 4, 6, n_heads=2, n_layers=1) + assert isinstance(hist, MarkovHistory) + + +def test_history_encoder_base_default_init_cache_and_step(): + """A HistoryEncoder subclass implementing only forward() must still get + working O(1) init_cache/step defaults from the base class.""" + + class _StubHistory(HistoryEncoder): + def forward(self, feat, has_prev): + return feat * 2 + + hist = _StubHistory() + assert hist.init_cache() is None + feat = torch.randn(2, 1, 4) + has_prev = torch.ones(2, 1, dtype=torch.bool) + out, cache = hist.step(feat, has_prev, "unused-cache") + assert torch.equal(out, hist.forward(feat, has_prev)) + assert cache == "unused-cache" + + # --- Stage2Autoregressive (v0.3.0 step 5) -----------------------------------