Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """Fine-tune the published Agron model (or any architecture) with the full zoo. | |
| Examples: | |
| python train_agron.py --steps 500 # continue the published LLaMA | |
| python train_agron.py --arch hybrid --recipe stable --steps 1000 | |
| python train_agron.py --arch mamba --optimizer muon --scheduler wsd \\ | |
| --algorithm gradient_clip --algorithm ema --steps 2000 | |
| python train_agron.py --engine hf --steps 500 # legacy HF Trainer | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import torch | |
| from transformers import ( | |
| DataCollatorForLanguageModeling, | |
| LlamaConfig, | |
| LlamaForCausalLM, | |
| 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 | |
| from train_from_scratch import PRESETS, load_texts, make_loader, tokenize_dataset | |
| MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| TOKENIZER_FILE = os.path.join(MODEL_DIR, "tokenizer.json") | |
| OUTPUT_DIR = os.getenv("AGRON_OUTPUT_DIR", os.path.join(MODEL_DIR, "agron-model")) | |
| def build_tokenizer(): | |
| tok = PreTrainedTokenizerFast(tokenizer_file=TOKENIZER_FILE) | |
| tok.pad_token = "<pad>" | |
| tok.bos_token = "<s>" | |
| tok.eos_token = "</s>" | |
| tok.unk_token = "<unk>" | |
| tok.mask_token = "<mask>" | |
| return tok | |
| def build_published_model(): | |
| """The exact config/weights that were published to the Hub.""" | |
| cfg = LlamaConfig( | |
| vocab_size=16384, | |
| hidden_size=512, | |
| intermediate_size=2048, | |
| num_hidden_layers=8, | |
| num_attention_heads=4, | |
| num_key_value_heads=4, | |
| max_position_embeddings=512, | |
| bos_token_id=0, | |
| eos_token_id=2, | |
| pad_token_id=1, | |
| tie_word_embeddings=False, | |
| ) | |
| model = LlamaForCausalLM(cfg) | |
| weights = os.path.join(MODEL_DIR, "model.safetensors") | |
| if os.path.exists(weights): | |
| from safetensors.torch import load_file | |
| model.load_state_dict(load_file(weights)) | |
| print("Loaded published weights from", weights) | |
| return model | |
| def main(): | |
| ap = argparse.ArgumentParser(description="Continue training an Agron model.") | |
| 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) | |
| ap.add_argument("--output", type=str, default=OUTPUT_DIR) | |
| ap.add_argument("--block-size", type=int, default=512) | |
| ap.add_argument("--engine", choices=["agron", "hf"], default="agron") | |
| ap.add_argument("--init", choices=["published", "scratch"], default="published", | |
| help="start from the published LLaMA weights or fresh weights") | |
| ap.add_argument("--num-workers", type=int, default=2) | |
| add_training_args(ap) | |
| args = ap.parse_args() | |
| tok = build_tokenizer() | |
| arch = args.arch or "llama" | |
| if args.init == "published" and arch == "llama": | |
| model = build_published_model() | |
| config = model.config | |
| else: | |
| if args.init == "published": | |
| print(f"--init published is only meaningful for --arch llama; initialising {arch} from scratch") | |
| model, config = build_model_from_args(args, preset=PRESETS["current"]) | |
| print(f"Model: {config.model_type} | {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) | |
| eval_ds = tokenize_dataset(tok, eval_texts, block_size=args.block_size) if eval_texts else None | |
| print(f"Training chunks: {len(train_ds):,}") | |
| out_dir = args.output | |
| 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, | |
| num_train_epochs=args.epochs, | |
| max_steps=args.steps, | |
| learning_rate=args.lr if args.lr is not None else 5e-05, | |
| lr_scheduler_type="linear", | |
| warmup_steps=args.warmup, | |
| optim="adamw_torch_fused", | |
| weight_decay=args.weight_decay if args.weight_decay is not None else 0.0, | |
| fp16=True if torch.cuda.is_available() else False, | |
| gradient_accumulation_steps=args.grad_accum, | |
| max_grad_norm=args.grad_clip or 0.0, | |
| logging_steps=max(1, args.log_every), | |
| logging_strategy="steps", | |
| save_strategy="steps", | |
| save_steps=max(100, args.steps // 2), | |
| save_total_limit=2, | |
| dataloader_num_workers=args.num_workers, | |
| remove_unused_columns=True, | |
| report_to=[], | |
| seed=args.seed, | |
| do_train=True, | |
| ) | |
| 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) | |
| report = trainer.fit( | |
| make_loader(train_ds, tok, batch_size=args.batch_size, num_workers=args.num_workers), | |
| eval_dataloader=( | |
| make_loader(eval_ds, tok, batch_size=args.batch_size, shuffle=False, | |
| num_workers=args.num_workers) | |
| if eval_ds is not None else None | |
| ), | |
| ) | |
| print(json.dumps({k: v for k, v in report.items() if k != "history"}, indent=2, default=str)) | |
| final_dir = os.path.join(out_dir, "final") | |
| os.makedirs(final_dir, exist_ok=True) | |
| model.save_pretrained(final_dir) | |
| tok.save_pretrained(final_dir) | |
| print("Training complete. Final model saved to", final_dir) | |
| if __name__ == "__main__": | |
| main() | |