Building an ML Model Training UI: SFT, DPO, and LoRA Config

Why a Training Page?

The XGen platform supports not only workflow execution but also direct model fine-tuning. We built a training page so users can configure and launch SFT (Supervised Fine-Tuning), DPO (Direct Preference Optimization), and LoRA runs entirely from the web UI.

Training Configuration Form Structure

The training configuration is fairly complex, spanning multiple sections: base settings, hyperparameters, dataset selection, and GPU settings.

interface TrainingConfig {
  // Base settings
  baseModel: string;          // 'gemma3-4b' | 'qwen3-8b' | etc.
  trainingMethod: 'sft' | 'dpo' | 'lora';
  
  // Hyperparameters
  learningRate: number;
  epochs: number;
  batchSize: number;
  warmupSteps: number;
  weightDecay: number;
  
  // LoRA settings (when trainingMethod === 'lora')
  loraRank?: number;
  loraAlpha?: number;
  loraDropout?: number;
  targetModules?: string[];
  
  // DPO settings
  dpoBeta?: number;
  
  // GPU settings
  gpuCount: number;
  gpuType: string;
}

Conditional Form Rendering

The fields rendered depend on the selected training method. We combined React Hook Form with Zod to implement dynamic validation.

const trainingSchema = z.discriminatedUnion('trainingMethod', [
  z.object({
    trainingMethod: z.literal('sft'),
    learningRate: z.number().min(1e-6).max(1e-2),
    epochs: z.number().int().min(1).max(100),
    batchSize: z.number().int().min(1).max(128),
  }),
  z.object({
    trainingMethod: z.literal('lora'),
    learningRate: z.number().min(1e-6).max(1e-2),
    loraRank: z.number().int().min(4).max(128),
    loraAlpha: z.number().int().min(8).max(256),
    loraDropout: z.number().min(0).max(0.5),
  }),
  z.object({
    trainingMethod: z.literal('dpo'),
    learningRate: z.number().min(1e-6).max(1e-2),
    dpoBeta: z.number().min(0.01).max(1),
  }),
]);

Real-Time Training Monitoring

Once training starts, the UI displays a live loss graph, the current epoch, and estimated time remaining. This data is streamed over SSE, and we use Recharts to render a line chart that updates in real time.

Training History Management

Past training runs are shown in a list, and a comparison view lets users examine the configuration and performance metrics of each run side by side. Overlaying the loss curves from multiple runs on a single chart makes it immediately clear which configuration performed better.

Preset System

For less experienced users, we provide recommended hyperparameter presets for each model. Selecting a preset such as "Gemma3 SFT Default" or "Qwen3 LoRA Fast" automatically populates all the relevant fields.