File size: 3,712 Bytes
ea7e326 | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
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)
# Prompt masking is approximated at the serialized boundary:
# train on the complete solution serialization while excluding PAD.
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()
|