|
|
| import argparse, yaml, random, os, json, math |
| from pathlib import Path |
|
|
| def main(): |
| ap=argparse.ArgumentParser() |
| ap.add_argument("--config",default="configs/lamm_tiny.yaml") |
| args=ap.parse_args() |
| cfg=yaml.safe_load(open(args.config)) |
|
|
| try: |
| import torch |
| from lamm.model import LAMMModel |
| from lamm.data import stream_hf, batched, format_sft |
| from lamm.tokenizer import ByteTokenizer |
| except Exception as e: |
| raise SystemExit("Install training dependencies: pip install -e '.[train]'") from e |
|
|
| seed=cfg["seed"] |
| random.seed(seed); torch.manual_seed(seed) |
| device=torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| tok=ByteTokenizer() |
| model=LAMMModel( |
| vocab_size=tok.vocab_size, |
| d_model=cfg["model"]["d_model"], |
| n_heads=cfg["model"]["n_heads"], |
| d_ff=cfg["model"]["d_ff"], |
| n_layers=cfg["model"]["n_layers"], |
| lam=cfg["model"]["state_lambda"], |
| ).to(device) |
|
|
| opt=torch.optim.AdamW( |
| model.parameters(), |
| lr=cfg["training"]["learning_rate"], |
| weight_decay=cfg["training"]["weight_decay"] |
| ) |
|
|
| out=Path(cfg["training"]["output_dir"]) |
| out.mkdir(parents=True,exist_ok=True) |
| (out/"config.json").write_text(json.dumps(cfg,indent=2)) |
|
|
| stream=stream_hf( |
| cfg["data"]["dataset"], |
| cfg["data"]["split"], |
| cfg["data"].get("source") |
| ) |
|
|
| max_len=cfg["model"]["max_length"] |
| batch_size=cfg["training"]["batch_size"] |
| grad_acc=cfg["training"]["gradient_accumulation"] |
| max_steps=cfg["training"]["max_steps"] |
| ckpt_every=cfg["training"]["checkpoint_every"] |
|
|
| def make_batch(records): |
| seqs=[tok.encode(format_sft(r),max_len) for r in records] |
| L=max(len(x) for x in seqs) |
| L=min(L,max_len) |
| ids=torch.full((len(seqs),L),tok.PAD,dtype=torch.long) |
| labels=torch.full_like(ids,-100) |
| for i,s in enumerate(seqs): |
| s=s[:L] |
| ids[i,:len(s)]=torch.tensor(s) |
| |
| |
| labels[i,:len(s)]=torch.tensor(s) |
| return ids.to(device), labels.to(device) |
|
|
| model.train() |
| running=0.0 |
| step=0 |
| seen=0 |
|
|
| for batch in batched(stream,batch_size): |
| ids,labels=make_batch(batch) |
| outp=model(ids,labels) |
| loss=outp["loss"]/grad_acc |
| loss.backward() |
| running += float(loss.item()) |
| seen += len(batch) |
|
|
| if (step+1)%grad_acc==0: |
| torch.nn.utils.clip_grad_norm_(model.parameters(),1.0) |
| opt.step(); opt.zero_grad(set_to_none=True) |
|
|
| step += 1 |
| if step%25==0: |
| print(json.dumps({ |
| "step":step, |
| "loss":running/max(1,min(25,step)), |
| "seen":seen, |
| "device":str(device) |
| })) |
| running=0.0 |
|
|
| if step%ckpt_every==0: |
| ckpt=out/f"step-{step}.pt" |
| torch.save({ |
| "model":model.state_dict(), |
| "optimizer":opt.state_dict(), |
| "step":step, |
| "seen":seen |
| },ckpt) |
| print("checkpoint",ckpt) |
|
|
| if step>=max_steps: |
| break |
|
|
| ckpt=out/"final.pt" |
| torch.save({ |
| "model":model.state_dict(), |
| "optimizer":opt.state_dict(), |
| "step":step, |
| "seen":seen |
| },ckpt) |
| print(json.dumps({ |
| "status":"TRAINING_COMPLETE", |
| "steps":step, |
| "samples_seen":seen, |
| "checkpoint":str(ckpt), |
| "device":str(device) |
| })) |
|
|
| if __name__=="__main__": |
| main() |
|
|
