Optimizing LLM Model Training Time: Tips and Techniques

#engineering #oxlo #ai

Training large models is expensive when the config is wrong. In this guide, I build a lightweight advisor agent that audits your training setup and suggests concrete optimizations before you burn GPU hours. We will run the reasoning engine on Oxlo.ai so you get fast, structured advice without managing infrastructure.

What you'll need

An Oxlo.ai API key
Python 3.10 or newer
The OpenAI SDK: pip install openai

Step 1: Bootstrap the client

First, I import the OpenAI SDK and point it at Oxlo.ai. I pull the API key from the environment so it never hits disk in the script.

import os
from openai import OpenAI

client = OpenAI(
    base_url="https://api.oxlo.ai/v1",
    api_key=os.environ.get("OXLO_API_KEY")
)

# Sanity check
assert client.api_key is not None, "Set the OXLO_API_KEY environment variable."

Step 2: Define the config schema

I use a typed dictionary to enforce the shape of the training configuration. This keeps the prompt construction deterministic and prevents bad inputs from reaching the model.

from typing import TypedDict

class TrainingConfig(TypedDict):
    model_name: str
    parameter_count_b: float
    gpu_type: str
    gpu_count: int
    dataset_tokens_b: float
    batch_size_per_device: int
    precision: str  # fp32, fp16, bf16
    gradient_accumulation_steps: int
    sequence_length: int
    current_estimated_hours: float

def get_sample_config() -> TrainingConfig:
    return {
        "model_name": "Llama-3.1-8B",
        "parameter_count_b": 8.0,
        "gpu_type": "NVIDIA A100 80GB",
        "gpu_count": 4,
        "dataset_tokens_b": 15.0,
        "batch_size_per_device": 2,
        "precision": "bf16",
        "gradient_accumulation_steps": 8,
        "sequence_length": 4096,
        "current_estimated_hours": 120.0,
    }

Step 3: Write the system prompt

The system prompt tells the model to act as a performance engineer. I ask for concrete changes, estimated time savings, and caveats so the output is actionable.

SYSTEM_PROMPT = """You are an HPC performance engineer specializing in LLM training optimization. 
Your job is to analyze the user's training configuration and suggest concrete changes that reduce wall-clock time without sacrificing convergence.

Rules:
- Recommend only changes that fit within the stated hardware.
- Prefer memory-efficient attention, distributed data parallelism adjustments, and precision tuning.
- Estimate the percentage of time saved for each suggestion.
- If gradient checkpointing or micro-batching is needed, say so explicitly.
- Keep the response structured with bullet points and a final summary.
- Do not suggest buying more hardware unless the config is physically impossible."""

Step 4: Build the advisor function

I format the config into a markdown table and send it to Oxlo.ai. I use deepseek-v3.2 here because it handles coding and reasoning well, and it sits on the Oxlo.ai free tier. Because Oxlo.ai charges per request instead of per token, I can stuff the prompt with the full hardware context and a detailed system message without the cost scaling with input length.

def build_user_message(config: TrainingConfig) -> str:
    return f"""Here is my current training configuration:

| Setting | Value |
|---|---|
| Model | {config['model_name']} |
| Parameters | {config['parameter_count_b']}B |
| GPUs | {config['gpu_count']}x {config['gpu_type']} |
| Dataset | {config['dataset_tokens_b']}B tokens |
| Global batch | {config['batch_size_per_device'] * config['gpu_count'] * config['gradient_accumulation_steps']} |
| Per-device batch | {config['batch_size_per_device']} |
| Grad accum steps | {config['gradient_accumulation_steps']} |
| Precision | {config['precision']} |
| Sequence length | {config['sequence_length']} |
| Current estimate | {config['current_estimated_hours']} hours |

Suggest optimizations to cut training time. Include numerical estimates where possible."""

def run_advisor(config: TrainingConfig) -> str:
    response = client.chat.completions.create(
        model="deepseek-v3.2",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": build_user_message(config)},
        ],
        temperature=0.2,
        max_tokens=1200,
    )
    return response.choices[0].message.content

Step 5: Wire up the CLI

I add a small entrypoint so I can run this from the terminal against the sample config. In production, you would swap get_sample_config() for a call to your experiment tracking API.

if __name__ == "__main__":
    config = get_sample_config()
    print("Analyzing training configuration...")
    print("-" * 40)
    advice = run_advisor(config)
    print(advice)

Run it

Export your key and execute the script:

export OXLO_API_KEY="sk-oxlo.ai-..."
python train_optimizer.py

Example output:

Analyzing training configuration...
----------------------------------------
* Increase per-device batch size from 2 to 4. With 80 GB A100s and bf16, you have headroom. Estimated savings: 15%.
* Reduce gradient accumulation from 8 to 4 to match the larger micro-batch. This lowers step overhead. Estimated savings: 8%.
* Enable torch.compile with max-autotune mode. On A100 this typically yields 10-20% throughput gain. Estimated savings: 12%.
* Use flash_attention_2 if not already enabled. For 4096 sequence length this removes the quadratic bottleneck. Estimated savings: 18%.
* Summary: Combined wall-clock reduction is roughly 40-45 hours, bringing the run down to ~75 hours. Monitor gradient norms after increasing batch size to ensure stable loss.

Wrap up

This advisor gives you a sanity check in seconds before you commit days of GPU time. Two concrete next steps: wire the script into your experiment launcher so it previews every run, and feed actual profiling logs from PyTorch profiler back into the prompt as context for iterative tuning.

DEV Community