Orion-ML

Navigate large model training and inference with confidence.

Diagnose bottlenecks. Get ranked optimization strategies. Understand cost-accuracy trade-offs. Works with PyTorch-centric stacks and common trainers. No infrastructure changes required.

Just as navigators used Orion to find their position and chart their course across open ocean, Orion-ML tells you where training or serving is bottlenecked and charts ranked strategies to improve throughput, memory, and cost.

• PyPI: orion-ml
• Tagline: Navigate large model training and inference with confidence
• Users: ML engineers and platform teams training LLMs and VLMs on multi-GPU and multi-node clusters

Components

Install

pip install orion-ml
# Optional: PyTorch, GPU telemetry, experiment trackers, HF Trainer, Lightning
pip install "orion-ml[torch,gpu,trackers,hf,lightning]"

Quick start: three lines in your training loop

from orion.train import TrainingTracker

tracker = TrainingTracker(job_name="my-run", output_dir="./orion_output")
with tracker.step():
    loss = model(inputs)
    loss.backward()
    optimizer.step()

Metrics are written to local SQLite first; optional W&B / MLflow export runs asynchronously and never blocks the step.

Hugging Face Trainer

from transformers import Trainer
from orion.train.callbacks import OrionTrainCallback

trainer = Trainer(..., callbacks=[OrionTrainCallback(job_name="hf-run", output_dir="./orion_output")])

Parallelism recommendation

from orion.train import ParallelismAdvisor, ParallelismConfig

advisor = ParallelismAdvisor()
rec = advisor.recommend(
    ParallelismConfig(
        model_params=7_000_000_000,
        gpu_memory_gb=80,
        num_gpus=8,
        num_nodes=1,
        interconnect="nvlink",
        precision="bf16",
        target_batch_size=64,
    )
)
print(rec.strategy, rec.reasoning)

Inference profiling

from orion.serve import InferenceProfiler

profiler = InferenceProfiler(model=model, tokenizer=tokenizer)
results = profiler.benchmark(prompts=test_prompts, max_new_tokens=256, batch_sizes=[1, 8, 32])
profiler.report()

Reports

from orion.sight import Reporter

reporter = Reporter(tracker)
reporter.report()  # terminal (rich)
reporter.report(format="html", output_path="report.html")
reporter.report(format="json", output_path="report.json")

Repository layout

orion/
├── orion/           # Python package
├── docs/            # Guides and API reference
├── examples/        # Runnable demos (may require torch/GPU)
└── tests/           # unit / integration / benchmarks

Output: sheetal@Sheetal-MBP project-orion % python gate3_real_bottleneck.py

================================================== GATE 3: REAL BOTTLENECK DETECTION

[1/2] Creating real data loader bottleneck... [transformers] Disabling PyTorch because PyTorch >= 2.4 is required but found 2.2.2 [transformers] PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used. [Orion-ML] Warning: PyTorch < 2.4 detected; some packages recommend torch>=2.4. Core TrainingTracker metrics are unchanged. ╭────── Orion-ML Training summary ──────╮ │ Primary bottleneck: DATA_LOADER_BOUND │ │ Efficiency score: 25/100 │ │ Anomalies flagged: 0 │ ╰───────────────────────────────────────╯ Top recommendations
┏━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Rank ┃ Strategy ┃ Accuracy risk ┃ ┡━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ 1 │ increase_num_workers │ NONE │ │ 2 │ pin_memory │ NONE │ │ 3 │ prefetch_factor │ NONE │ │ 4 │ streaming_dataset_or_dali │ LOW │ └──────┴───────────────────────────┴───────────────┘ Data loader overhead: 92.0% ✅ Real data loader bottleneck correctly detected

[2/2] Creating real gradient explosion... [Orion-ML] Warning: PyTorch < 2.4 detected; some packages recommend torch>=2.4. Core TrainingTracker metrics are unchanged. ╭────── Orion-ML Training summary ──────╮ │ Primary bottleneck: DATA_LOADER_BOUND │ │ Efficiency score: 25/100 │ │ Anomalies flagged: 1 │ ╰───────────────────────────────────────╯ Top recommendations
┏━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Rank ┃ Strategy ┃ Accuracy risk ┃ ┡━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ 1 │ increase_num_workers │ NONE │ │ 2 │ pin_memory │ NONE │ │ 3 │ prefetch_factor │ NONE │ │ 4 │ streaming_dataset_or_dali │ LOW │ └──────┴───────────────────────────┴───────────────┘ Anomalies
┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Type ┃ Severity ┃ Message ┃ ┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ LOSS_DIVERGENCE │ CRITICAL │ Loss increased for 3+ consecutive steps — check │ │ │ │ LR, batch, or bad data. │ └─────────────────┴──────────┴─────────────────────────────────────────────────┘ Anomalies detected: ['LOSS_DIVERGENCE'] ✅ Loss divergence detected (same root cause — acceptable)

================================================== GATE 3: ✅ ALL PASSED — Ready for Lambda GPU demo

License

Apache 2.0 — see LICENSE.

Contributing

See CONTRIBUTING.md.