Upload train.py

train.py

@@ -309,7 +309,7 @@ def _format_example(ex: dict, tok, text_column: str = "auto", include_reasoning:
     return str(ex)
 
 
-def build_dataset(seq_len: int, max_samples=None, split: str = "train",
+def build_dataset(seq_len: int, max_samples=None, max_tokens=None, split: str = "train",
                   dataset_name: str = "roneneldan/TinyStories",
                   dataset_config: str = None,
                   text_column: str = "auto",
@@ -322,6 +322,7 @@ def build_dataset(seq_len: int, max_samples=None, split: str = "train",
       - Messages/chat format (auto-detected, uses apply_chat_template)
       - Category filtering (comma-separated substrings)
       - Streaming for huge datasets
+      - Pre-allocated token buffer to avoid OOM on billion-token datasets
     """
     from datasets import load_dataset
     from chimera import ChimeraTokenizer
@@ -341,30 +342,77 @@ def build_dataset(seq_len: int, max_samples=None, split: str = "train",
         cat_filters = [c.strip() for c in category_filter.split(",") if c.strip()]
         print(f"[DATA] Filtering categories: {cat_filters}")
 
-    all_ids = []
-    target = max_samples * (seq_len + 1) if max_samples else float('inf')
+    # Determine token budget
+    if max_tokens is not None:
+        token_budget = max_tokens
+    elif max_samples is not None:
+        token_budget = max_samples * (seq_len + 1)
+    else:
+        token_budget = None
+
     processed = 0
     skipped = 0
 
-    for i, ex in enumerate(ds):
-        # Category filter
-        if cat_filters and not _matches_category_filter(ex, cat_filters):
-            skipped += 1
-            continue
-
-        text = _format_example(ex, tok, text_column, include_reasoning)
-        if not text or not text.strip():
-            skipped += 1
-            continue
-
-        all_ids.extend(tok.encode(text, add_special_tokens=False))
-        all_ids.append(tok.eos_token_id)
-        processed += 1
-
-        if len(all_ids) >= target:
-            break
-        if (processed + 1) % 10000 == 0:
-            print(f"  {processed:,} examples, {len(all_ids):,} tokens...")
+    if token_budget is not None and token_budget > 0:
+        # Pre-allocated flat buffer — avoids Python list overhead (~28 bytes/token)
+        buffer = torch.empty(token_budget, dtype=torch.long)
+        buf_idx = 0
+
+        for i, ex in enumerate(ds):
+            if cat_filters and not _matches_category_filter(ex, cat_filters):
+                skipped += 1
+                continue
+
+            text = _format_example(ex, tok, text_column, include_reasoning)
+            if not text or not text.strip():
+                skipped += 1
+                continue
+
+            ids = tok.encode(text, add_special_tokens=False)
+            ids.append(tok.eos_token_id)
+            n_ids = len(ids)
+
+            # Truncate if we would exceed the buffer
+            if buf_idx + n_ids > token_budget:
+                n_ids = token_budget - buf_idx
+                if n_ids <= 0:
+                    break
+                ids = ids[:n_ids]
+
+            if n_ids > 0:
+                buffer[buf_idx:buf_idx + n_ids] = torch.tensor(ids, dtype=torch.long)
+                buf_idx += n_ids
+            processed += 1
+
+            if buf_idx >= token_budget:
+                break
+            if (processed + 1) % 10000 == 0:
+                print(f"  {processed:,} examples, {buf_idx:,} tokens...")
+
+        all_ids = buffer[:buf_idx]
+    else:
+        # Fallback: old list approach for unbounded collection
+        all_ids = []
+        target = max_samples * (seq_len + 1) if max_samples else float('inf')
+        for i, ex in enumerate(ds):
+            if cat_filters and not _matches_category_filter(ex, cat_filters):
+                skipped += 1
+                continue
+
+            text = _format_example(ex, tok, text_column, include_reasoning)
+            if not text or not text.strip():
+                skipped += 1
+                continue
+
+            all_ids.extend(tok.encode(text, add_special_tokens=False))
+            all_ids.append(tok.eos_token_id)
+            processed += 1
+
+            if len(all_ids) >= target:
+                break
+            if (processed + 1) % 10000 == 0:
+                print(f"  {processed:,} examples, {len(all_ids):,} tokens...")
+        all_ids = torch.tensor(all_ids, dtype=torch.long)
 
     print(f"[DATA] Processed {processed:,} examples, skipped {skipped:,} (category/text mismatch)")
 
@@ -374,7 +422,6 @@ def build_dataset(seq_len: int, max_samples=None, split: str = "train",
             f"category_filter={category_filter}, text_column={text_column}"
         )
 
-    all_ids = torch.tensor(all_ids, dtype=torch.long)
     n = len(all_ids) // (seq_len + 1)
     if max_samples:
         n = min(n, max_samples)
@@ -487,6 +534,8 @@ def train(args):
     print(f"IPEX:         {HAS_IPEX}")
     print(f"Tokenizer:    splintr o200k_base ({config['vocab_size']} tokens)")
     print(f"Dataset:      {args.dataset_name} / {args.dataset_split}")
+    if args.dataset_config:
+        print(f"Dataset config: {args.dataset_config}")
     if args.category_filter:
         print(f"Category filter: {args.category_filter}")
     if args.include_reasoning:
@@ -530,6 +579,7 @@ def train(args):
     dataset, tok = build_dataset(
         args.seq_len,
         max_samples=args.max_samples,
+        max_tokens=args.max_tokens,
         split=args.dataset_split,
         dataset_name=args.dataset_name,
         dataset_config=args.dataset_config,
@@ -710,7 +760,10 @@ if __name__ == "__main__":
     p.add_argument("--lr", type=float, default=1e-3)
     p.add_argument("--warmup", type=int, default=200)
     p.add_argument("--max_steps", type=int, default=5000)
-    p.add_argument("--max_samples", type=int, default=None)
+    p.add_argument("--max_samples", type=int, default=None,
+                   help="Maximum number of chunks to generate")
+    p.add_argument("--max_tokens", type=int, default=None,
+                   help="Maximum total tokens to collect (pre-allocated buffer, prevents OOM on huge datasets)")
 
     # CPU Optimizations
     p.add_argument("--bf16", action="store_true", default=True,