File size: 2,290 Bytes
d782871 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | #!/usr/bin/env python3
"""Train one TinyLlama variant from a YAML config."""
import argparse
import yaml
import torch
from transformers import AutoTokenizer, set_seed
from exp import TinyLlamaConfig, TinyLlamaForCausalLM, build_dataset, create_trainer
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", required=True, help="Path to YAML config")
parser.add_argument("--push", action="store_true", help="Push final model to HF Hub")
args = parser.parse_args()
with open(args.config) as f:
cfg = yaml.safe_load(f)
# Explicit seed before any randomness
seed = cfg.get("training", {}).get("seed", 42)
set_seed(seed)
model_cfg = cfg["model"]
train_cfg = cfg.get("training", {})
# Tokenizer
tok_name = model_cfg.pop("tokenizer_name", "meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained(tok_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# Model – the config must contain mlp_type and activation
tiny_config = TinyLlamaConfig(**model_cfg)
# Choose backend: "sdpa" (default, fast) or "flash_attention_2" (if installed)
tiny_config._attn_implementation = "sdpa"
model = TinyLlamaForCausalLM(tiny_config)
model = model.to(torch.bfloat16)
n_params = sum(p.numel() for p in model.parameters()) / 1e6
print(f"Model: {n_params:.2f}M params | MLP type: {tiny_config.mlp_type} | Activation: {tiny_config.activation}")
print(f"Attention implementation: {model.config._attn_implementation}")
# Optional: verify backend if you expect a specific one
# assert model.config._attn_implementation == "sdpa", "Backend mismatch"
# Data
msl = model_cfg.get("max_position_embeddings", 512)
train_ds = build_dataset(tokenizer, max_seq_len=msl, split="train")
eval_ds = build_dataset(tokenizer, max_seq_len=msl, split="validation")
# Train
trainer = create_trainer(model, tokenizer, cfg, train_ds, eval_ds)
trainer.train()
# Save & push
out = train_cfg.get("output_dir", "./out")
trainer.save_model(out)
if args.push or train_cfg.get("push_to_hub", False):
trainer.push_to_hub()
print(f"Done. Artifacts in {out}")
if __name__ == "__main__":
main()
|