diff options
Diffstat (limited to 'src/zbp_scaling/model.py')
| -rw-r--r-- | src/zbp_scaling/model.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/src/zbp_scaling/model.py b/src/zbp_scaling/model.py index 61ea4de..35c1548 100644 --- a/src/zbp_scaling/model.py +++ b/src/zbp_scaling/model.py @@ -61,7 +61,9 @@ class Block(nn.Module): super().__init__() self.norm_a = RMSNorm(d) self.qkv = nn.Linear(d, 3 * d, bias=False) - core = ZBPBlock(_Core(d, heads), cfg, batch_dims=1, name=f"L{li}.core") + # the score-space estimator materialises chunk*B*H*T^2 score tensors: give the core its own small chunk + core_cfg = cfg if cfg.mode == "bp" else cfg.replace(probe_chunk=min(cfg.probe_chunk or 4, 4)) + core = ZBPBlock(_Core(d, heads), core_cfg, batch_dims=1, name=f"L{li}.core") core.score_probe = {"d": d, "heads": heads, "dk": d, "dv": d, "window": None, "qk": True} self.core = core self.proj = nn.Linear(d, d, bias=False) |