| """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: |
| |
| |
| 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 |
|
|
