mlnomad commited on
Commit
53d9696
·
verified ·
1 Parent(s): a931476

Fix: lazy non-persistent RoPE buffers — HF from_pretrained's meta-init was leaving persistent rope_cos/rope_sin uninitialised, producing NaN logits. Buffers are now computed on first forward.

Browse files
Files changed (1) hide show
  1. torch_gpt.py +22 -13
torch_gpt.py CHANGED
@@ -315,17 +315,27 @@ class GELU_GPT(nn.Module):
315
  }
316
  )
317
 
318
- # Precompute RoPE tables (non-trainable, stored as buffers).
319
- # Flax uses rotary_seq_len = sequence_len * 10; mirror that so long
320
- # contexts + parity tests both work.
321
- rotary_seq_len = config.sequence_len * 10
322
- cos, sin = precompute_rotary_embeddings(
323
- rotary_seq_len, config.head_dim, base=config.rope_base
324
- )
325
- # persistent=True so HF from_pretrained (which can meta-init
326
- # non-persistent buffers as NaN) restores them correctly.
327
- self.register_buffer("rope_cos", cos, persistent=True)
328
- self.register_buffer("rope_sin", sin, persistent=True)
 
 
 
 
 
 
 
 
 
 
329
 
330
  # ------------------------------------------------------------------
331
  # Forward
@@ -334,8 +344,7 @@ class GELU_GPT(nn.Module):
334
  B, T = idx.shape
335
  config = self.config
336
 
337
- cos = self.rope_cos[:, :T].to(dtype=self.wte.weight.dtype)
338
- sin = self.rope_sin[:, :T].to(dtype=self.wte.weight.dtype)
339
 
340
  x = self.wte(idx)
341
  x = rms_norm(x)
 
315
  }
316
  )
317
 
318
+ # Lazy RoPE: HF from_pretrained can leave persistent buffers
319
+ # uninitialised (meta-init NaN). We instead register an empty
320
+ # placeholder and compute the table on the first forward.
321
+ self._rope_max_len = config.sequence_len * 10
322
+ self._rope_head_dim = config.head_dim
323
+ self._rope_base = config.rope_base
324
+ self.register_buffer("rope_cos", torch.empty(0), persistent=False)
325
+ self.register_buffer("rope_sin", torch.empty(0), persistent=False)
326
+ self._rope_initialized = False
327
+
328
+ def _get_rope(self, T, dtype, device):
329
+ if (not self._rope_initialized
330
+ or self.rope_cos.numel() == 0
331
+ or self.rope_cos.shape[1] < T):
332
+ cos, sin = precompute_rotary_embeddings(
333
+ max(T, self._rope_max_len), self._rope_head_dim, base=self._rope_base
334
+ )
335
+ self.rope_cos = cos.to(device)
336
+ self.rope_sin = sin.to(device)
337
+ self._rope_initialized = True
338
+ return self.rope_cos[:, :T].to(dtype), self.rope_sin[:, :T].to(dtype)
339
 
340
  # ------------------------------------------------------------------
341
  # Forward
 
344
  B, T = idx.shape
345
  config = self.config
346
 
347
+ cos, sin = self._get_rope(T, self.wte.weight.dtype, self.wte.weight.device)
 
348
 
349
  x = self.wte(idx)
350
  x = rms_norm(x)