rayllm-orchestrator 0.1.2

Train and serve LLMs on any hardware with one command each

RayLLM-Orchestrator

Fine-tune an LLM on your data and serve it behind an OpenAI-compatible API — one command each.

# Train
rayllm train --model gpt2 --dataset my-data.jsonl --epochs 3

# Serve (in another terminal)
rayllm serve --model ./checkpoints/gpt2 --port 8000

# Test it
curl -X POST http://localhost:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"messages":[{"role":"user","content":"Hello"}]}'

Both rayllm <cmd> (installed console script) and python orchestrator.py <cmd> (from a source checkout) work identically.

---

What It Does

Train — real PyTorch fine-tuning with:

• Model loading from HuggingFace IDs or local paths
• Dataset tokenization (JSONL or HuggingFace datasets)
• LoRA / QLoRA (4-bit / 8-bit base weights)
• Checkpointing and resume from a saved epoch
• Optional multi-GPU via Ray + PyTorch FSDP (needs >1 CUDA GPU)
• On Apple Silicon, an MLX-LoRA path when mlx-lm is installed and the model is MLX-loadable; otherwise it falls back to the PyTorch path

Serve — one OpenAI-compatible HTTP API over several backends. You pick the backend (a recommended default is offered); it is not chosen silently. See Choosing a backend.

---

Installation

pip install rayllm-orchestrator

This installs the core stack: torch, transformers, datasets, peft, accelerate. The serving backends below are optional and installed separately — you only need the one(s) you want:

Backend	Extra install	Where it runs
transformers	(included)	CPU / CUDA / Apple MPS — works everywhere
MLX	pip install mlx mlx-lm	Apple Silicon (arm64 macOS)
llama.cpp	pip install llama-cpp-python	CPU, Metal on macOS — GGUF models
vLLM	pip install vllm	NVIDIA GPU (Linux)
Ollama	install Ollama + pip install ollama	uses the local Ollama daemon

From source

git clone https://github.com/agastya-choudhary123/rayllm-orchestrator
cd rayllm-orchestrator
pip install -e .
python orchestrator.py train --model gpt2 --dataset my-data.jsonl

Docker

The repo ships a Dockerfile (and Dockerfile.cpu) you build locally — there is no prebuilt image published to a registry:

docker build -t rayllm .
docker run -it rayllm train --model gpt2 --dataset my-data.jsonl --epochs 3

---

Quick Start

1. Prepare your data

JSONL, one object per line. Supported shapes: {"text": ...}, {"prompt": ..., "completion": ...}, or {"messages": [...]}.

{"text": "Your training example here."}
{"text": "Another example."}

Or point at a HuggingFace dataset id:

rayllm train --model gpt2 --dataset wikitext --epochs 1

2. Train

rayllm train --model gpt2 --dataset my-data.jsonl --epochs 3 --quant 4bit

Downloads the model, tokenizes your data, trains (saving a checkpoint after each epoch), and prints the checkpoint path. With --quant, base weights load in 4/8-bit and only LoRA adapters train (QLoRA); the PyTorch path merges the adapters into a plain HuggingFace checkpoint the serving layer loads directly.

3. Serve

rayllm serve --model ./checkpoints/gpt2 --port 8000

Prompts you to choose a backend from what's installed (or pass --backend), then starts an OpenAI-compatible API on port 8000.

4. Use it

curl -X POST http://localhost:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"messages":[{"role":"user","content":"Write a haiku about coding"}],"max_tokens":50}'

---

Choosing a backend

List what's actually usable on your machine:

rayllm backends

Example output — availability is probed live (installed and runnable here), and a format-aware recommendation is marked:

1. [✗] ollama         needs the ollama daemon running
2. [✓] mlx            Apple-Silicon-native            <- recommended for this model
3. [✗] llama_cpp      `pip install llama-cpp-python`
4. [✗] vllm           installed but no CUDA GPU found
5. [✓] transformers   universal fallback

Then either let serve/run prompt you (it offers the recommended one as the Enter-default), or force a choice:

rayllm serve --model ./checkpoints/gpt2 --backend transformers

---

Common Use Cases

Apple Silicon (MLX)

pip install mlx mlx-lm
rayllm serve --model mlx-community/Qwen2.5-0.5B-Instruct-4bit --backend mlx

MLX is Apple's Metal-native framework and keeps 4-bit models memory-safe in unified memory (a guard caps its cache so training can't starve the OS). See measured numbers.

Fine-tune with quantization (QLoRA)

rayllm train --model meta-llama/Llama-3.2-3B --dataset my-data.jsonl --quant 4bit

Loads base weights in 4-bit and trains only the LoRA adapters — far fewer trainable parameters and lower memory than full fine-tuning.

