File size: 2,603 Bytes
492340a 5792a99 492340a 5792a99 492340a 5792a99 492340a 5792a99 492340a 5792a99 492340a | 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 | """HF Jobs compatible tiny-train: reads /data, writes /data/outputs + Hub."""
import argparse, json, os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
from pathlib import Path
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--data", default=os.getenv("ALFA_DATA", "/data/processed/train.jsonl"))
ap.add_argument("--out", default=os.getenv("ALFA_OUT", "/data/outputs/tiny-50M"))
ap.add_argument("--epochs", type=int, default=1)
ap.add_argument("--cpu", action="store_true", help="force CPU (no_cuda)")
args = ap.parse_args()
import torch
from transformers import LlamaConfig, LlamaForCausalLM, Trainer, TrainingArguments
from torch.utils.data import Dataset
class JsonlDS(Dataset):
def __init__(self, path, tok_len=256):
rows = [json.loads(l) for l in open(path, encoding="utf-8")]
self.rows, self.tok_len = rows, tok_len
def __len__(self): return len(self.rows)
def __getitem__(self, i):
r = self.rows[i]
s = (r.get("prompt", "") + "\n" + r.get("answer", ""))[:2000]
ids = [min(ord(c), 30000) for c in s][:self.tok_len]
ids += [0] * (self.tok_len - len(ids))
return {"input_ids": torch.tensor(ids, dtype=torch.long),
"labels": torch.tensor(ids, dtype=torch.long)}
print(f"data: {args.data} exists={Path(args.data).exists()}")
cfg = LlamaConfig(hidden_size=512, num_hidden_layers=8, num_attention_heads=8,
num_key_value_heads=4, intermediate_size=1376,
vocab_size=32000, max_position_embeddings=2048)
model = LlamaForCausalLM(cfg)
ds = JsonlDS(args.data)
use_cuda = torch.cuda.is_available() and not args.cpu
targs = TrainingArguments(output_dir=args.out, per_device_train_batch_size=1,
gradient_accumulation_steps=4,
num_train_epochs=args.epochs, logging_steps=2,
save_steps=50, save_total_limit=1,
fp16=use_cuda, use_cpu=args.cpu, report_to="none",
gradient_checkpointing=True,
gradient_checkpointing_kwargs={"use_reentrant": False},
dataloader_num_workers=0, dataloader_pin_memory=False,
max_grad_norm=0.0)
Trainer(model=model, args=targs, train_dataset=ds).train()
model.save_pretrained(args.out)
print(f"saved -> {args.out}")
if __name__ == "__main__":
main()
|