Building praque_trainer: MLflow-Integrated Model Training Module

Project Background

praque_trainer is the model training backend module for the XGen platform. It was developed in August 2025 across 7 commits.

Architecture

graph TD
    A[XGen Training UI] -->|API 호출| B[praque_trainer]
    B --> C[데이터 로더]
    B --> D[PEFT Config]
    B --> E[StableAdamW]
    B --> F[MLflow]

    C --> G[학습 실행]
    D --> G
    E --> G
    G --> F

    F --> H[메트릭 기록]
    F --> I[모델 아티팩트 저장]

Core Components

1. StableAdamW Optimizer

Drawing on lessons from prj_ecellm, StableAdamW was adopted as the default optimizer:

from lib.optimizers import StableAdamW

optimizer = StableAdamW(
    model.parameters(),
    lr=2e-5,
    weight_decay=0.001,  # value learned from ecellm
    eps=1e-8
)

2. PEFT Configuration Loader

Dynamically loads configurations for LoRA, QLoRA, and similar methods:

def load_peft_config(config_path: str):
    config = load_yaml(config_path)
    return PeftConfig(
        r=config.get("lora_r", 16),
        alpha=config.get("lora_alpha", 32),
        target_modules=config.get("target_modules", ["q_proj", "v_proj"]),
        dropout=config.get("dropout", 0.05)
    )

3. MLflow Integration

Reads the MLflow URL from an environment variable and tracks experiments:

import mlflow

mlflow_url = os.environ.get("MLFLOW_URL")
mlflow.set_tracking_uri(mlflow_url)

with mlflow.start_run():
    mlflow.log_params(training_config)
    for epoch in range(num_epochs):
        metrics = train_one_epoch(model, data)
        mlflow.log_metrics(metrics, step=epoch)
    mlflow.log_artifact(model_path)

Refactoring Notes

The final commit included a significant refactor:

• Constants and loaders were reorganized into separate modules
• Unused files were removed
• Environment variable references were replaced with constants

All 7 commits landed on the same day (8/4) — a focused, single-day sprint.