1159 lines
52 KiB
Python
1159 lines
52 KiB
Python
"""
|
||
Unsloth Training Backend
|
||
Integrates Unsloth training capabilities with the FastAPI backend
|
||
"""
|
||
import os
|
||
# Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
import torch
|
||
from utils.hardware import clear_gpu_cache, safe_num_proc
|
||
torch._dynamo.config.recompile_limit = 64
|
||
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
|
||
from unsloth.chat_templates import get_chat_template
|
||
|
||
import json
|
||
import threading
|
||
import math
|
||
import logging
|
||
import time
|
||
from typing import Optional, Callable
|
||
from dataclasses import dataclass
|
||
import pandas as pd
|
||
from datasets import Dataset, load_dataset
|
||
|
||
# Add the parent directory to sys.path to import unsloth modules
|
||
#sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
||
from utils.models import is_vision_model
|
||
from utils.datasets import format_and_template_dataset
|
||
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
|
||
from trl import SFTTrainer, SFTConfig
|
||
|
||
# Import Unsloth trainers
|
||
#from unsloth_compiled_cache.UnslothSFTTrainer import _UnslothSFTTrainer as SFTTrainer
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
@dataclass
|
||
class TrainingProgress:
|
||
"""Training progress tracking"""
|
||
epoch: float = 0
|
||
step: int = 0
|
||
total_steps: int = 0
|
||
loss: float = 0.0
|
||
learning_rate: float = 0.0
|
||
is_training: bool = False
|
||
is_completed: bool = False
|
||
error: Optional[str] = None
|
||
status_message: str = "Ready to train" # Current stage message
|
||
elapsed_seconds: Optional[float] = None
|
||
eta_seconds: Optional[float] = None
|
||
grad_norm: Optional[float] = None
|
||
num_tokens: Optional[int] = None
|
||
eval_loss: Optional[float] = None
|
||
|
||
class UnslothTrainer:
|
||
"""
|
||
Unsloth Training Backend
|
||
"""
|
||
|
||
def __init__(self):
|
||
self.model = None
|
||
self.tokenizer = None
|
||
self.trainer = None
|
||
self.training_thread = None
|
||
self.training_progress = TrainingProgress()
|
||
self.progress_callbacks = []
|
||
self.is_training = False
|
||
self.should_stop = False
|
||
self.save_on_stop = True
|
||
|
||
# Model state tracking
|
||
self.is_vlm = False
|
||
self.model_name = None
|
||
|
||
# Training metrics tracking
|
||
self.training_start_time: Optional[float] = None
|
||
self.batch_size: Optional[int] = None
|
||
self.max_seq_length: Optional[int] = None
|
||
self.gradient_accumulation_steps: Optional[int] = None
|
||
|
||
# Thread safety
|
||
self._lock = threading.Lock()
|
||
|
||
# Store training context for later transfer
|
||
self.training_context = {
|
||
'base_model_name': None,
|
||
'output_dir': None,
|
||
'is_lora': True, # Default to LoRA
|
||
}
|
||
|
||
def add_progress_callback(self, callback: Callable[[TrainingProgress], None]):
|
||
"""Add callback for training progress updates"""
|
||
self.progress_callbacks.append(callback)
|
||
|
||
def _update_progress(self, **kwargs):
|
||
"""Update training progress and notify callbacks"""
|
||
with self._lock:
|
||
for key, value in kwargs.items():
|
||
if hasattr(self.training_progress, key):
|
||
setattr(self.training_progress, key, value)
|
||
|
||
# Notify all callbacks
|
||
for callback in self.progress_callbacks:
|
||
try:
|
||
callback(self.training_progress)
|
||
except Exception as e:
|
||
logger.error(f"Error in progress callback: {e}")
|
||
|
||
def load_model(self,
|
||
model_name: str,
|
||
max_seq_length: int = 2048,
|
||
load_in_4bit: bool = True,
|
||
hf_token: Optional[str] = None,
|
||
is_dataset_multimodal: bool = False) -> bool:
|
||
"""Load model for training (supports both text and vision models)"""
|
||
try:
|
||
if self.model is not None:
|
||
del self.model
|
||
if self.tokenizer is not None:
|
||
del self.tokenizer
|
||
|
||
if self.trainer is not None:
|
||
del self.trainer
|
||
|
||
print("\nClearing GPU memory before training...")
|
||
clear_gpu_cache()
|
||
|
||
# Remove stale compiled cache so the new model gets a fresh one
|
||
from utils.cache_cleanup import clear_unsloth_compiled_cache
|
||
clear_unsloth_compiled_cache()
|
||
|
||
# Detect if this is a vision model AND dataset is multimodal
|
||
# A vision-capable model with a text-only dataset should use FastLanguageModel
|
||
self.is_vlm = is_vision_model(model_name) and is_dataset_multimodal
|
||
self.model_name = model_name
|
||
|
||
logger.info(f"Model architecture is vision: {is_vision_model(model_name)}")
|
||
logger.info(f"Dataset is multimodal: {is_dataset_multimodal}")
|
||
logger.info(f"Using VLM path: {self.is_vlm}")
|
||
|
||
# Reset training state for new run
|
||
self._update_progress(
|
||
is_training=True,
|
||
is_completed=False,
|
||
error=None,
|
||
step=0,
|
||
loss=0.0,
|
||
epoch=0
|
||
)
|
||
|
||
# Update UI immediately with loading message
|
||
model_display = model_name.split('/')[-1] if '/' in model_name else model_name
|
||
self._update_progress(
|
||
status_message=f"Loading {'vision' if self.is_vlm else 'text'} model... {model_display}"
|
||
)
|
||
|
||
print(f"\nLoading {'vision' if self.is_vlm else 'text'} model: {model_name}")
|
||
|
||
# Set HF token if provided
|
||
if hf_token:
|
||
os.environ["HF_TOKEN"] = hf_token
|
||
|
||
|
||
# Branch based on model type
|
||
if self.is_vlm:
|
||
# Load vision model - returns (model, tokenizer)
|
||
self.model, self.tokenizer = FastVisionModel.from_pretrained(
|
||
model_name=model_name,
|
||
max_seq_length=max_seq_length,
|
||
dtype=None, # Auto-detect
|
||
load_in_4bit=load_in_4bit,
|
||
token=hf_token,
|
||
)
|
||
logger.info("Loaded vision model")
|
||
else:
|
||
# Load text model - returns (model, tokenizer)
|
||
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
|
||
model_name=model_name,
|
||
max_seq_length=max_seq_length,
|
||
dtype=None, # Auto-detect
|
||
load_in_4bit=load_in_4bit,
|
||
token=hf_token,
|
||
)
|
||
logger.info("Loaded text model")
|
||
|
||
if self.should_stop:
|
||
return False
|
||
|
||
self._update_progress(status_message="Model loaded successfully")
|
||
print("Model loaded successfully")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error loading model: {e}")
|
||
self._update_progress(error=str(e), is_training=False)
|
||
return False
|
||
|
||
def prepare_model_for_training(self,
|
||
use_lora: bool = True,
|
||
# Vision-specific LoRA parameters (only used if is_vlm=True)
|
||
finetune_vision_layers: bool = True,
|
||
finetune_language_layers: bool = True,
|
||
finetune_attention_modules: bool = True,
|
||
finetune_mlp_modules: bool = True,
|
||
# Standard LoRA parameters
|
||
target_modules: list = None,
|
||
lora_r: int = 16,
|
||
lora_alpha: int = 16,
|
||
lora_dropout: float = 0.0,
|
||
use_gradient_checkpointing: str = "unsloth",
|
||
use_rslora: bool = False,
|
||
use_loftq: bool = False) -> bool:
|
||
"""
|
||
Prepare model for training (with optional LoRA).
|
||
"""
|
||
try:
|
||
if self.model is None:
|
||
raise ValueError("Model not loaded. Call load_model() first.")
|
||
|
||
|
||
# Full finetuning mode - skip PEFT entirely
|
||
if not use_lora:
|
||
self._update_progress(status_message="Full finetuning mode - no LoRA adapters")
|
||
print("Full finetuning mode - training all parameters\n")
|
||
return True
|
||
|
||
# LoRA/QLoRA mode - apply PEFT
|
||
# "all-linear" is a PEFT keyword that targets every linear layer
|
||
if isinstance(target_modules, list) and "all-linear" in target_modules:
|
||
if len(target_modules) == 1:
|
||
target_modules = "all-linear"
|
||
else:
|
||
target_modules = [m for m in target_modules if m != "all-linear"]
|
||
elif target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0):
|
||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
|
||
"gate_proj", "up_proj", "down_proj"]
|
||
|
||
# Validate and normalize gradient_checkpointing
|
||
# Must be one of: True, False, or "unsloth"
|
||
if isinstance(use_gradient_checkpointing, str):
|
||
use_gradient_checkpointing = use_gradient_checkpointing.strip().lower()
|
||
if use_gradient_checkpointing == "" or use_gradient_checkpointing == "unsloth":
|
||
use_gradient_checkpointing = "unsloth"
|
||
elif use_gradient_checkpointing in ("true", "1", "yes"):
|
||
use_gradient_checkpointing = True
|
||
elif use_gradient_checkpointing in ("false", "0", "no"):
|
||
use_gradient_checkpointing = False
|
||
else:
|
||
# Invalid value, default to "unsloth"
|
||
logger.warning(f"Invalid gradient_checkpointing value: {use_gradient_checkpointing}, defaulting to 'unsloth'")
|
||
use_gradient_checkpointing = "unsloth"
|
||
elif use_gradient_checkpointing not in (True, False, "unsloth"):
|
||
# Invalid type or value, default to "unsloth"
|
||
logger.warning(f"Invalid gradient_checkpointing type/value: {use_gradient_checkpointing}, defaulting to 'unsloth'")
|
||
use_gradient_checkpointing = "unsloth"
|
||
|
||
# Verify model is loaded
|
||
if self.model is None:
|
||
error_msg = "Model is None - model was not loaded properly"
|
||
logger.error(error_msg)
|
||
self._update_progress(error=error_msg)
|
||
return False
|
||
|
||
# Check if model has the expected attributes
|
||
if not hasattr(self.model, 'config'):
|
||
error_msg = "Model does not have config attribute - model may not be loaded correctly"
|
||
logger.error(error_msg)
|
||
self._update_progress(error=error_msg)
|
||
return False
|
||
|
||
print(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n")
|
||
print(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n")
|
||
|
||
# Branch based on vision vs text
|
||
if self.is_vlm:
|
||
# Vision model LoRA
|
||
print(f"Vision model LoRA configuration:")
|
||
print(f" - Finetune vision layers: {finetune_vision_layers}")
|
||
print(f" - Finetune language layers: {finetune_language_layers}")
|
||
print(f" - Finetune attention modules: {finetune_attention_modules}")
|
||
print(f" - Finetune MLP modules: {finetune_mlp_modules}\n")
|
||
|
||
self.model = FastVisionModel.get_peft_model(
|
||
self.model,
|
||
finetune_vision_layers=finetune_vision_layers,
|
||
finetune_language_layers=finetune_language_layers,
|
||
finetune_attention_modules=finetune_attention_modules,
|
||
finetune_mlp_modules=finetune_mlp_modules,
|
||
r=lora_r,
|
||
target_modules=target_modules,
|
||
lora_alpha=lora_alpha,
|
||
lora_dropout=lora_dropout,
|
||
bias="none",
|
||
use_gradient_checkpointing=use_gradient_checkpointing,
|
||
random_state=3407,
|
||
use_rslora=use_rslora,
|
||
loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
|
||
)
|
||
else:
|
||
# Text model LoRA
|
||
print(f"Text model LoRA configuration:")
|
||
print(f" - Target modules: {target_modules}\n")
|
||
|
||
self.model = FastLanguageModel.get_peft_model(
|
||
self.model,
|
||
r=lora_r,
|
||
target_modules=target_modules,
|
||
lora_alpha=lora_alpha,
|
||
lora_dropout=lora_dropout,
|
||
bias="none",
|
||
use_gradient_checkpointing=use_gradient_checkpointing,
|
||
random_state=3407,
|
||
use_rslora=use_rslora,
|
||
loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
|
||
)
|
||
|
||
# Check if stopped during LoRA preparation
|
||
if self.should_stop:
|
||
print("Stopped during LoRA configuration\n")
|
||
return False
|
||
|
||
self._update_progress(status_message="LoRA adapters configured")
|
||
print("LoRA adapters configured successfully\n")
|
||
return True
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
import sys
|
||
error_details = f"{type(e).__name__}: {str(e)}" if str(e) else f"{type(e).__name__} (no message)"
|
||
full_traceback = traceback.format_exc()
|
||
logger.error(f"Error preparing model: {error_details}")
|
||
logger.error(f"Full traceback:\n{full_traceback}")
|
||
print(f"\n[ERROR] Error preparing model: {error_details}", file=sys.stderr, flush=True)
|
||
print(f"[ERROR] Full traceback:\n{full_traceback}", file=sys.stderr, flush=True)
|
||
self._update_progress(error=error_details)
|
||
return False
|
||
|
||
def load_and_format_dataset(self,
|
||
dataset_source: str,
|
||
format_type: str = "auto",
|
||
local_datasets: list = None,
|
||
custom_format_mapping: dict = None,
|
||
subset: str = None,
|
||
train_split: str = "train",
|
||
eval_split: str = None) -> Optional[tuple]:
|
||
"""
|
||
Load and prepare dataset for training.
|
||
|
||
Strategy: format first, then split — ensures both train and eval
|
||
portions are properly formatted and templated.
|
||
|
||
Returns:
|
||
Tuple of (dataset_info, eval_dataset) or None on error.
|
||
eval_dataset may be None if no eval split is available.
|
||
"""
|
||
try:
|
||
dataset = None
|
||
eval_dataset = None
|
||
has_separate_eval_source = False # True if eval comes from a separate HF split
|
||
|
||
if local_datasets:
|
||
# Load local datasets
|
||
all_data = []
|
||
for dataset_file in local_datasets:
|
||
# dataset_file may already be an absolute path from routes/training.py
|
||
if os.path.isabs(dataset_file):
|
||
file_path = dataset_file
|
||
else:
|
||
# Fallback: try relative to assets/datasets
|
||
script_dir = Path(__file__).parent.parent
|
||
assets_datasets_dir = script_dir / "assets" / "datasets"
|
||
file_path = assets_datasets_dir / dataset_file
|
||
|
||
if str(file_path).endswith('.json'):
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
if isinstance(data, list):
|
||
all_data.extend(data)
|
||
else:
|
||
all_data.append(data)
|
||
elif str(file_path).endswith('.csv'):
|
||
df = pd.read_csv(file_path)
|
||
all_data.extend(df.to_dict('records'))
|
||
|
||
if all_data:
|
||
dataset = Dataset.from_list(all_data)
|
||
|
||
# Check if stopped during dataset loading
|
||
if self.should_stop:
|
||
print("Stopped during dataset loading\n")
|
||
return None
|
||
|
||
self._update_progress(status_message=f"Loaded {len(all_data)} samples from local files")
|
||
print(f"Loaded {len(all_data)} samples from local files\n")
|
||
|
||
elif dataset_source:
|
||
# Load from Hugging Face
|
||
load_kwargs = {"path": dataset_source, "split": train_split or "train"}
|
||
if subset:
|
||
load_kwargs["name"] = subset
|
||
dataset = load_dataset(**load_kwargs)
|
||
|
||
# Check if stopped during dataset loading
|
||
if self.should_stop:
|
||
print("Stopped during dataset loading\n")
|
||
return None
|
||
|
||
self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}")
|
||
print(f"Loaded dataset from Hugging Face: {dataset_source}\n")
|
||
|
||
# Resolve eval split from a separate HF split (explicit or auto-detected)
|
||
if eval_split:
|
||
# Explicit eval split provided - load it directly
|
||
print(f"Loading explicit eval split: '{eval_split}'\n")
|
||
eval_load_kwargs = {"path": dataset_source, "split": eval_split}
|
||
if subset:
|
||
eval_load_kwargs["name"] = subset
|
||
eval_dataset = load_dataset(**eval_load_kwargs)
|
||
has_separate_eval_source = True
|
||
print(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n")
|
||
else:
|
||
# Auto-detect eval split from HF (returns a separate dataset, or None)
|
||
eval_dataset = self._auto_detect_eval_split_from_hf(
|
||
dataset_source=dataset_source,
|
||
subset=subset,
|
||
)
|
||
if eval_dataset is not None:
|
||
has_separate_eval_source = True
|
||
|
||
if dataset is None:
|
||
raise ValueError("No dataset provided")
|
||
|
||
# Check if stopped before applying template
|
||
if self.should_stop:
|
||
print("Stopped before applying chat template\n")
|
||
return None
|
||
|
||
# ========== FORMAT FIRST ==========
|
||
print(f"Formatting dataset with format_type='{format_type}'...\n")
|
||
|
||
dataset_info = format_and_template_dataset(
|
||
dataset,
|
||
model_name=self.model_name,
|
||
tokenizer=self.tokenizer,
|
||
is_vlm=self.is_vlm,
|
||
format_type=format_type,
|
||
dataset_name=dataset_source,
|
||
custom_format_mapping=custom_format_mapping,
|
||
)
|
||
|
||
# Check if stopped during formatting
|
||
if self.should_stop:
|
||
print("Stopped during dataset formatting\n")
|
||
return None
|
||
|
||
self._update_progress(status_message=f"Dataset formatted and ready for training")
|
||
print(f"Dataset formatted successfully\n")
|
||
|
||
# ========== THEN SPLIT ==========
|
||
if has_separate_eval_source and eval_dataset is not None:
|
||
# Eval came from a separate HF split — format it too
|
||
print(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n")
|
||
eval_info = format_and_template_dataset(
|
||
eval_dataset,
|
||
model_name=self.model_name,
|
||
tokenizer=self.tokenizer,
|
||
is_vlm=self.is_vlm,
|
||
format_type=format_type,
|
||
dataset_name=dataset_source,
|
||
custom_format_mapping=custom_format_mapping,
|
||
)
|
||
eval_dataset = eval_info["dataset"]
|
||
print(f"Eval dataset formatted successfully\n")
|
||
elif not has_separate_eval_source:
|
||
# No separate eval source — split the already-formatted dataset
|
||
formatted_dataset = dataset_info["dataset"]
|
||
split_result = self._resolve_eval_split_from_dataset(formatted_dataset)
|
||
if split_result is not None:
|
||
train_portion, eval_dataset = split_result
|
||
dataset_info["dataset"] = train_portion
|
||
|
||
return (dataset_info, eval_dataset)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error loading dataset: {e}")
|
||
self._update_progress(error=str(e))
|
||
return None
|
||
|
||
def _auto_detect_eval_split_from_hf(self, dataset_source: str,
|
||
subset: str) -> Optional[Dataset]:
|
||
"""Auto-detect an eval split from HF dataset (separate named split only)."""
|
||
try:
|
||
from datasets import get_dataset_split_names
|
||
load_kwargs = {"path": dataset_source}
|
||
if subset:
|
||
load_kwargs["name"] = subset
|
||
available_splits = get_dataset_split_names(**load_kwargs)
|
||
print(f"Available splits: {available_splits}\n")
|
||
|
||
# Check for common eval split names
|
||
for candidate in ["eval", "validation", "valid", "val", "test"]:
|
||
if candidate in available_splits:
|
||
eval_load_kwargs = {"path": dataset_source, "split": candidate}
|
||
if subset:
|
||
eval_load_kwargs["name"] = subset
|
||
candidate_ds = load_dataset(**eval_load_kwargs)
|
||
if len(candidate_ds) >= 16:
|
||
print(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n")
|
||
return candidate_ds
|
||
else:
|
||
print(f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n")
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Could not check dataset splits: {e}")
|
||
|
||
# No separate HF eval split found — caller will handle programmatic splitting
|
||
return None
|
||
|
||
def _resolve_eval_split_from_dataset(self, dataset) -> Optional[tuple]:
|
||
"""Split a dataset into train and eval portions.
|
||
|
||
Returns:
|
||
Tuple of (train_dataset, eval_dataset), or None if dataset too small.
|
||
"""
|
||
MIN_EVAL_ROWS = 16
|
||
MIN_TOTAL_ROWS = 32 # Need at least 16 train + 16 eval
|
||
|
||
n = len(dataset)
|
||
if n < MIN_TOTAL_ROWS:
|
||
print(f"Dataset too small ({n} rows) for eval split, skipping eval\n")
|
||
return None
|
||
|
||
eval_size = max(MIN_EVAL_ROWS, min(128, int(0.05 * n)))
|
||
# Ensure we don't take more than half the dataset
|
||
eval_size = min(eval_size, n // 2)
|
||
|
||
print(f"Auto-splitting: {eval_size} rows for eval from {n} total\n")
|
||
split_result = dataset.train_test_split(test_size=eval_size, seed=3407)
|
||
print(f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n")
|
||
return (split_result['train'], split_result['test'])
|
||
|
||
def start_training(self,
|
||
dataset: Dataset,
|
||
eval_dataset: Dataset = None,
|
||
eval_steps: float = 0.01,
|
||
output_dir: str = "./outputs",
|
||
num_epochs: int = 3,
|
||
learning_rate: float = 5e-5,
|
||
batch_size: int = 2,
|
||
gradient_accumulation_steps: int = 4,
|
||
warmup_steps: int = None,
|
||
warmup_ratio: float = None,
|
||
max_steps: int = 0,
|
||
save_steps: int = 0,
|
||
weight_decay: float = 0.01,
|
||
random_seed: int = 3407,
|
||
packing: bool = False,
|
||
train_on_completions: bool = False,
|
||
enable_wandb: bool = False,
|
||
wandb_project: str = "unsloth-training",
|
||
wandb_token: str = None,
|
||
enable_tensorboard: bool = False,
|
||
tensorboard_dir: str = "runs",
|
||
**kwargs) -> bool:
|
||
"""Start training in a separate thread"""
|
||
|
||
if self.is_training:
|
||
logger.warning("Training already in progress")
|
||
return False
|
||
|
||
|
||
if self.model is None or self.tokenizer is None:
|
||
self._update_progress(error="Model not loaded")
|
||
return False
|
||
|
||
# Start training in separate thread
|
||
self.training_thread = threading.Thread(
|
||
target=self._train_worker,
|
||
args=(dataset,),
|
||
kwargs={
|
||
'output_dir': output_dir,
|
||
'num_epochs': num_epochs,
|
||
'learning_rate': learning_rate,
|
||
'batch_size': batch_size,
|
||
'gradient_accumulation_steps': gradient_accumulation_steps,
|
||
'warmup_steps': warmup_steps,
|
||
'warmup_ratio': warmup_ratio,
|
||
'max_steps': max_steps,
|
||
'save_steps': save_steps,
|
||
'weight_decay': weight_decay,
|
||
'random_seed': random_seed,
|
||
'packing': packing,
|
||
'train_on_completions': train_on_completions,
|
||
'enable_wandb': enable_wandb,
|
||
'wandb_project': wandb_project,
|
||
'wandb_token': wandb_token,
|
||
'enable_tensorboard': enable_tensorboard,
|
||
'tensorboard_dir': tensorboard_dir,
|
||
'eval_dataset': eval_dataset,
|
||
'eval_steps': eval_steps,
|
||
**kwargs
|
||
}
|
||
)
|
||
|
||
self.should_stop = False
|
||
self.is_training = True
|
||
try:
|
||
self.training_thread.start()
|
||
return True
|
||
except Exception as e:
|
||
self.is_training = False
|
||
logger.error(f"Failed to start training thread: {e}")
|
||
return False
|
||
|
||
def _train_worker(self, dataset: Dataset, **training_args):
|
||
"""Worker function for training (runs in separate thread)"""
|
||
try:
|
||
# Store training parameters for metrics calculation
|
||
self.batch_size = training_args.get('batch_size', 2)
|
||
self.max_seq_length = training_args.get('max_seq_length', 2048)
|
||
self.gradient_accumulation_steps = training_args.get('gradient_accumulation_steps', 4)
|
||
|
||
# Set training start time
|
||
self.training_start_time = time.time()
|
||
|
||
self._update_progress(is_training=True, error=None)
|
||
|
||
# Setup logging
|
||
if training_args.get('enable_wandb', False) and training_args.get('wandb_token'):
|
||
os.environ["WANDB_API_KEY"] = training_args['wandb_token']
|
||
import wandb
|
||
wandb.init(project=training_args.get('wandb_project', 'unsloth-training'))
|
||
|
||
# Create output directory
|
||
output_dir = training_args.get('output_dir', './outputs')
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# ========== DATA COLLATOR SELECTION ==========
|
||
# Detect special model types
|
||
model_name_lower = self.model_name.lower()
|
||
is_deepseek_ocr = "deepseek" in model_name_lower and "ocr" in model_name_lower
|
||
|
||
print("Configuring data collator...\n")
|
||
|
||
data_collator = None # Default to built-in data collator
|
||
if is_deepseek_ocr:
|
||
# Special DeepSeek OCR collator - auto-install if needed
|
||
print("Detected DeepSeek OCR model\n")
|
||
# Ensure DeepSeek OCR module is installed
|
||
if not _ensure_deepseek_ocr_installed():
|
||
error_msg = (
|
||
"Failed to install DeepSeek OCR module. "
|
||
"Please install manually: "
|
||
"from huggingface_hub import snapshot_download; "
|
||
"snapshot_download('unsloth/DeepSeek-OCR', local_dir='deepseek_ocr')"
|
||
)
|
||
logger.error(error_msg)
|
||
self._update_progress(error=error_msg, is_training=False)
|
||
return
|
||
|
||
try:
|
||
from backend.data_utils import DeepSeekOCRDataCollator
|
||
|
||
print("Configuring DeepSeek OCR data collator...\n")
|
||
FastVisionModel.for_training(self.model)
|
||
data_collator = DeepSeekOCRDataCollator(
|
||
tokenizer=self.tokenizer,
|
||
model=self.model,
|
||
image_size=640,
|
||
base_size=1024,
|
||
crop_mode=True,
|
||
train_on_responses_only=training_args.get('train_on_completions', False),
|
||
)
|
||
print("DeepSeek OCR data collator configured successfully\n")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to configure DeepSeek OCR collator: {e}")
|
||
error_msg = f"Error configuring DeepSeek OCR: {str(e)}"
|
||
self._update_progress(error=error_msg, is_training=False)
|
||
return
|
||
|
||
elif self.is_vlm:
|
||
# Standard VLM collator
|
||
print("Using UnslothVisionDataCollator for vision model\n")
|
||
from unsloth.trainer import UnslothVisionDataCollator
|
||
|
||
FastVisionModel.for_training(self.model)
|
||
data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
|
||
print("Vision data collator configured\n")
|
||
|
||
# ========== TRAINING CONFIGURATION ==========
|
||
# Handle epochs vs max_steps properly
|
||
max_steps_val = training_args.get('max_steps', 0)
|
||
num_epochs_val = training_args.get('num_epochs', 3)
|
||
|
||
# Handle warmup_steps vs warmup_ratio
|
||
warmup_steps_val = training_args.get('warmup_steps', None)
|
||
warmup_ratio_val = training_args.get('warmup_ratio', None)
|
||
|
||
config_args = {
|
||
"per_device_train_batch_size": training_args.get('batch_size', 2),
|
||
"gradient_accumulation_steps": training_args.get('gradient_accumulation_steps', 4),
|
||
"num_train_epochs": training_args.get('num_epochs', 3), # Default to epochs
|
||
"learning_rate": training_args.get('learning_rate', 2e-4),
|
||
"fp16": not is_bfloat16_supported(),
|
||
"bf16": is_bfloat16_supported(),
|
||
"logging_steps": 1,
|
||
"weight_decay": training_args.get('weight_decay', 0.01),
|
||
"seed": training_args.get('random_seed', 3407),
|
||
"output_dir": output_dir,
|
||
"report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none",
|
||
"include_num_input_tokens_seen": True, # Enable token counting
|
||
"dataset_num_proc": safe_num_proc(max(1, os.cpu_count() // 4)),
|
||
}
|
||
|
||
# Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps
|
||
if warmup_ratio_val is not None:
|
||
config_args["warmup_ratio"] = warmup_ratio_val
|
||
print(f"Using warmup_ratio: {warmup_ratio_val}\n")
|
||
elif warmup_steps_val is not None:
|
||
config_args["warmup_steps"] = warmup_steps_val
|
||
print(f"Using warmup_steps: {warmup_steps_val}\n")
|
||
else:
|
||
# Default to warmup_steps if neither provided
|
||
config_args["warmup_steps"] = 5
|
||
print(f"Using default warmup_steps: 5\n")
|
||
|
||
# Add save_steps if specified
|
||
save_steps_val = training_args.get('save_steps', 0)
|
||
if save_steps_val and save_steps_val > 0:
|
||
config_args["save_steps"] = save_steps_val
|
||
config_args["save_strategy"] = "steps"
|
||
|
||
# If max_steps is specified, use it instead of epochs
|
||
max_steps_val = training_args.get('max_steps', 0)
|
||
if max_steps_val and max_steps_val > 0:
|
||
del config_args["num_train_epochs"] # Remove epochs
|
||
config_args["max_steps"] = max_steps_val # Use steps instead
|
||
print(f"Training for {max_steps_val} steps\n")
|
||
else:
|
||
print(f"Training for {config_args['num_train_epochs']} epochs\n")
|
||
|
||
# ========== EVAL CONFIGURATION ==========
|
||
eval_dataset = training_args.get('eval_dataset', None)
|
||
eval_steps_val = training_args.get('eval_steps', 0.01)
|
||
if eval_dataset is not None:
|
||
config_args["eval_strategy"] = "steps"
|
||
config_args["eval_steps"] = eval_steps_val
|
||
print(f"Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n")
|
||
print(f"Eval dataset: {len(eval_dataset)} rows\n")
|
||
else:
|
||
print("No eval dataset — evaluation disabled\n")
|
||
|
||
# Add model-specific parameters
|
||
# Use optim and lr_scheduler_type from training_args if provided, otherwise use defaults
|
||
optim_value = training_args.get('optim', "adamw_8bit")
|
||
lr_scheduler_type_value = training_args.get('lr_scheduler_type', "linear")
|
||
|
||
if self.is_vlm:
|
||
# Vision-specific config
|
||
print("Configuring vision model training parameters\n")
|
||
# Use provided values or defaults for vision models
|
||
optim_value = training_args.get('optim', "adamw_torch_fused")
|
||
lr_scheduler_type_value = training_args.get('lr_scheduler_type', "cosine")
|
||
config_args.update({
|
||
"optim": optim_value,
|
||
"lr_scheduler_type": lr_scheduler_type_value,
|
||
"gradient_checkpointing": True,
|
||
"gradient_checkpointing_kwargs": {"use_reentrant": False},
|
||
"max_grad_norm": 0.3, # Recommended for vision models
|
||
"remove_unused_columns": False,
|
||
"dataset_text_field": "",
|
||
"dataset_kwargs": {"skip_prepare_dataset": True},
|
||
"max_length": training_args.get('max_seq_length', 2048),
|
||
})
|
||
else:
|
||
print("Configuring text model training parameters\n")
|
||
config_args.update({
|
||
"optim": optim_value,
|
||
"lr_scheduler_type": lr_scheduler_type_value,
|
||
"dataset_text_field": "text",
|
||
})
|
||
|
||
# Only add packing for text models (not DeepSeek OCR which is VLM)
|
||
if not is_deepseek_ocr:
|
||
packing_enabled = training_args.get('packing', False)
|
||
config_args["packing"] = packing_enabled
|
||
print(f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n")
|
||
|
||
print(f"The configuration is: {config_args}")
|
||
|
||
print("Training configuration prepared\n")
|
||
# ========== TRAINER INITIALIZATION ==========
|
||
if self.is_vlm:
|
||
trainer_kwargs = {
|
||
"model": self.model,
|
||
"train_dataset": dataset['dataset'],
|
||
"processing_class": self.tokenizer,
|
||
"data_collator": data_collator,
|
||
"args": SFTConfig(**config_args),
|
||
}
|
||
if eval_dataset is not None:
|
||
trainer_kwargs["eval_dataset"] = eval_dataset
|
||
self.trainer = SFTTrainer(**trainer_kwargs)
|
||
else:
|
||
# For text-only training, if the tokenizer is actually a Processor
|
||
# (e.g., Gemma-3 returns a ProcessorMixin even for text), we must
|
||
# unwrap to the raw tokenizer. Otherwise Unsloth's SFTTrainer detects
|
||
# ProcessorMixin → sets _is_vlm=True → skips _prepare_dataset entirely,
|
||
# and the 'text' column never gets tokenized to 'input_ids'.
|
||
from transformers import ProcessorMixin
|
||
sft_tokenizer = self.tokenizer
|
||
if isinstance(self.tokenizer, ProcessorMixin) and hasattr(self.tokenizer, 'tokenizer'):
|
||
print(f" ⚠️ Unwrapping Processor → raw tokenizer for text-only SFTTrainer")
|
||
sft_tokenizer = self.tokenizer.tokenizer
|
||
|
||
trainer_kwargs = {
|
||
"model": self.model,
|
||
"tokenizer": sft_tokenizer,
|
||
"train_dataset": dataset['dataset'],
|
||
"data_collator": data_collator,
|
||
"args": SFTConfig(**config_args),
|
||
}
|
||
if eval_dataset is not None:
|
||
trainer_kwargs["eval_dataset"] = eval_dataset
|
||
self.trainer = SFTTrainer(**trainer_kwargs)
|
||
print("Trainer initialized\n")
|
||
|
||
# ========== TRAIN ON RESPONSES ONLY ==========
|
||
# Determine if we should train on responses only
|
||
instruction_part = None
|
||
response_part = None
|
||
train_on_responses_enabled = training_args.get('train_on_completions', False)
|
||
|
||
# DeepSeek OCR handles this internally in its collator, so skip
|
||
if train_on_responses_enabled and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'):
|
||
try:
|
||
print("Configuring train on responses only...\n")
|
||
|
||
# Get the template mapping for this model
|
||
model_name_lower = self.model_name.lower()
|
||
|
||
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
|
||
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
|
||
print(f"Detected template: {template_name}\n")
|
||
|
||
if template_name in TEMPLATE_TO_RESPONSES_MAPPER:
|
||
instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["instruction"]
|
||
response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"]
|
||
|
||
print(f"Instruction marker: {instruction_part[:50]}...\n")
|
||
print(f"Response marker: {response_part[:50]}...\n")
|
||
else:
|
||
print(f"No response mapping found for template: {template_name}\n")
|
||
train_on_responses_enabled = False
|
||
else:
|
||
print(f"No template mapping found for model: {self.model_name}\n")
|
||
train_on_responses_enabled = False
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Could not configure train on responses: {e}")
|
||
train_on_responses_enabled = False
|
||
|
||
# Apply train on responses only if we have valid parts
|
||
if train_on_responses_enabled and instruction_part and response_part and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'):
|
||
try:
|
||
from unsloth.chat_templates import train_on_responses_only
|
||
|
||
self.trainer = train_on_responses_only(
|
||
self.trainer,
|
||
instruction_part=instruction_part,
|
||
response_part=response_part,
|
||
num_proc=config_args.get("dataset_num_proc", safe_num_proc(max(1, os.cpu_count() // 4))),
|
||
)
|
||
print("Train on responses only configured successfully\n")
|
||
|
||
# ── Safety net: check if all samples were filtered out ──
|
||
# Unsloth's train_on_responses_only masks non-response
|
||
# tokens with -100. If max_seq_length is too short and the
|
||
# response portion gets truncated away, EVERY sample ends
|
||
# up with all labels == -100 and Unsloth removes them,
|
||
# leaving 0 usable training samples.
|
||
filtered_len = len(self.trainer.train_dataset)
|
||
original_len = len(dataset["dataset"])
|
||
dropped = original_len - filtered_len
|
||
drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0
|
||
|
||
if filtered_len == 0 or drop_pct > 30:
|
||
max_seq = training_args.get('max_seq_length', 2048)
|
||
error_msg = (
|
||
f"{dropped}/{original_len} samples ({drop_pct}%) "
|
||
f"were dropped after applying 'train on responses "
|
||
f"only' — only {filtered_len} remain. This usually "
|
||
f"means max_seq_length ({max_seq}) is too short "
|
||
f"and the response portion is being truncated "
|
||
f"away. Try increasing max_seq_length (e.g. 8192) "
|
||
f"or disabling 'Train on completions'."
|
||
)
|
||
logger.error(error_msg)
|
||
self._update_progress(error=error_msg, is_training=False)
|
||
return
|
||
|
||
if dropped > 0:
|
||
print(
|
||
f"⚠️ {dropped}/{original_len} samples "
|
||
f"({drop_pct}%) were dropped (all labels "
|
||
f"masked). {filtered_len} samples remain.\n"
|
||
)
|
||
print(f"Post-filter dataset size: {filtered_len} samples\n")
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Failed to apply train on responses only: {e}")
|
||
train_on_responses_enabled = False
|
||
else:
|
||
if train_on_responses_enabled and is_deepseek_ocr:
|
||
print("Train on responses handled by DeepSeek OCR collator\n")
|
||
else:
|
||
print("Training on full sequences (including prompts)\n")
|
||
|
||
# Add custom callback for progress tracking
|
||
from transformers import TrainerCallback
|
||
|
||
class ProgressCallback(TrainerCallback):
|
||
def __init__(self, trainer_instance):
|
||
self.trainer_instance = trainer_instance
|
||
|
||
def on_train_begin(self, args, state, control, **kwargs):
|
||
"""Called at the beginning of training"""
|
||
pass
|
||
|
||
def on_log(self, args, state, control, logs=None, **kwargs):
|
||
"""Called when logging occurs"""
|
||
if logs:
|
||
# Get loss from either 'loss' or 'train_loss' key
|
||
loss_value = logs.get('loss', logs.get('train_loss', 0.0))
|
||
current_step = state.global_step
|
||
|
||
# Extract grad_norm from logs (available when gradient clipping is enabled)
|
||
grad_norm = logs.get('grad_norm', None)
|
||
|
||
# Calculate elapsed_seconds
|
||
elapsed_seconds = None
|
||
if self.trainer_instance.training_start_time is not None:
|
||
elapsed_seconds = time.time() - self.trainer_instance.training_start_time
|
||
|
||
# Calculate eta_seconds
|
||
eta_seconds = None
|
||
if elapsed_seconds is not None and current_step > 0:
|
||
total_steps = self.trainer_instance.training_progress.total_steps
|
||
if total_steps > 0:
|
||
steps_remaining = total_steps - current_step
|
||
if steps_remaining > 0:
|
||
time_per_step = elapsed_seconds / current_step
|
||
eta_seconds = time_per_step * steps_remaining
|
||
|
||
# Extract num_tokens from TRL SFTTrainer state (real counter)
|
||
# Requires include_num_input_tokens_seen=True in SFTConfig
|
||
num_tokens = getattr(state, "num_input_tokens_seen", None)
|
||
|
||
self.trainer_instance._update_progress(
|
||
step=current_step,
|
||
epoch=round(state.epoch, 2) if state.epoch else 0,
|
||
loss=loss_value,
|
||
learning_rate=logs.get('learning_rate', 0.0),
|
||
elapsed_seconds=elapsed_seconds,
|
||
eta_seconds=eta_seconds,
|
||
grad_norm=grad_norm,
|
||
num_tokens=num_tokens,
|
||
eval_loss=logs.get('eval_loss', None),
|
||
status_message=""
|
||
)
|
||
|
||
def on_epoch_end(self, args, state, control, **kwargs):
|
||
"""Called at the end of each epoch"""
|
||
self.trainer_instance._update_progress(
|
||
epoch=state.epoch,
|
||
step=state.global_step
|
||
)
|
||
|
||
def on_step_end(self, args, state, control, **kwargs):
|
||
"""Called at the end of each step"""
|
||
# Check if we should stop training
|
||
if self.trainer_instance.should_stop:
|
||
print(f"Stop detected at step {state.global_step}\n")
|
||
control.should_training_stop = True
|
||
return control
|
||
|
||
# ========== PROGRESS TRACKING ==========
|
||
progress_callback = ProgressCallback(self)
|
||
self.trainer.add_callback(progress_callback)
|
||
|
||
num_samples = len(self.trainer.train_dataset)
|
||
batch_size = training_args.get('batch_size', 2)
|
||
grad_accum = training_args.get('gradient_accumulation_steps', 4)
|
||
num_epochs = training_args.get('num_epochs', 3)
|
||
max_steps_val = training_args.get('max_steps', 0)
|
||
|
||
# Step 1: Calculate dataloader length (number of batches)
|
||
len_dataloader = math.ceil(num_samples / batch_size)
|
||
|
||
# Step 2: Calculate steps per epoch (following transformers logic)
|
||
num_update_steps_per_epoch = max(
|
||
len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0),
|
||
1
|
||
)
|
||
|
||
# Step 3: Determine total steps based on max_steps or epochs
|
||
if max_steps_val and max_steps_val > 0:
|
||
# Use max_steps if specified
|
||
total_steps = max_steps_val
|
||
print(f"Progress tracking: {total_steps} steps (max_steps)\n")
|
||
else:
|
||
# Calculate from epochs
|
||
total_steps = num_update_steps_per_epoch * num_epochs
|
||
print(f"Progress tracking: {total_steps} steps ({num_epochs} epochs × {num_update_steps_per_epoch} steps/epoch)\n")
|
||
|
||
self._update_progress(total_steps=total_steps)
|
||
|
||
# ========== START TRAINING ==========
|
||
self._update_progress(status_message="Starting training...")
|
||
print("Starting training...\n")
|
||
self.trainer.train()
|
||
|
||
# ========== SAVE MODEL ==========
|
||
if self.should_stop and self.save_on_stop:
|
||
# Stopped by user — save model at current checkpoint
|
||
self.trainer.save_model()
|
||
self.tokenizer.save_pretrained(output_dir)
|
||
print(f"\nTraining stopped. Model saved to {output_dir}\n")
|
||
self._update_progress(
|
||
is_training=False,
|
||
status_message=f"Training stopped. Model saved to {output_dir}",
|
||
)
|
||
elif self.should_stop:
|
||
# Cancelled by user — don't save
|
||
print("\nTraining cancelled.\n")
|
||
self._update_progress(
|
||
is_training=False,
|
||
status_message="Training cancelled.",
|
||
)
|
||
else:
|
||
# Normal completion
|
||
self.trainer.save_model()
|
||
self.tokenizer.save_pretrained(output_dir)
|
||
print(f"\nTraining completed! Model saved to {output_dir}\n")
|
||
self._update_progress(
|
||
is_training=False,
|
||
is_completed=True,
|
||
status_message=f"Training completed! Model saved to {output_dir}",
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Training error: {e}")
|
||
self._update_progress(is_training=False, error=str(e))
|
||
|
||
finally:
|
||
self.is_training = False
|
||
|
||
def stop_training(self, save: bool = True):
|
||
"""Stop ongoing training"""
|
||
print(f"\nStopping training (save={save})...")
|
||
self.should_stop = True
|
||
self.save_on_stop = save
|
||
stop_msg = (
|
||
"Stopping training and saving checkpoint..."
|
||
if save
|
||
else "Cancelling training..."
|
||
)
|
||
self._update_progress(status_message=stop_msg)
|
||
|
||
# If trainer exists, try to stop it gracefully
|
||
if self.trainer:
|
||
try:
|
||
# The callback will catch should_stop flag and stop the training loop
|
||
print("Training will stop at next step...\n")
|
||
except Exception as e:
|
||
logger.error(f"Error stopping trainer: {e}")
|
||
|
||
def get_training_progress(self) -> TrainingProgress:
|
||
"""Get current training progress"""
|
||
with self._lock:
|
||
return self.training_progress
|
||
|
||
def cleanup(self):
|
||
"""Cleanup resources"""
|
||
if self.trainer:
|
||
self.trainer = None
|
||
if self.model:
|
||
self.model = None
|
||
if self.tokenizer:
|
||
self.tokenizer = None
|
||
|
||
# Clear GPU memory
|
||
clear_gpu_cache()
|
||
|
||
|
||
def _ensure_deepseek_ocr_installed():
|
||
"""
|
||
Auto-install DeepSeek OCR module if not available.
|
||
Downloads from HuggingFace hub as a local module.
|
||
|
||
Returns:
|
||
bool: True if available (either already installed or just installed)
|
||
"""
|
||
try:
|
||
# Try importing to see if already available
|
||
from deepseek_ocr.modeling_deepseekocr import format_messages
|
||
logger.info("DeepSeek OCR module already available")
|
||
return True
|
||
except ImportError:
|
||
pass
|
||
|
||
try:
|
||
logger.info("DeepSeek OCR module not found. Auto-installing from HuggingFace...")
|
||
print("\n Downloading DeepSeek OCR module from HuggingFace...\n")
|
||
|
||
from huggingface_hub import snapshot_download
|
||
import sys
|
||
import os
|
||
|
||
# Get the script directory to install locally
|
||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||
parent_dir = os.path.dirname(script_dir) # Go up to project root
|
||
|
||
# Download to project root as 'deepseek_ocr' folder
|
||
local_dir = os.path.join(parent_dir, "deepseek_ocr")
|
||
|
||
snapshot_download(
|
||
"unsloth/DeepSeek-OCR",
|
||
local_dir=local_dir,
|
||
local_dir_use_symlinks=False
|
||
)
|
||
|
||
# Add to sys.path if not already there
|
||
if parent_dir not in sys.path:
|
||
sys.path.insert(0, parent_dir)
|
||
|
||
# Try importing again
|
||
from deepseek_ocr.modeling_deepseekocr import format_messages
|
||
|
||
logger.info("DeepSeek OCR module installed successfully")
|
||
print("DeepSeek OCR module installed successfully!\n")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to install DeepSeek OCR module: {e}")
|
||
print(f"\n❌ Failed to install DeepSeek OCR module: {e}\n")
|
||
return False
|
||
|
||
# Global trainer instance
|
||
_trainer_instance = None
|
||
|
||
def get_trainer() -> UnslothTrainer:
|
||
"""Get global trainer instance"""
|
||
global _trainer_instance
|
||
if _trainer_instance is None:
|
||
_trainer_instance = UnslothTrainer()
|
||
return _trainer_instance
|