fix(generate): only feed the new token to forward() after the first step The loop was passing the full growing input_ids to forward() every step while also maintaining a KV cache, so at step i it reprocessed all i tokens on top of a cache that already held them — O(n^2) memory in the number of generated tokens. Short generations that hit EOS early survived; anything longer OOM'd (25.8 GiB allocated for a 158 MB model). Pass only ids[:, -1:] after the first step; the RoPE positions already account for the cache length via `start`. Verified: model.generate() now runs to full length (120+ tokens) with no OOM.

#1
by Compactbot - opened
Files changed (1) hide show
  1. modeling_tinystories.py +2 -1
modeling_tinystories.py CHANGED
@@ -167,7 +167,8 @@ class TinyStoriesGPT(PreTrainedModel):
167
  ids = input_ids
168
  past = None
169
  for _ in range(max_new_tokens):
170
- out = self(ids, past_key_values=past, use_cache=True)
 
171
  past = out.past_key_values
172
  logits = out.logits[:, -1, :]
173
  if temperature and temperature != 1.0:
 
167
  ids = input_ids
168
  past = None
169
  for _ in range(max_new_tokens):
170
+ inp = ids if past is None else ids[:, -1:]
171
+ out = self(inp, past_key_values=past, use_cache=True)
172
  past = out.past_key_values
173
  logits = out.logits[:, -1, :]
174
  if temperature and temperature != 1.0: