#!/usr/bin/env python3
"""Long streamed ConvGPT-v2 2D-only trainer on PleIAs/SYNTH with chat-template reasoning eval."""
from __future__ import annotations
import argparse
import json
import os
import random
import time
from dataclasses import dataclass
import math
from datetime import datetime
from pathlib import Path
from typing import Any
import torch
import torch.nn.functional as F
from datasets import Dataset, load_dataset
from torch.utils.data import IterableDataset
from transformers import (
AutoTokenizer,
EarlyStoppingCallback,
Trainer,
TrainerCallback,
TrainingArguments,
set_seed,
)
from src.models.convgpt_v2 import ConvGPTV2Config, ConvGPTV2ForCausalLM
try:
import wandb
WANDB_AVAILABLE = True
except Exception:
wandb = None
WANDB_AVAILABLE = False
SEED = 42
DEFAULT_CACHE_DIR = Path("/home/mkurman/gitlab/ai/hf_cache")
DEFAULT_OUTPUT_ROOT = Path("/home/mkurman/gitlab/ai/model-trainer/results/convgpt_v2_long_2d")
DEFAULT_TOKENIZER = Path("/home/mkurman/gitlab/ai/model-trainer/neuroblast-merged")
DEFAULT_PLEIAS_FILES = ["synth_500.parquet", "synth_499.parquet"]
def normalize_pleias_files(files: list[str] | None) -> list[str] | None:
"""Return None when the full PleIAs/SYNTH train split should be streamed."""
if not files:
return None
normalized = [str(x).strip() for x in files if str(x).strip()]
if any(x.lower() in {"all", "full", "*"} for x in normalized):
return None
return normalized
def load_pleias_stream(files: list[str] | None, cache_dir: Path):
kwargs = {"streaming": True, "split": "train", "cache_dir": str(cache_dir)}
if files is not None:
kwargs["data_files"] = files
return load_dataset("PleIAs/SYNTH", **kwargs)
@dataclass
class CurriculumState:
lengths: list[int]
stage_steps: list[int]
bag_sizes: list[int]
current_length: int
current_bag_size: int = 1
total_tokens_seen: int = 0
def update_for_step(self, step: int) -> tuple[int, int]:
if not self.lengths:
return self.current_length, self.current_bag_size
idx = 0
for i, boundary in enumerate(self.stage_steps):
if step >= boundary:
idx = i
idx = min(idx, len(self.lengths) - 1)
self.current_length = self.lengths[idx]
self.current_bag_size = self.bag_sizes[idx] if self.bag_sizes else 1
return self.current_length, self.current_bag_size
@dataclass
class ChatDataCollator:
tokenizer: Any
max_length: int
curriculum_state: CurriculumState | None = None
def active_length(self) -> int:
if self.curriculum_state is not None:
return min(self.max_length, self.curriculum_state.current_length)
return self.max_length
def __call__(self, features: list[dict[str, Any]]):
pad_id = self.tokenizer.pad_token_id
input_ids = [torch.tensor(f["input_ids"], dtype=torch.long) for f in features]
attention_mask = [torch.tensor(f["attention_mask"], dtype=torch.long) for f in features]
labels = [torch.tensor(f["labels"], dtype=torch.long) for f in features]
bag_size = self.curriculum_state.current_bag_size if self.curriculum_state is not None else 1
input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=pad_id)
attention_mask = torch.nn.utils.rnn.pad_sequence(attention_mask, batch_first=True, padding_value=0)
labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)
active_length = self.active_length()
if input_ids.size(1) > active_length:
input_ids = input_ids[:, -active_length:]
attention_mask = attention_mask[:, -active_length:]
labels = labels[:, -active_length:]
elif input_ids.size(1) < active_length:
pad_width = active_length - input_ids.size(1)
# Fixed right-padding keeps the compiled training step shape-stable
# across batches. Labels use -100 only for these padding tokens.
input_ids = F.pad(input_ids, (0, pad_width), value=pad_id)
attention_mask = F.pad(attention_mask, (0, pad_width), value=0)
labels = F.pad(labels, (0, pad_width), value=-100)
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
"tst_bag_size": torch.full((input_ids.size(0),), bag_size, dtype=torch.long),
}
def extract_generation_prompt_text(item: Any) -> str:
"""Return the user prompt text from supported generation-prompt schemas."""
if isinstance(item, str):
return item.strip()
if isinstance(item, dict):
text = str(item.get("prompt") or item.get("content") or "").strip()
if text:
return text
messages = item.get("messages")
if isinstance(messages, list):
for message in messages:
if isinstance(message, dict) and message.get("role") == "user":
return str(message.get("content") or "").strip()
return ""
class CurriculumAndWandbCallback(TrainerCallback):
def __init__(
self,
tokenizer,
prompts: list[str],
max_new_tokens: int,
every_n_steps: int,
curriculum_state: CurriculumState,
effective_batch_tokens_multiplier: int,
):
self.tokenizer = tokenizer
self.prompts = prompts
self.max_new_tokens = max_new_tokens
self.every_n_steps = every_n_steps
self.curriculum_state = curriculum_state
self.effective_batch_tokens_multiplier = effective_batch_tokens_multiplier
self._last_step = 0
def _tokens_through_step(self, step: int) -> int:
"""Estimate curriculum tokens consumed through a resumed global step."""
total = 0
starts = list(self.curriculum_state.stage_steps)
lengths = list(self.curriculum_state.lengths)
for i, start in enumerate(starts):
end = starts[i + 1] if i + 1 < len(starts) else step
bounded_end = min(step, end)
if bounded_end > start:
total += (bounded_end - start) * self.effective_batch_tokens_multiplier * lengths[i]
if step < end:
break
return total
def on_train_begin(self, args, state, control, **kwargs):
if state.global_step > 0:
self.curriculum_state.update_for_step(state.global_step)
self.curriculum_state.total_tokens_seen = self._tokens_through_step(state.global_step)
self._last_step = state.global_step
print(
f"[resume] global_step={state.global_step} "
f"seq_len={self.curriculum_state.current_length} "
f"bag_size={self.curriculum_state.current_bag_size} "
f"tokens_seen={self.curriculum_state.total_tokens_seen}"
)
def _log_wandb(self, payload: dict, step: int):
if WANDB_AVAILABLE and wandb.run is not None:
wandb.log(payload, step=step)
def on_step_begin(self, args, state, control, **kwargs):
old_len = self.curriculum_state.current_length
old_bag = self.curriculum_state.current_bag_size
new_len, new_bag = self.curriculum_state.update_for_step(state.global_step)
if new_len != old_len or new_bag != old_bag:
print(f"[curriculum] step={state.global_step} seq_len={new_len} bag_size={new_bag}")
def on_step_end(self, args, state, control, **kwargs):
step_delta = max(0, state.global_step - self._last_step)
if step_delta:
self.curriculum_state.total_tokens_seen += (
step_delta * self.effective_batch_tokens_multiplier * self.curriculum_state.current_length
)
self._last_step = state.global_step
self._log_wandb(
{
"curriculum/seq_len": self.curriculum_state.current_length,
"curriculum/tst_bag_size": self.curriculum_state.current_bag_size,
"tokens/total_tokens_seen": self.curriculum_state.total_tokens_seen,
"tokens/total_tokens_seen_b": self.curriculum_state.total_tokens_seen / 1e9,
"tokens/tokens_per_optimizer_step": self.effective_batch_tokens_multiplier * self.curriculum_state.current_length,
},
state.global_step,
)
if state.global_step <= 0 or state.global_step % self.every_n_steps != 0:
return
model = kwargs.get("model")
if model is None:
return
device = next(model.parameters()).device
was_training = model.training
model.eval()
print(f"\n[gen-eval] --- step {state.global_step} ---")
rows = []
for idx, prompt_record in enumerate(self.prompts[:10], 1):
prompt = extract_generation_prompt_text(prompt_record)
if not prompt:
print(f"[gen-eval] SKIP_PROMPT {idx}: empty/unusable prompt record={prompt_record!r}")
continue
messages = [{"role": "user", "content": prompt}]
text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = self.tokenizer(text, return_tensors="pt", add_special_tokens=False).to(device)
with torch.no_grad():
out = model.generate(
**inputs,
max_new_tokens=self.max_new_tokens,
do_sample=False,
pad_token_id=self.tokenizer.pad_token_id,
eos_token_id=self.tokenizer.eos_token_id,
use_cache=True,
)
generated = self.tokenizer.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False)
full_text = self.tokenizer.decode(out[0], skip_special_tokens=False)
print(f"[gen-eval] PROMPT {idx}: {prompt}")
print(f"[gen-eval] RENDERED_PROMPT {idx}: {text[:500]}")
print(f"[gen-eval] OUTPUT {idx}: {generated[:2000]}")
rows.append([state.global_step, self.curriculum_state.current_length, self.curriculum_state.current_bag_size, self.curriculum_state.total_tokens_seen, idx, prompt, text, generated, full_text])
if WANDB_AVAILABLE and wandb.run is not None:
table = wandb.Table(columns=["step", "seq_len", "bag_size", "total_tokens_seen", "idx", "prompt", "rendered_prompt", "completion", "full_text"])
for row in rows:
table.add_data(*row)
wandb.log({"generations/samples": table}, step=state.global_step)
if was_training:
model.train()
class MaxRuntimeCallback(TrainerCallback):
def __init__(self, max_seconds: int):
self.max_seconds = int(max_seconds)
self.started_at: float | None = None
def on_train_begin(self, args, state, control, **kwargs):
self.started_at = time.monotonic()
def on_step_end(self, args, state, control, **kwargs):
if self.max_seconds <= 0 or self.started_at is None:
return control
elapsed = time.monotonic() - self.started_at
if elapsed >= self.max_seconds:
print(f"[runtime] max_train_seconds={self.max_seconds} reached at step={state.global_step}; saving and stopping")
control.should_save = True
control.should_training_stop = True
return control
class PleiasEncodedIterable(IterableDataset):
def __init__(self, tokenizer, files: list[str] | None, cache_dir: Path, max_length: int, seed: int, shuffle_buffer_size: int, per_pass_limit: int, curriculum_state: CurriculumState | None = None):
super().__init__()
self.tokenizer = tokenizer
self.files = files
self.cache_dir = cache_dir
self.max_length = max_length
self.seed = seed
self.shuffle_buffer_size = shuffle_buffer_size
self.per_pass_limit = per_pass_limit
self.curriculum_state = curriculum_state
def __iter__(self):
epoch = 0
while True:
ds = load_pleias_stream(self.files, self.cache_dir)
if self.shuffle_buffer_size > 0:
ds = ds.shuffle(seed=self.seed + epoch, buffer_size=self.shuffle_buffer_size)
produced = 0
for item in ds:
q = str(item.get("query") or "").strip()
r = str(item.get("synthetic_reasoning") or "").strip()
a = str(item.get("synthetic_answer") or "").strip()
if not q or not r or not a:
continue
ex = encode_example(
self.tokenizer,
[
{"role": "user", "content": q},
{"role": "assistant", "content": f"\n{r}\n\n\n{a}"},
],
min(self.max_length, self.curriculum_state.current_length) if self.curriculum_state is not None else self.max_length,
)
if ex is None:
continue
yield ex
produced += 1
if self.per_pass_limit > 0 and produced >= self.per_pass_limit:
break
epoch += 1
def build_model(args, tokenizer) -> ConvGPTV2ForCausalLM:
config = ConvGPTV2Config(
vocab_size=len(tokenizer),
hidden_size=args.hidden_size,
intermediate_size=args.hidden_size * 4,
num_hidden_layers=args.layers,
max_position_embeddings=args.max_position_embeddings or args.max_length,
grid_size=args.grid_size,
packing=args.packing,
conv2d_backend=args.conv2d_backend,
conv2d_chunk_size=args.conv2d_chunk_size,
two_d_every=args.two_d_every,
two_d_start_layer=args.two_d_start_layer,
position_embedding_type=getattr(args, "position_embedding_type", "learned"),
rope_theta=getattr(args, "rope_theta", 10000.0),
router_rope_fraction=getattr(args, "router_rope_fraction", 1.0),
use_row_col_embeddings=getattr(args, "use_row_col_embeddings", True),
conv1d_residual_gate_init=args.conv1d_gate_init,
conv2d_residual_gate_init=args.conv2d_gate_init,
branch_dropout=args.branch_dropout,
use_1d_branch=args.use_1d_branch,
use_2d_branch=args.use_2d_branch,
fusion="gated",
conv1d_kernel_size=3,
conv1d_dilations=[1, 2, 4, 8],
conv2d_kernel_size=3,
conv2d_dilations=[1, 2, 4, 8],
router_type=args.router_type,
retrieval_every=args.retrieval_every,
retrieval_num_slots=args.retrieval_num_slots,
retrieval_top_k=args.retrieval_top_k,
chunk_memory_size=args.chunk_memory_size,
chunk_memory_top_k=args.chunk_memory_top_k,
chunk_memory_token_top_k=args.chunk_memory_token_top_k,
chunk_memory_gate_init=args.chunk_memory_gate_init,
chunk_memory_include_current_chunk=args.chunk_memory_include_current_chunk,
dropout=args.dropout,
tie_word_embeddings=True,
pad_token_id=tokenizer.pad_token_id,
bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id,
)
model = ConvGPTV2ForCausalLM(config)
if hasattr(model, "gradient_checkpointing_enable") and args.gradient_checkpointing:
model.gradient_checkpointing_enable()
if hasattr(model.config, "use_cache"):
model.config.use_cache = False
return model
def build_stream_rows(files: list[str] | None, cache_dir: Path, limit: int):
ds = load_pleias_stream(files, cache_dir)
count = 0
for item in ds:
q = str(item.get("query") or "").strip()
r = str(item.get("synthetic_reasoning") or "").strip()
a = str(item.get("synthetic_answer") or "").strip()
if not q or not r or not a:
continue
yield {
"messages": [
{"role": "user", "content": q},
{"role": "assistant", "content": f"\n{r}\n\n\n{a}"},
]
}
count += 1
if limit > 0 and count >= limit:
break
def encode_example(tokenizer, messages: list[dict[str, str]], max_length: int):
full_ids: list[int] = []
full_labels: list[int] = []
for i in range(len(messages)):
prefix = messages[:i]
current = messages[: i + 1]
prefix_text = tokenizer.apply_chat_template(prefix, tokenize=False, add_generation_prompt=False) if prefix else ""
current_text = tokenizer.apply_chat_template(current, tokenize=False, add_generation_prompt=False)
prefix_ids = tokenizer(prefix_text, add_special_tokens=False)["input_ids"] if prefix_text else []
current_ids = tokenizer(current_text, add_special_tokens=False)["input_ids"]
delta_ids = current_ids[len(prefix_ids):]
full_ids.extend(delta_ids)
# Full-token training: every non-padding token contributes to the
# next-token loss. Padding is the only source of -100, added later by
# ChatDataCollator when examples in a batch have different lengths.
full_labels.extend(delta_ids)
if len(full_ids) > max_length:
# Preserve the chat-template / assistant boundary instead of blindly
# taking the tail. Very short curriculum stages otherwise see only a
# random answer suffix and lose the generation scaffold.
first_label = next((i for i, x in enumerate(full_labels) if x != -100), 0)
max_start = max(0, len(full_ids) - max_length)
start = min(first_label, max_start)
full_ids = full_ids[start : start + max_length]
full_labels = full_labels[start : start + max_length]
if not full_ids:
return None
if len(full_ids) < max_length:
pad_width = max_length - len(full_ids)
pad_id = tokenizer.pad_token_id
full_ids = full_ids + [pad_id] * pad_width
full_labels = full_labels + [-100] * pad_width
return {
"input_ids": full_ids,
"attention_mask": [0 if label == -100 else 1 for label in full_labels],
"labels": full_labels,
}
def build_eval_rows(tokenizer, files: list[str] | None, cache_dir: Path, max_length: int, eval_limit: int, seed: int, shuffle_buffer_size: int):
ds = load_pleias_stream(files, cache_dir)
if shuffle_buffer_size > 0:
ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer_size)
eval_encoded = []
for item in ds:
q = str(item.get("query") or "").strip()
r = str(item.get("synthetic_reasoning") or "").strip()
a = str(item.get("synthetic_answer") or "").strip()
if not q or not r or not a:
continue
ex = encode_example(
tokenizer,
[
{"role": "user", "content": q},
{"role": "assistant", "content": f"\n{r}\n\n\n{a}"},
],
max_length,
)
if ex is None:
continue
eval_encoded.append(ex)
if len(eval_encoded) >= eval_limit:
break
return eval_encoded
class TSTTrainer(Trainer):
# The TST loss already returns a per-batch mean and intentionally ignores
# Trainer's num_items_in_batch. Tell Trainer to apply the normal
# gradient-accumulation division; otherwise GA=32 logs/backprops ~32x loss.
model_accepts_loss_kwargs = False
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
if getattr(self.args, "torch_compile", False) and torch.cuda.is_available():
torch.compiler.cudagraph_mark_step_begin()
bag_size_tensor = inputs.pop("tst_bag_size", None)
bag_size = 1
if bag_size_tensor is not None:
bag_size = int(bag_size_tensor.flatten()[0].item())
labels = inputs.get("labels")
if bag_size <= 1 or labels is None:
return super().compute_loss(model, inputs, return_outputs=return_outputs, num_items_in_batch=num_items_in_batch)
# A next-bag objective needs at least one source position with future
# labels. Clamp overly aggressive bag sizes instead of producing fake
# zero-loss stages.
seq_len = int(labels.shape[1])
bag_size = min(bag_size, max(1, seq_len - 1))
if bag_size <= 1:
return super().compute_loss(model, inputs, return_outputs=return_outputs, num_items_in_batch=num_items_in_batch)
outputs = model(**inputs)
logits = outputs.logits.float()
labels = labels.to(logits.device)
# Multi-hot cross entropy for the next bag. For position t, predict
# labels t+1 ... t+bag_size, averaged uniformly over valid labels.
log_probs = torch.log_softmax(logits[:, :-bag_size, :], dim=-1)
if log_probs.shape[1] <= 0:
return super().compute_loss(model, inputs, return_outputs=return_outputs, num_items_in_batch=num_items_in_batch)
total = logits.new_tensor(0.0)
valid_count = logits.new_tensor(0.0)
for offset in range(1, bag_size + 1):
target = labels[:, offset : offset + log_probs.shape[1]]
valid = target.ne(-100)
safe_target = target.masked_fill(~valid, 0)
token_loss = -log_probs.gather(-1, safe_target.unsqueeze(-1)).squeeze(-1)
total = total + token_loss.masked_fill(~valid, 0).sum()
valid_count = valid_count + valid.sum()
if valid_count.item() == 0:
return super().compute_loss(model, inputs, return_outputs=return_outputs, num_items_in_batch=num_items_in_batch)
loss = total / valid_count
return (loss, outputs) if return_outputs else loss
def parse_args():
p = argparse.ArgumentParser(description="Long streamed ConvGPT-v2 2D-only trainer")
p.add_argument("--tokenizer", type=Path, default=DEFAULT_TOKENIZER)
p.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT)
p.add_argument("--run-name", type=str, default=None)
p.add_argument("--cache-dir", type=Path, default=DEFAULT_CACHE_DIR)
p.add_argument("--pleias-files", nargs="*", default=DEFAULT_PLEIAS_FILES)
p.add_argument("--train-limit", type=int, default=0, help="Max encoded examples to draw per streaming pass before reopening the stream; 0 = unbounded")
p.add_argument("--eval-limit", type=int, default=1000)
p.add_argument("--max-length", type=int, default=2048)
p.add_argument("--max-position-embeddings", type=int, default=None, help="Model positional capacity; defaults to --max-length. Use larger value when initializing from longer-context checkpoints.")
p.add_argument("--per-device-train-batch-size", type=int, default=1)
p.add_argument("--per-device-eval-batch-size", type=int, default=1)
p.add_argument("--gradient-accumulation-steps", type=int, default=32)
p.add_argument("--learning-rate", type=float, default=3e-4)
p.add_argument("--warmup-ratio", type=float, default=0.02)
p.add_argument("--max-steps", type=int, default=200000)
p.add_argument("--save-steps", type=int, default=1000)
p.add_argument("--eval-steps", type=int, default=1000)
p.add_argument("--logging-steps", type=int, default=20)
p.add_argument("--save-total-limit", type=int, default=5)
p.add_argument("--seed", type=int, default=SEED)
p.add_argument("--hidden-size", type=int, default=256)
p.add_argument("--layers", type=int, default=10)
p.add_argument("--grid-size", type=int, default=64)
p.add_argument("--packing", type=str, default="hilbert")
p.add_argument("--conv2d-backend", type=str, default="chunked_gather")
p.add_argument("--conv2d-chunk-size", type=int, default=1024)
p.add_argument("--use-1d-branch", action=argparse.BooleanOptionalAction, default=False)
p.add_argument("--use-2d-branch", action=argparse.BooleanOptionalAction, default=True)
p.add_argument("--two-d-every", type=int, default=1)
p.add_argument("--two-d-start-layer", type=int, default=0)
p.add_argument("--position-embedding-type", type=str, default="learned", choices=["learned", "nope", "rope_nope"])
p.add_argument("--rope-theta", type=float, default=10000.0)
p.add_argument("--router-rope-fraction", type=float, default=1.0)
p.add_argument("--use-row-col-embeddings", action=argparse.BooleanOptionalAction, default=True)
p.add_argument("--conv1d-gate-init", type=float, default=-12.0)
p.add_argument("--conv2d-gate-init", type=float, default=2.0)
p.add_argument("--router-type", type=str, default="none", choices=["none", "topk_memory", "chunk_memory", "chunk_token_memory"])
p.add_argument("--retrieval-every", type=int, default=0)
p.add_argument("--retrieval-top-k", type=int, default=4)
p.add_argument("--retrieval-num-slots", type=int, default=64)
p.add_argument("--chunk-memory-size", type=int, default=64)
p.add_argument("--chunk-memory-top-k", type=int, default=4)
p.add_argument("--chunk-memory-token-top-k", type=int, default=0)
p.add_argument("--chunk-memory-gate-init", type=float, default=-4.0)
p.add_argument("--chunk-memory-include-current-chunk", action="store_true", default=False)
p.add_argument("--branch-dropout", type=float, default=0.0)
p.add_argument("--dropout", type=float, default=0.0)
p.add_argument("--gradient-checkpointing", action=argparse.BooleanOptionalAction, default=True)
p.add_argument("--bf16", action=argparse.BooleanOptionalAction, default=torch.cuda.is_available())
p.add_argument("--torch-compile-model", action="store_true", default=False, help="Compile the ConvGPT model with torch.compile after moving it to CUDA. Experimental; use for fixed-shape throughput runs.")
p.add_argument("--torch-compile-mode", type=str, default="reduce-overhead", choices=["default", "reduce-overhead", "max-autotune"], help="torch.compile mode when --torch-compile-model is set.")
p.add_argument("--report-to", type=str, default="wandb")
p.add_argument("--wandb-mode", type=str, default=os.environ.get("WANDB_MODE", "online"))
p.add_argument("--dataloader-num-workers", type=int, default=0)
p.add_argument("--early-stopping-patience", type=int, default=0)
p.add_argument("--max-train-seconds", type=int, default=0, help="Wall-clock training budget after trainer start; 0 disables. Saves and stops when reached.")
p.add_argument("--gen-prompts-json", type=Path, required=True)
p.add_argument("--gen-max-new-tokens", type=int, default=256)
p.add_argument("--shuffle-buffer-size", type=int, default=10000)
p.add_argument("--curriculum-lengths", type=str, default="", help="Comma-separated sequence lengths, e.g. 8,16,32,...,2048. Empty disables curriculum.")
p.add_argument("--curriculum-stage-steps", type=str, default="", help="Comma-separated global-step starts for each curriculum length. Empty = evenly spaced stages over max_steps.")
p.add_argument("--tst-bag-sizes", type=str, default="", help="Comma-separated token-superposition bag sizes per curriculum stage. Use 1 for standard recovery.")
p.add_argument("--resume-from-checkpoint", type=Path, default=None, help="Optional HF Trainer checkpoint directory to resume from.")
p.add_argument("--init-from-checkpoint", type=Path, default=None, help="Optional model checkpoint to initialize weights from without resuming optimizer/trainer state.")
p.add_argument("--ignore-data-skip", action="store_true", default=False, help="When resuming streaming datasets, do not replay/skip prior batches before training.")
return p.parse_args()
def parse_int_list(spec: str) -> list[int]:
return [int(x.strip()) for x in spec.split(',') if x.strip()]
def build_curriculum_state(args) -> CurriculumState:
lengths = parse_int_list(args.curriculum_lengths)
explicit_steps = parse_int_list(args.curriculum_stage_steps) if args.curriculum_stage_steps.strip() else []
bag_sizes = parse_int_list(args.tst_bag_sizes) if args.tst_bag_sizes.strip() else []
if not lengths:
lengths = [args.max_length]
lengths = [min(args.max_length, max(1, x)) for x in lengths]
if lengths[-1] != args.max_length:
lengths.append(args.max_length)
if explicit_steps and len(explicit_steps) == len(lengths) - 1:
explicit_steps.append(args.max_steps)
if bag_sizes and len(bag_sizes) == len(lengths) - 1:
bag_sizes.append(1)
if explicit_steps:
stage_steps = explicit_steps
if len(stage_steps) != len(lengths):
raise ValueError("--curriculum-stage-steps must have the same length as the effective curriculum lengths")
else:
span = max(1, args.max_steps // len(lengths))
stage_steps = [i * span for i in range(len(lengths))]
if not bag_sizes:
bag_sizes = [1] * len(lengths)
if len(bag_sizes) != len(lengths):
raise ValueError("--tst-bag-sizes must have the same length as the effective curriculum lengths")
bag_sizes = [max(1, x) for x in bag_sizes]
return CurriculumState(lengths=lengths, stage_steps=stage_steps, bag_sizes=bag_sizes, current_length=lengths[0], current_bag_size=bag_sizes[0])
def estimate_curriculum_tokens(args, curriculum_state: CurriculumState) -> int:
starts = list(curriculum_state.stage_steps)
lengths = list(curriculum_state.lengths)
total = 0
mult = args.per_device_train_batch_size * args.gradient_accumulation_steps
for i, start in enumerate(starts):
end = starts[i + 1] if i + 1 < len(starts) else args.max_steps
if end <= start:
continue
total += (end - start) * mult * lengths[i]
return total
def load_generation_prompts(path: Path) -> list[str]:
raw = json.loads(path.read_text(encoding="utf-8"))
prompts: list[Any] = []
for item in raw:
text = extract_generation_prompt_text(item)
if text:
prompts.append(text)
if not prompts:
raise ValueError(f"No usable generation prompts found in {path}")
return prompts
def main():
args = parse_args()
set_seed(args.seed)
if args.run_name is None:
args.run_name = f"convgpt_v2_2d_pleias_long_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
args.pleias_files = normalize_pleias_files(args.pleias_files)
out_dir = args.output_root / args.run_name
out_dir.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ["WANDB_MODE"] = args.wandb_mode
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
tokenizer.padding_side = "right"
prompts = load_generation_prompts(args.gen_prompts_json)
curriculum_state = build_curriculum_state(args)
effective_token_multiplier = args.per_device_train_batch_size * args.gradient_accumulation_steps
estimated_curriculum_tokens = estimate_curriculum_tokens(args, curriculum_state)
print(f"[curriculum] lengths={curriculum_state.lengths}")
print(f"[curriculum] stage_steps={curriculum_state.stage_steps}")
print(f"[curriculum] tst_bag_sizes={curriculum_state.bag_sizes}")
print(f"[curriculum] estimated_tokens={estimated_curriculum_tokens}")
print(f"[curriculum] estimated_tokens_b={estimated_curriculum_tokens / 1e9:.3f}")
train_ds = PleiasEncodedIterable(
tokenizer=tokenizer,
files=args.pleias_files,
cache_dir=args.cache_dir,
max_length=args.max_length,
seed=args.seed,
shuffle_buffer_size=args.shuffle_buffer_size,
per_pass_limit=args.train_limit,
curriculum_state=curriculum_state,
)
eval_encoded = build_eval_rows(
tokenizer,
args.pleias_files,
args.cache_dir,
args.max_length,
args.eval_limit,
args.seed,
args.shuffle_buffer_size,
)
eval_ds = Dataset.from_list(eval_encoded)
print(f"[data] tokenizer={args.tokenizer}")
print(f"[data] pleias_files={args.pleias_files}")
print(f"[data] train_stream=streaming_iterable eval_examples={len(eval_ds)}")
print(f"[data] seq_len={args.max_length}")
print(f"[data] approx_target_tokens={args.max_steps * args.per_device_train_batch_size * args.gradient_accumulation_steps * args.max_length}")
print(f"[data] target_tokens_billions={(args.max_steps * args.per_device_train_batch_size * args.gradient_accumulation_steps * args.max_length) / 1e9:.3f}")
model = build_model(args, tokenizer)
if args.init_from_checkpoint is not None:
print(f"[model] init_from_checkpoint={args.init_from_checkpoint}")
model = ConvGPTV2ForCausalLM.from_pretrained(
str(args.init_from_checkpoint),
config=model.config,
ignore_mismatched_sizes=True,
dtype=torch.float32,
)
if hasattr(model, "gradient_checkpointing_enable") and args.gradient_checkpointing:
model.gradient_checkpointing_enable()
if hasattr(model.config, "use_cache"):
model.config.use_cache = False
if torch.cuda.is_available():
model = model.to("cuda")
if args.torch_compile_model:
if not torch.cuda.is_available():
raise ValueError("--torch-compile-model requires CUDA")
print(f"[model] torch_compile=True mode={args.torch_compile_mode}")
total_params = sum(p.numel() for p in model.parameters())
print(f"[model] params={total_params:,}")
print(
f"[model] 1d_branch={args.use_1d_branch} 2d_branch={args.use_2d_branch} "
f"packing={args.packing} grid_size={args.grid_size} "
f"conv1d_gate_init={args.conv1d_gate_init} conv2d_gate_init={args.conv2d_gate_init}"
)
collator = ChatDataCollator(tokenizer, args.max_length, curriculum_state=curriculum_state)
report_to = [] if args.report_to.strip().lower() in {"", "none", "off", "disabled"} else args.report_to
training_args = TrainingArguments(
output_dir=str(out_dir),
#overwrite_output_dir=True,
per_device_train_batch_size=args.per_device_train_batch_size,
per_device_eval_batch_size=args.per_device_eval_batch_size,
gradient_accumulation_steps=args.gradient_accumulation_steps,
learning_rate=args.learning_rate,
lr_scheduler_type="cosine",
warmup_ratio=args.warmup_ratio,
max_steps=args.max_steps,
bf16=args.bf16 and torch.cuda.is_available(),
fp16=False,
logging_nan_inf_filter=False,
torch_compile=args.torch_compile_model,
torch_compile_backend="inductor" if args.torch_compile_model else None,
torch_compile_mode=args.torch_compile_mode if args.torch_compile_model else None,
logging_steps=args.logging_steps,
save_steps=args.save_steps,
eval_steps=args.eval_steps,
eval_strategy="steps",
save_strategy="steps",
save_total_limit=args.save_total_limit,
report_to=report_to,
dataloader_num_workers=args.dataloader_num_workers,
remove_unused_columns=False,
gradient_checkpointing=args.gradient_checkpointing,
run_name=args.run_name,
seed=args.seed,
load_best_model_at_end=(args.early_stopping_patience > 0),
ignore_data_skip=args.ignore_data_skip,
metric_for_best_model="eval_loss",
greater_is_better=False,
)
callbacks = [CurriculumAndWandbCallback(tokenizer, prompts, args.gen_max_new_tokens, args.eval_steps, curriculum_state, effective_token_multiplier)]
if args.early_stopping_patience > 0:
callbacks.append(EarlyStoppingCallback(early_stopping_patience=args.early_stopping_patience))
if args.max_train_seconds > 0:
callbacks.append(MaxRuntimeCallback(args.max_train_seconds))
trainer = TSTTrainer(
model=model,
args=training_args,
train_dataset=train_ds,
eval_dataset=eval_ds,
data_collator=collator,
processing_class=tokenizer,
callbacks=callbacks,
)
trainer.model_accepts_loss_kwargs = False
manifest = {
"run_name": args.run_name,
"tokenizer": str(args.tokenizer),
"pleias_files": args.pleias_files,
"pleias_full_dataset": args.pleias_files is None,
"train_examples": "streaming_iterable",
"eval_examples": len(eval_ds),
"train_limit_per_pass": args.train_limit,
"shuffle_buffer_size": args.shuffle_buffer_size,
"max_length": args.max_length,
"max_position_embeddings": args.max_position_embeddings or args.max_length,
"hidden_size": args.hidden_size,
"layers": args.layers,
"grid_size": args.grid_size,
"packing": args.packing,
"conv2d_backend": args.conv2d_backend,
"two_d_every": args.two_d_every,
"two_d_start_layer": args.two_d_start_layer,
"position_embedding_type": getattr(args, "position_embedding_type", "learned"),
"rope_theta": getattr(args, "rope_theta", 10000.0),
"router_rope_fraction": getattr(args, "router_rope_fraction", 1.0),
"use_row_col_embeddings": getattr(args, "use_row_col_embeddings", True),
"conv1d_gate_init": args.conv1d_gate_init,
"conv2d_gate_init": args.conv2d_gate_init,
"router_type": args.router_type,
"retrieval_every": args.retrieval_every,
"retrieval_num_slots": args.retrieval_num_slots,
"retrieval_top_k": args.retrieval_top_k,
"chunk_memory_size": args.chunk_memory_size,
"chunk_memory_top_k": args.chunk_memory_top_k,
"chunk_memory_token_top_k": args.chunk_memory_token_top_k,
"chunk_memory_gate_init": args.chunk_memory_gate_init,
"chunk_memory_include_current_chunk": args.chunk_memory_include_current_chunk,
"torch_compile_model": args.torch_compile_model,
"torch_compile_mode": args.torch_compile_mode,
"use_1d_branch": args.use_1d_branch,
"use_2d_branch": args.use_2d_branch,
"router_type_effective": args.router_type,
"max_steps": args.max_steps,
"per_device_train_batch_size": args.per_device_train_batch_size,
"gradient_accumulation_steps": args.gradient_accumulation_steps,
"approx_target_tokens": args.max_steps * args.per_device_train_batch_size * args.gradient_accumulation_steps * args.max_length,
"approx_target_tokens_billions": (args.max_steps * args.per_device_train_batch_size * args.gradient_accumulation_steps * args.max_length) / 1e9,
"curriculum_lengths": curriculum_state.lengths,
"curriculum_stage_steps": curriculum_state.stage_steps,
"tst_bag_sizes": curriculum_state.bag_sizes,
"estimated_curriculum_tokens": estimated_curriculum_tokens,
"estimated_curriculum_tokens_billions": estimated_curriculum_tokens / 1e9,
"gen_eval_prompt_count": len(prompts),
"gen_max_new_tokens": args.gen_max_new_tokens,
"max_train_seconds": args.max_train_seconds,
"init_from_checkpoint": str(args.init_from_checkpoint) if args.init_from_checkpoint else None,
}
(out_dir / "run_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
result = trainer.train(resume_from_checkpoint=str(args.resume_from_checkpoint) if args.resume_from_checkpoint else None)
print(f"[train] metrics={result.metrics}")
final_dir = out_dir / "final"
trainer.save_model(str(final_dir))
tokenizer.save_pretrained(str(final_dir))
print(f"[train] saved={final_dir}")
if __name__ == "__main__":
main()