File size: 5,590 Bytes
e648e16 | 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 136 137 138 139 140 141 142 143 144 145 146 147 | """Training utilities for Delta Ultra Mini using HuggingFace Trainer."""
from __future__ import annotations
import logging
import os
import json
from pathlib import Path
from typing import Any
import torch
from transformers import Trainer, TrainerCallback, TrainingArguments
from transformers.utils import SAFE_WEIGHTS_NAME
from delta.dataset import DeltaDataCollator, DeltaDataset
from delta.model import DeltaConfig, DeltaModel
from delta.tokenizer import load_tokenizer
logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)
class DeltaTrainer(Trainer):
"""Trainer override that saves tied weights safely with safetensors."""
def _save(self, output_dir: str | None = None, state_dict: dict[str, torch.Tensor] | None = None) -> None:
output_dir = output_dir or self.args.output_dir
os.makedirs(output_dir, exist_ok=True)
state_dict = state_dict or self.model.state_dict()
weights_path = os.path.join(output_dir, SAFE_WEIGHTS_NAME)
try:
import safetensors.torch
safetensors.torch.save_file(state_dict, weights_path, metadata={"format": "pt"})
except RuntimeError as exc:
# When embeddings are tied, tensors share the same underlying storage.
# `save_file(state_dict, ...)` refuses that; `save_model(model, ...)` handles it.
if "share memory" not in str(exc):
raise
import safetensors.torch
safetensors.torch.save_model(self.model, weights_path, metadata={"format": "pt"})
class DeltaProgressCallback(TrainerCallback):
"""Print step heartbeats before tqdm can advance."""
def __init__(self, every: int = 10) -> None:
self.every = max(1, every)
self._step_start: float | None = None
def _should_print(self, step: int) -> bool:
return step <= 5 or step % self.every == 0
def on_train_begin(self, args: TrainingArguments, state, control, **kwargs):
total = int(state.max_steps or 0)
print(
f"Delta training started: {total} steps, device={args.device}, "
f"batch_size={args.per_device_train_batch_size}",
flush=True,
)
def on_step_begin(self, args: TrainingArguments, state, control, **kwargs):
import time
step = int(state.global_step) + 1
self._step_start = time.perf_counter()
if self._should_print(step):
total = int(state.max_steps or 0)
print(f"Delta step {step}/{total} running...", flush=True)
def on_step_end(self, args: TrainingArguments, state, control, **kwargs):
import time
step = int(state.global_step)
if self._step_start is None or not self._should_print(step):
return
elapsed = time.perf_counter() - self._step_start
total = int(state.max_steps or 0)
print(f"Delta step {step}/{total} done in {elapsed:.1f}s", flush=True)
def train(config_dict: dict[str, Any]) -> Trainer:
"""Train Delta Ultra Mini.
Args:
config_dict: Training and model configuration.
Returns:
The configured Trainer after training.
"""
data_path = Path(config_dict["data_path"])
output_dir = Path(config_dict["output_dir"])
tokenizer_path = Path(config_dict.get("tokenizer_path", output_dir / "tokenizer.json"))
epochs = float(config_dict.get("epochs", 1))
batch_size = int(config_dict.get("batch_size", 2))
resume_from_checkpoint = config_dict.get("resume_from_checkpoint")
model_config = DeltaConfig.from_dict(config_dict.get("model", config_dict))
tokenizer = load_tokenizer(tokenizer_path)
dataset = DeltaDataset(
data_path=data_path,
tokenizer=tokenizer,
max_seq_len=model_config.max_seq_len,
stride=int(config_dict.get("stride", 256)),
)
model = DeltaModel(model_config)
training_args = TrainingArguments(
output_dir=str(output_dir),
num_train_epochs=epochs,
per_device_train_batch_size=batch_size,
learning_rate=float(config_dict.get("learning_rate", 3e-4)),
weight_decay=float(config_dict.get("weight_decay", 0.01)),
warmup_steps=int(config_dict.get("warmup_steps", 500)),
lr_scheduler_type="cosine",
max_grad_norm=float(config_dict.get("max_grad_norm", 1.0)),
fp16=torch.cuda.is_available(),
save_steps=int(config_dict.get("save_steps", 1000)),
logging_steps=int(config_dict.get("logging_steps", 50)),
report_to=[],
remove_unused_columns=False,
dataloader_pin_memory=torch.cuda.is_available(),
logging_first_step=True,
)
trainer = DeltaTrainer(
model=model,
args=training_args,
train_dataset=dataset,
data_collator=DeltaDataCollator(tokenizer.pad_token_id),
callbacks=[DeltaProgressCallback(int(config_dict.get("progress_every", 10)))],
)
trainer.train(resume_from_checkpoint=resume_from_checkpoint)
trainer.save_model(str(output_dir))
output_dir.mkdir(parents=True, exist_ok=True)
model.save_checkpoint(
output_dir / "delta_checkpoint.pt",
optimizer=trainer.optimizer,
scheduler=trainer.lr_scheduler,
step=int(trainer.state.global_step),
)
with (output_dir / "config.json").open("w", encoding="utf-8") as handle:
json.dump(model_config.to_dict(), handle, ensure_ascii=False, indent=2)
logger.info("Saved Delta checkpoint to %s", output_dir / "delta_checkpoint.pt")
return trainer
|