Multi-GPU training (Ray + FSDP)

rayllm train --model meta-llama/Llama-3.2-3B --dataset my-data.jsonl \
  --strategy fsdp-ray --num-workers 4

Requires Ray installed and more than one CUDA GPU; otherwise it logs a note and runs the single-device loop. (deepspeed-ray is accepted but currently runs the same Ray+FSDP path — there is no separate DeepSpeed backend.)

Speculative decoding (MLX)

rayllm serve --model <large-mlx-model> --draft-model <small-mlx-model> --backend mlx

A small draft model proposes tokens the large model verifies — faster generation with identical output. MLX backend only.

---

Command Reference

train

rayllm train --model <id-or-path> --dataset <id-or-path.jsonl>
  [--epochs N] [--strategy single|fsdp-ray|deepspeed-ray]
  [--quant none|8bit|4bit] [--out DIR] [--num-workers N]

Defaults: --epochs 1, --strategy fsdp-ray (falls back to single without Ray + multiple GPUs), --quant none, --out ./checkpoints, --num-workers 0 (auto).

serve

rayllm serve --model <checkpoint-or-id>
  [--backend ollama|mlx|llama_cpp|vllm|transformers]
  [--quant ...] [--port 8000] [--host 0.0.0.0]
  [--max-model-len 4096] [--draft-model <id>]

If --backend is omitted you're prompted to choose (interactive terminals only).

run

rayllm run --model <id-or-path> --data <id-or-path.jsonl>
  [--epochs N] [--port 8000] [--backend ...] [--no-serve] [--no-fast]

Fine-tune, then serve in one shot.

backends

rayllm backends

Scan this machine and list which serving backends are available.

---

Hardware Support

Hardware	Train	Serve	Notes
CPU (any OS)	✓ (torch)	✓ (transformers / llama.cpp)	Works everywhere; slow on large models
Apple Silicon	✓ (MLX or torch)	✓ (MLX / transformers)	MLX needs mlx-lm; native Metal GPU
NVIDIA GPU (Linux)	✓ (FSDP w/ Ray)	✓ (vLLM / transformers)	vLLM/Ray installed separately
Multi-GPU	✓ (Ray + FSDP)	✓ (vLLM)	Needs >1 GPU

The vLLM and multi-GPU/FSDP paths require real NVIDIA hardware and are not exercised in this project's own testing — treat them as supported-but-unverified here.

---

Measured Performance

Real serving throughput measured with the repo's benchmark_serving helper (orchestrator/backends_mlx.py) on one Apple Silicon Mac (~17 GB unified memory), MLX backend, 4-bit weights, 120-token generations:

Model	Throughput	Peak memory
Qwen2.5-0.5B-Instruct-4bit	~276 tok/s	0.33 GB
Qwen3-0.6B-4bit	~240 tok/s	0.44 GB
Qwen3-8B-4bit	~22 tok/s	4.72 GB

These are single-machine snapshots, not guarantees — throughput and memory depend heavily on model, quantization, prompt length, and hardware. Benchmark on your own setup before relying on any number. No cross-backend speed comparisons (e.g. MLX vs. PyTorch MPS) are claimed here because none have been measured.

---

Troubleshooting

No module named 'torch' — install the core stack: pip install rayllm-orchestrator (or pip install torch transformers datasets peft accelerate).

Out of memory during training — use quantization: add --quant 4bit.

"No serving backend is available" — run rayllm backends to see what's installed, then install one (e.g. pip install mlx-lm on a Mac) or pass --backend transformers.

Serve doesn't respond — check it's up: curl http://localhost:8000/v1/models.

---

Architecture

orchestrator.py            ← thin shim (python orchestrator.py ...)
orchestrator/
├── cli.py                ← the `rayllm` command (train / serve / run / backends)
├── train.py              ← training loop + MLX/torch dispatch, fallback
├── serve.py              ← serving + backend dispatch (OpenAI-compatible API)
├── models.py             ← model loading + backend detection/selection
├── data.py               ← tokenization, packing/bucketing
├── fast.py               ← bf16, prefetch, packing, LoRA optimizations
├── backends_mlx.py       ← Apple Silicon native train + serve (MLX)
└── util.py               ← logging, capability detection

Core dependencies: torch, transformers, datasets, peft, accelerate. Backends (MLX, vLLM, llama.cpp, Ollama) and Ray are optional and installed as needed.

---

License

MIT. Contributions welcome — see CONTRIBUTING.md.

See Also

• vLLM — GPU serving with continuous batching
• MLX — Apple Silicon ML framework
• Ray Train — distributed training
• Hugging Face — models & datasets