AbstractPhil commited on
Commit
64f2a91
·
verified ·
1 Parent(s): 24c682b

mirror alephlm.py from geolip 0.7.0 @0c946b3 (hub_ckpt; default 0 = v1 verbatim)

Browse files
Files changed (1) hide show
  1. alephlm.py +17 -1
alephlm.py CHANGED
@@ -15,6 +15,9 @@ from typing import NamedTuple, Optional
15
  import torch
16
  import torch.nn as nn
17
  import torch.nn.functional as F
 
 
 
18
 
19
 
20
  class LMOutput(NamedTuple):
@@ -36,6 +39,7 @@ class Block(nn.Module):
36
  super().__init__()
37
  d = cfg.d_model
38
  self.is_hub = layer_idx in cfg.hub_layers
 
39
  self.n1 = nn.LayerNorm(d)
40
  self.n2 = nn.LayerNorm(d)
41
  if self.is_hub:
@@ -48,8 +52,20 @@ class Block(nn.Module):
48
  cfg.gate_init)
49
 
50
  def forward(self, x, disable_bank=False, disable_hub=False):
 
 
 
 
 
51
  if not (disable_hub and self.is_hub):
52
- x = x + self.attn(self.n1(x))
 
 
 
 
 
 
 
53
  return x + self.bank(self.n2(x), disable_dispatch=disable_bank)
54
 
55
  def prefill(self, x):
 
15
  import torch
16
  import torch.nn as nn
17
  import torch.nn.functional as F
18
+ import torch.utils.checkpoint # noqa: F401 — torch<=2.4 does NOT auto-import
19
+ # this submodule; without it the hub_ckpt path AttributeErrors at step 1
20
+ # (the recorded A40 landmine, 2026-08-06)
21
 
22
 
23
  class LMOutput(NamedTuple):
 
39
  super().__init__()
40
  d = cfg.d_model
41
  self.is_hub = layer_idx in cfg.hub_layers
42
+ self.ckpt_mode = getattr(cfg, "hub_ckpt", 0)
43
  self.n1 = nn.LayerNorm(d)
44
  self.n2 = nn.LayerNorm(d)
45
  if self.is_hub:
 
52
  cfg.gate_init)
53
 
54
  def forward(self, x, disable_bank=False, disable_hub=False):
55
+ # hub_ckpt (2026-08-26): recompute-in-backward for the heavy
56
+ # branches — at v2 scale the retained scan tensors alone exceed a
57
+ # 95GB card (measured OOM). Training-path only; eval/decode and
58
+ # the census replay (which calls attn/bank directly) are untouched.
59
+ ckpt = self.ckpt_mode and self.training and torch.is_grad_enabled()
60
  if not (disable_hub and self.is_hub):
61
+ if ckpt:
62
+ x = x + torch.utils.checkpoint.checkpoint(
63
+ lambda t: self.attn(t), self.n1(x), use_reentrant=False)
64
+ else:
65
+ x = x + self.attn(self.n1(x))
66
+ if ckpt and self.ckpt_mode >= 2 and not disable_bank:
67
+ return x + torch.utils.checkpoint.checkpoint(
68
+ lambda t: self.bank(t), self.n2(x), use_reentrant=False)
69
  return x + self.bank(self.n2(x), disable_dispatch=disable_bank)
70
 
71
  def prefill(self, x):