File size: 4,056 Bytes
a006332 | 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 | """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, 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"})
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(),
)
trainer = DeltaTrainer(
model=model,
args=training_args,
train_dataset=dataset,
data_collator=DeltaDataCollator(tokenizer.pad_token_id),
)
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
|