diff options
| author | yurenh <blackhao0426@gmail.com> | 2026-08-31 18:15:39 -0500 |
|---|---|---|
| committer | yurenh <blackhao0426@gmail.com> | 2026-08-31 18:15:39 -0500 |
| commit | 7db653a60d5125774d60da8d38ee3d49a787be91 (patch) | |
| tree | 685f199605ba3db2c0bdf32f8dcfae6a446cb76b | |
| parent | 6a544fabfc2af22e4d5823410dd2387b5af89ea9 (diff) | |
cap score-core probe chunk (memory ~ chunk*B*H*T^2)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe
| -rw-r--r-- | src/zbp_scaling/__pycache__/model.cpython-313.pyc | bin | 9121 -> 9337 bytes | |||
| -rw-r--r-- | src/zbp_scaling/model.py | 4 |
2 files changed, 3 insertions, 1 deletions
diff --git a/src/zbp_scaling/__pycache__/model.cpython-313.pyc b/src/zbp_scaling/__pycache__/model.cpython-313.pyc Binary files differindex 6b9985f..df0135f 100644 --- a/src/zbp_scaling/__pycache__/model.cpython-313.pyc +++ b/src/zbp_scaling/__pycache__/model.cpython-313.pyc 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) |