File size: 4,409 Bytes
24ebd71 | 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 127 128 129 130 131 132 133 134 135 | from __future__ import annotations
import argparse
import json
import os
import time
import trackio
from snip_common import (
ARTIFACT_DIR,
DATA_DIR,
build_tokenizer,
make_model,
parameter_count,
read_texts,
texts_to_blocks,
)
from transformers import (
DataCollatorForLanguageModeling,
Trainer,
TrainerCallback,
TrainingArguments,
set_seed,
)
class DiagnosticCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
if not logs:
return
loss = logs.get("loss")
if loss is not None and loss != loss:
trackio.alert(
title="NaN loss",
text=f"Training produced NaN at step {state.global_step}.",
level=trackio.AlertLevel.ERROR,
)
if loss is not None and state.global_step >= 200 and loss > 7:
trackio.alert(
title="High loss",
text=f"Loss is {loss:.4f} at step {state.global_step}.",
level=trackio.AlertLevel.WARN,
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--max-steps", type=int, default=800)
parser.add_argument("--batch-size", type=int, default=16)
parser.add_argument("--learning-rate", type=float, default=8e-4)
parser.add_argument("--resume-from-checkpoint")
args = parser.parse_args()
set_seed(42)
os.environ.setdefault("TRACKIO_PROJECT", "snip-model-foundry")
train_texts = read_texts(DATA_DIR / "train.jsonl")
eval_texts = read_texts(DATA_DIR / "eval.jsonl")
tokenizer = build_tokenizer(train_texts)
train_dataset = texts_to_blocks(train_texts, tokenizer)
eval_dataset = texts_to_blocks(eval_texts, tokenizer)
model = make_model(tokenizer)
parameters = parameter_count(model)
ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
started = time.perf_counter()
training_args = TrainingArguments(
output_dir=str(ARTIFACT_DIR / "checkpoints"),
max_steps=args.max_steps,
per_device_train_batch_size=args.batch_size,
per_device_eval_batch_size=args.batch_size,
gradient_accumulation_steps=1,
learning_rate=args.learning_rate,
warmup_steps=max(1, int(args.max_steps * 0.05)),
weight_decay=0.01,
lr_scheduler_type="cosine",
eval_strategy="steps",
eval_steps=100,
logging_steps=20,
save_strategy="steps",
save_steps=200,
save_total_limit=2,
report_to="trackio",
project="snip-model-foundry",
run_name="snip-0.4m-pretrain-v1",
use_cpu=True,
dataloader_num_workers=0,
remove_unused_columns=False,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
processing_class=tokenizer,
callbacks=[DiagnosticCallback()],
)
result = trainer.train(
resume_from_checkpoint=args.resume_from_checkpoint or None,
)
elapsed = time.perf_counter() - started
trainer.save_model(ARTIFACT_DIR)
tokenizer.save_pretrained(ARTIFACT_DIR)
logged_losses = [
float(entry["loss"]) for entry in trainer.state.log_history if "loss" in entry
]
evaluations = [entry for entry in trainer.state.log_history if "eval_loss" in entry]
if not evaluations:
raise RuntimeError("Training completed without a recorded evaluation.")
final_evaluation = evaluations[-1]
summary = {
"model": "SNIP-0.4M",
"parameters": parameters,
"train_examples": len(train_dataset),
"eval_examples": len(eval_dataset),
"max_steps": args.max_steps,
"train_loss": sum(logged_losses) / len(logged_losses),
"trainer_reported_loss": float(result.training_loss),
"eval_loss": float(final_evaluation["eval_loss"]),
"continuation_elapsed_seconds": elapsed,
"resumed_from": args.resume_from_checkpoint,
"tokens_seen": args.max_steps * args.batch_size * 128,
}
(ARTIFACT_DIR / "training_summary.json").write_text(
json.dumps(summary, indent=2),
encoding="utf-8",
)
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()
|