POLAR Trainer: Building a General-Purpose Multi-Training Framework
Background
As the demand for training diverse AI models in the e-commerce domain grew, writing new training code from scratch every time became increasingly inefficient. We built POLAR Trainer to handle all training types — classification, language model fine-tuning, DPO alignment, sentence embedding training, and more — within a single framework.
Supported Training Types
| Type | Description |
|---|---|
| Classification | Training text/image classification models |
| MLM (Masked Language Modeling) | Pre-training BERT-family models |
| CLM (Causal Language Modeling) | Training GPT-family language models |
| DPO (Direct Preference Optimization) | Aligning models to human preferences as an RLHF alternative |
| SFT (Supervised Fine-Tuning) | Fine-tuning with supervised learning |
| Multimodal | Text + image multimodal training |
| Sentence Transformer | Training sentence embedding models |
Core Technology Stack
DeepSpeed ZeRO Integration
We integrated DeepSpeed ZeRO to enable large-model training on a single GPU. The framework supports ZeRO Stage 2/3, distributing optimizer states and gradients across data-parallel groups to reduce memory usage.
# DeepSpeed 설정 예시
deepspeed_config = {
"zero_optimization": {
"stage": 2,
"offload_optimizer": {"device": "cpu"},
"allgather_partitions": True
},
"fp16": {"enabled": True},
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 8
}
PEFT/LoRA Support
Rather than fine-tuning the entire model, we use LoRA (Low-Rank Adaptation) to train only a small subset of parameters, improving efficiency. QLoRA (quantization + LoRA) is also supported.
from peft import LoraConfig, get_peft_model
lora_config = LoraConfig(
r=16, # Rank
lora_alpha=32, # Alpha
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
task_type="CAUSAL_LM"
)
model = get_peft_model(base_model, lora_config)
MLflow Experiment Tracking
Hyperparameters, metrics, and model artifacts for every training run are tracked with MLflow. This makes it straightforward to compare experiments and identify the best configuration.
MinIO Model Registry
Once training completes, models are automatically uploaded to MinIO object storage. HuggingFace Hub integration is also supported, allowing trained models to be deployed or shared immediately.
Integration with the XGen Platform
POLAR Trainer is integrated with the Training page of the XGen platform. When a user selects a model and training configuration through the UI, POLAR Trainer runs on the backend to execute the job. Vast.ai GPU instances are provisioned dynamically, so we get access to powerful GPUs on demand while keeping costs under control.
Retrospective
The framework skeleton was built in 7 key commits, and the Training UI on the XGen frontend was developed over 30+ commits. The biggest challenge was providing a unified interface across diverse training types without losing the type-specific configuration details each one requires. We addressed this by building on top of the HuggingFace Transformers Trainer API and layering in custom logic for each training type.