Training large models wastes hours when batch sizes, precision, or parallelism are misconfigured. In this tutorial, I will build a Python agent that ingests a training configuration and returns concrete changes to cut wall-clock time. We will run the agent through Oxlo.ai, where request-based pricing keeps costs flat even when we feed the model long hardware logs and detailed JSON configs.
What you'll need
The only dependencies are Python 3.10+ and the OpenAI SDK. You will also need an Oxlo.ai API key.
pip install openaiGrab your key from https://portal.oxlo.ai.
Step 1: Craft the system prompt
I want structured output so I can parse it programmatically. The system prompt forces the model to return JSON with four fields: bottleneck, changes, estimated_speedup, and reasoning.
SYSTEM_PROMPT = """You are a training optimization engineer. Analyze the user's distributed training configuration and identify the single biggest bottleneck. Respond ONLY in JSON with this structure:
{
"bottleneck": "string describing the primary bottleneck",
"changes": [
{"parameter": "string", "current": "value", "recommended": "value", "rationale": "string"}
],
"estimated_speedup": "string like 1.3x or 2.1x",
"reasoning": "string with your chain of thought"
}
Be specific. Reference hardware limits, memory bandwidth, and numerical precision. Do not suggest changes that exceed the stated GPU memory."""Step 2: Bootstrap the Oxlo.ai client
The OpenAI SDK works as a drop-in replacement. I point the base URL at Oxlo.ai and load the key from the environment.
import os
from openai import OpenAI
client = OpenAI(
base_url="https://api.oxlo.ai/v1",
api_key=os.getenv("OXLO_API_KEY")
)Step 3: Model the training config
I will use a typed dictionary so the rest of the script stays readable. This schema covers the variables that matter most for throughput.
from typing import TypedDict
class TrainingConfig(TypedDict):
model_name: str
parameters_b: float
dataset_tokens: int
gpu_type: str
num_gpus: int
batch_size_per_device: int
gradient_accumulation_steps: int
precision: str # fp32, fp16, bf16
sequence_length: int
optimizer_state_bytes_per_param: float # e.g., 12 for AdamW
def format_config(config: TrainingConfig) -> str:
lines = [f"{k}: {v}" for k, v in config.items()]
return "\n".join(lines)Step 4: Build the analysis function
This is where we call Oxlo.ai. I use the Qwen 3 32B model because optimization requires reasoning about memory bandwidth and compute tradeoffs. Because Oxlo.ai charges per request rather than per token, I can pass the full config as a large prompt without watching metered tokens tick up.
import json
def optimize_training(config: TrainingConfig) -> dict:
user_message = (
"Analyze the following training configuration and suggest optimizations.\n\n"
+ format_config(config)
)
response = client.chat.completions.create(
model="qwen-3-32b",
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_message},
],
)
content = response.choices[0].message.content
cleaned = content.strip()
if cleaned.startswith("
```json"):
cleaned = cleaned[7:]
elif cleaned.startswith("```
"):
cleaned = cleaned[3:]
if cleaned.endswith("
```
"):
cleaned = cleaned[:-3]
return json.loads(cleaned.strip())Step 5: Add a validation layer
LLMs can hallucinate keys. I will wrap the parser in a small validator that prints a human-readable report and falls back gracefully if JSON is malformed.
def print_report(result: dict):
print(f"Bottleneck: {result.get('bottleneck', 'Unknown')}")
print(f"Estimated speedup: {result.get('estimated_speedup', 'N/A')}")
print("\nRecommended changes:")
for change in result.get("changes", []):
print(
f" - {change['parameter']}: {change['current']} -> {change['recommended']}"
f" ({change['rationale']})"
)
print(f"\nReasoning: {result.get('reasoning', '')}")
def safe_optimize(config: TrainingConfig):
try:
result = optimize_training(config)
print_report(result)
except Exception as e:
print(f"Failed to parse optimization result: {e}")Step 6: Run it
Here is a realistic misconfiguration: an 8B parameter model on four A100s with a tiny per-device batch size and fp32 precision. This is exactly the kind of long-context engineering prompt that benefits from Oxlo.ai's flat per-request pricing.
if __name__ == "__main__":
config: TrainingConfig = {
"model_name": "meta-llama/Meta-Llama-3.1-8B",
"parameters_b": 8.0,
"dataset_tokens": 1_000_000_000,
"gpu_type": "NVIDIA A100-SXM4-80GB",
"num_gpus": 4,
"batch_size_per_device": 1,
"gradient_accumulation_steps": 64,
"precision": "fp32",
"sequence_length": 4096,
"optimizer_state_bytes_per_param": 12.0,
}
safe_optimize(config)Example output:
Bottleneck: Memory-bound training caused by fp32 precision and unnecessarily small micro-batch
Estimated speedup: 2.8x
Recommended changes:
- precision: fp32 -> bf16 (Halves activation and weight memory, doubles effective tensor-core throughput on A100)
- batch_size_per_device: 1 -> 4 (Increases occupancy without exceeding 80GB when combined with bf16 and gradient checkpointing)
- gradient_accumulation_steps: 64 -> 16 (Keeps global batch constant while reducing pipeline bubble overhead)
- activation_checkpointing: off -> on (Trades ~20% compute for 2x memory savings, enabling larger micro-batches)
Reasoning: At fp32 the model weights alone consume ~32GB, leaving insufficient headroom for activations with batch_size > 1. Switching to bf16 reduces weights to ~16GB and unlocks bfloat16 tensor-core paths on Ampere. A micro-batch of 4 fits comfortably with activation checkpointing, moving the bottleneck from memory bandwidth to compute and improving SM utilization from 34% to 81%.Next steps
Wire this agent into your experiment tracker so it can ingest live throughput metrics from Weights & Biases and suggest dynamic adjustments during a run. If you need faster latency for interactive use, swap the model ID to llama-3.3-70b or deepseek-v3.2 on Oxlo.ai, both available on the same flat per-request plan.
Top comments (0)