AbstractPhil commited on
Commit
5e71e4f
·
verified ·
1 Parent(s): fa82fde

load under transformers' remote-code loader: defer the imports it cannot resolve

Browse files

The loader resolves every relative import to a FLAT file beside the module,
whatever the nesting. Three `from .fusion import ...` sites in alephlm.py and
one `from .data.curriculum import ...` in presets.py therefore made it demand
fusion.py and data.curriculum.py in this repo and refuse to load the model
(OSError / FileNotFoundError on transformers 5).

Both imports are inert for this craft: the config carries no `fusion` key, so
that branch is never taken, and the curriculum is a training-side concern. They
are now resolved by name at the point of use, so nothing is imported unless the
code path is actually run. No weights, config or math changed.

Proved before push, in a virtualenv with transformers 5.13 and no geolip
installed: AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)
loads 112.5M params with no missing or unexpected keys and runs a forward
pass over a byte prompt.

Files changed (2) hide show
  1. alephlm.py +19 -3
  2. presets.py +12 -1
alephlm.py CHANGED
@@ -101,7 +101,17 @@ class AlephLM(nn.Module):
101
  # units and the front/back blocks stay at byte resolution.
102
  spec = getattr(cfg, "fusion", None)
103
  if spec:
104
- from .fusion import Fusion
 
 
 
 
 
 
 
 
 
 
105
  self.fusion = Fusion(spec, cfg.d_model)
106
  else:
107
  self.fusion = None
@@ -219,7 +229,10 @@ class AlephLM(nn.Module):
219
  across rows (the hub's prefix state has one shape per row)."""
220
  from .embedding import PAD_ROW
221
  from .governor import raw_block
222
- from .fusion import cat_caches, zero_cache_rows
 
 
 
223
  front, middle, back = self._ranges()
224
  for b in middle:
225
  assert raw_block(b).is_hub, "fused decode needs hub blocks in the middle"
@@ -265,7 +278,10 @@ class AlephLM(nn.Module):
265
  @torch.no_grad()
266
  def _decode_step_fused(self, next_id, cache, t):
267
  from .embedding import PAD_ROW
268
- from .fusion import snapshot, restore_rows
 
 
 
269
  front, middle, back = self._ranges()
270
  prev1, prev2, prev3 = cache["prev1"], cache["prev2"], cache["prev3"]
271
  st = self.fusion.starts_step(next_id, prev1, prev2, prev3, PAD_ROW) # (B,) closes the open unit
 
101
  # units and the front/back blocks stay at byte resolution.
102
  spec = getattr(cfg, "fusion", None)
103
  if spec:
104
+ # VENDORED: the live package reads
105
+ # from .fusion import Fusion
106
+ # transformers' remote-code loader resolves every relative
107
+ # import to a flat file beside this one, so written that way
108
+ # it demands a fusion.py here and refuses to load the model.
109
+ # This craft is fusion-free (its config carries no `fusion`
110
+ # key, so this branch is never taken); fusion belongs to the
111
+ # training side and itself imports the data subpackage, which
112
+ # a flat package cannot carry. Deferred and resolved by name.
113
+ from importlib import import_module
114
+ Fusion = import_module("geolip.alephllm.model.fusion").Fusion
115
  self.fusion = Fusion(spec, cfg.d_model)
116
  else:
117
  self.fusion = None
 
229
  across rows (the hub's prefix state has one shape per row)."""
230
  from .embedding import PAD_ROW
231
  from .governor import raw_block
232
+ # VENDORED: from .fusion import cat_caches, zero_cache_rows (see __init__)
233
+ from importlib import import_module
234
+ _f = import_module("geolip.alephllm.model.fusion")
235
+ cat_caches, zero_cache_rows = _f.cat_caches, _f.zero_cache_rows
236
  front, middle, back = self._ranges()
237
  for b in middle:
238
  assert raw_block(b).is_hub, "fused decode needs hub blocks in the middle"
 
278
  @torch.no_grad()
279
  def _decode_step_fused(self, next_id, cache, t):
280
  from .embedding import PAD_ROW
281
+ # VENDORED: from .fusion import snapshot, restore_rows (see __init__)
282
+ from importlib import import_module
283
+ _f = import_module("geolip.alephllm.model.fusion")
284
+ snapshot, restore_rows = _f.snapshot, _f.restore_rows
285
  front, middle, back = self._ranges()
286
  prev1, prev2, prev3 = cache["prev1"], cache["prev2"], cache["prev3"]
287
  st = self.fusion.starts_step(next_id, prev1, prev2, prev3, PAD_ROW) # (B,) closes the open unit
presets.py CHANGED
@@ -256,7 +256,18 @@ def make_v3_preset(n_layers: int = 24, d_model: int = 1024,
256
  epoch_cap / rebalance_to: the data-plane decisions (the trainer
257
  refuses a scaled stage without a rebalance rule); under 'hold' the
258
  stages stay at 1x and the held budget goes to fineweb_main."""
259
- from .data.curriculum import curriculum_phases, _BASE_STAGE_TOKENS
 
 
 
 
 
 
 
 
 
 
 
260
  if name is None:
261
  name = "mini-beatrix-3" if n_layers == 24 and d_model == 1024 \
262
  else f"mini-beatrix-3-d{d_model}-l{n_layers}"
 
256
  epoch_cap / rebalance_to: the data-plane decisions (the trainer
257
  refuses a scaled stage without a rebalance rule); under 'hold' the
258
  stages stay at 1x and the held budget goes to fineweb_main."""
259
+ # VENDORED: the live package reads
260
+ # from .data.curriculum import curriculum_phases, _BASE_STAGE_TOKENS
261
+ # Written that way here it is a relative import of a SUBPACKAGE, and
262
+ # transformers' remote-code loader resolves every relative import to a
263
+ # flat file beside this one — it would demand "data.curriculum.py" and
264
+ # refuse to load the model. This package ships the inference path only;
265
+ # the curriculum lives in the installable geolip-alephllm, so the
266
+ # import is deferred and resolved by name when that package is there.
267
+ from importlib import import_module
268
+ _curriculum = import_module("geolip.alephllm.data.curriculum")
269
+ curriculum_phases = _curriculum.curriculum_phases
270
+ _BASE_STAGE_TOKENS = _curriculum._BASE_STAGE_TOKENS
271
  if name is None:
272
  name = "mini-beatrix-3" if n_layers == 24 and d_model == 1024 \
273
  else f"mini-beatrix-3-d{d_model}-l{n_layers}"