Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """ | |
| Train an Agron model FROM SCRATCH (random initialization). | |
| Every architecture, optimiser, scheduler and algorithm from ``agron_nets`` is | |
| available on the command line. | |
| Usage examples: | |
| python train_from_scratch.py --size small --steps 500 | |
| python train_from_scratch.py --size small --arch mamba --optimizer muon \\ | |
| --recipe orthogonal --steps 1000 | |
| python train_from_scratch.py --arch moe --optimizer lion --scheduler wsd \\ | |
| --algorithm gradient_clip --algorithm ema --steps 2000 | |
| python train_from_scratch.py --size medium --dataset file --file corpus.txt --steps 1000 | |
| python train_from_scratch.py --arch gru --engine hf --steps 200 # legacy HF Trainer | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import torch | |
| from datasets import Dataset, load_dataset | |
| from transformers import ( | |
| DataCollatorForLanguageModeling, | |
| PreTrainedTokenizerFast, | |
| Trainer, | |
| TrainingArguments, | |
| ) | |
| from agron_nets.training import add_training_args, build_model_from_args, trainer_from_args | |
| from agron_nets.utils import count_parameters, format_parameters | |
| # Base architecture presets | |
| PRESETS = { | |
| "small": {"vocab_size": 16384, "hidden_size": 256, "intermediate_size": 1024, "num_hidden_layers": 4, "num_attention_heads": 4, "num_key_value_heads": 4}, | |
| "current": {"vocab_size": 16384, "hidden_size": 512, "intermediate_size": 2048, "num_hidden_layers": 8, "num_attention_heads": 4, "num_key_value_heads": 4}, | |
| "medium": {"vocab_size": 16384, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 12, "num_attention_heads": 6, "num_key_value_heads": 6}, | |
| } | |
| ROOT = os.path.dirname(os.path.abspath(__file__)) | |
| TOKENIZER_FILE = os.path.join(ROOT, "tokenizer.json") | |
| def _work_dir(): | |
| """Writable scratch dir for training artifacts.""" | |
| for cand in [os.environ.get("AGRON_WORK_DIR"), "/content", os.path.join(ROOT, "workspace")]: | |
| if not cand: | |
| continue | |
| try: | |
| os.makedirs(cand, exist_ok=True) | |
| probe = os.path.join(cand, ".agron_write_probe") | |
| with open(probe, "w") as f: | |
| f.write("ok") | |
| os.remove(probe) | |
| return cand | |
| except Exception: | |
| continue | |
| return ROOT | |
| def tokenize_dataset(tok, text_list, block_size=256): | |
| tokenized = tok(text_list, truncation=False)["input_ids"] | |
| chunks = [] | |
| buf = [] | |
| for ids in tokenized: | |
| buf.extend(ids) | |
| while len(buf) >= block_size: | |
| chunks.append(buf[:block_size]) | |
| buf = buf[block_size:] | |
| if len(buf) >= block_size // 2: | |
| chunks.append(buf + [tok.pad_token_id] * (block_size - len(buf))) | |
| if not chunks: | |
| raise RuntimeError("no training chunks produced") | |
| return Dataset.from_dict({"input_ids": chunks, "labels": [c[:] for c in chunks]}) | |
| def make_loader(dataset, tok, batch_size=8, shuffle=True, num_workers=0): | |
| """Torch DataLoader whose batches are dicts the Agron engine understands.""" | |
| def collate(rows): | |
| input_ids = torch.tensor([r["input_ids"] for r in rows], dtype=torch.long) | |
| labels = input_ids.clone() | |
| labels = labels.masked_fill(labels == tok.pad_token_id, -100) | |
| return { | |
| "input_ids": input_ids, | |
| "labels": labels, | |
| "attention_mask": torch.ones_like(input_ids), | |
| } | |
| return torch.utils.data.DataLoader( | |
| dataset, | |
| batch_size=batch_size, | |
| shuffle=shuffle, | |
| num_workers=num_workers, | |
| collate_fn=collate, | |
| drop_last=True, | |
| ) | |
| def load_texts(args): | |
| if args.dataset == "wikitext": | |
| ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1") | |
| train_texts = [x["text"] for x in ds["train"] if x["text"].strip()] | |
| eval_texts = [x["text"] for x in ds["validation"] if x["text"].strip()] | |
| return train_texts, eval_texts | |
| if not args.file or not os.path.exists(args.file): | |
| raise RuntimeError("Provide existing --file for --dataset file") | |
| with open(args.file, "r", encoding="utf-8") as f: | |
| texts = [ln.strip() for ln in f if ln.strip()] | |
| split = max(1, int(len(texts) * 0.95)) | |
| return texts[:split], texts[split:] | |
| def main(): | |
| ap = argparse.ArgumentParser(description="Train an Agron architecture from scratch.") | |
| ap.add_argument("--size", choices=list(PRESETS), default="small") | |
| ap.add_argument("--steps", type=int, default=500) | |
| ap.add_argument("--epochs", type=int, default=1) | |
| ap.add_argument("--dataset", choices=["wikitext", "file"], default="wikitext") | |
| ap.add_argument("--file", type=str, default=None, help="text file for --dataset file (one doc per line or paragraph)") | |
| ap.add_argument("--output", type=str, default=None) | |
| ap.add_argument("--block-size", type=int, default=256) | |
| ap.add_argument("--engine", choices=["agron", "hf"], default="agron", | |
| help="agron = compose optimiser/scheduler/algorithms (default); hf = legacy Trainer") | |
| ap.add_argument("--num-workers", type=int, default=0) | |
| add_training_args(ap) | |
| args = ap.parse_args() | |
| tok = PreTrainedTokenizerFast(tokenizer_file=TOKENIZER_FILE) | |
| tok.pad_token = "<pad>"; tok.bos_token = "<s>"; tok.eos_token = "</s>"; tok.unk_token = "<unk>" | |
| model, config = build_model_from_args(args, preset=PRESETS[args.size]) | |
| print(f"Initialized from-scratch model: {args.size} / {config.model_type} | " | |
| f"{format_parameters(count_parameters(model))} params") | |
| train_texts, eval_texts = load_texts(args) | |
| train_ds = tokenize_dataset(tok, train_texts, block_size=args.block_size) | |
| print(f"Training chunks: {len(train_ds):,}") | |
| out_dir = args.output or os.path.join(_work_dir(), f"agron-scratch-{args.size}-{config.model_type}") | |
| os.makedirs(out_dir, exist_ok=True) | |
| if args.engine == "hf": | |
| if config.model_type != "llama": | |
| raise SystemExit("--engine hf only supports the llama architecture; use --engine agron") | |
| ta = TrainingArguments( | |
| output_dir=out_dir, | |
| per_device_train_batch_size=args.batch_size, | |
| max_steps=args.steps, | |
| learning_rate=args.lr if args.lr is not None else 5e-05, | |
| optim="adamw_torch_fused", | |
| fp16=True if torch.cuda.is_available() else False, | |
| logging_steps=max(1, args.log_every), | |
| save_strategy="steps", | |
| save_steps=max(100, args.steps // 2), | |
| save_total_limit=2, | |
| report_to=[], | |
| remove_unused_columns=True, | |
| seed=args.seed, | |
| ) | |
| collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False) | |
| trainer = Trainer(model=model, args=ta, train_dataset=train_ds, data_collator=collator, processing_class=tok) | |
| trainer.train() | |
| else: | |
| trainer = trainer_from_args(args, model, output_dir=out_dir) | |
| eval_loader = None | |
| if eval_texts: | |
| eval_ds = tokenize_dataset(tok, eval_texts, block_size=args.block_size) | |
| eval_loader = make_loader(eval_ds, tok, batch_size=args.batch_size, shuffle=False, | |
| num_workers=args.num_workers) | |
| report = trainer.fit( | |
| make_loader(train_ds, tok, batch_size=args.batch_size, num_workers=args.num_workers), | |
| eval_dataloader=eval_loader, | |
| ) | |
| print(json.dumps({k: v for k, v in report.items() if k != "history"}, indent=2, default=str)) | |
| final = os.path.join(out_dir, "final") | |
| os.makedirs(final, exist_ok=True) | |
| model.save_pretrained(final) | |
| tok.save_pretrained(final) | |
| print("From-scratch training complete:", final) | |
| if __name__ == "__main__": | |
| main() | |