neotune

Simple LLM fine-tuning with LoRA and DeepSpeed.

Three inputs: a model, your datasets, and hyperparameters.

Installation

pip install neotune

With optional extras:

pip install "neotune[ray]"       # Ray distributed training
pip install "neotune[logging]"   # MLflow + Weights & Biases
pip install "neotune[all]"       # everything

Quick Start

from neotune import finetune

results = finetune(
    model="google/gemma-3-4b-pt",
    datasets={"train": train_ds, "validation": val_ds},
    hyperparameters={"learning_rate": 2e-4, "num_train_epochs": 3},
)

Each dataset is a HuggingFace Dataset with a "text" column containing fully-formatted prompt/response strings, or pre-tokenized columns (input_ids, attention_mask, labels).

Preparing Your Datasets

neotune expects you to bring your own HuggingFace Dataset objects.

From a HuggingFace dataset

from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-4b-pt")

ds = load_dataset("tatsu-lab/alpaca", split="train")

def format_example(example):
    messages = [
        {"role": "user", "content": example["instruction"]},
        {"role": "assistant", "content": example["output"]},
    ]
    return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}

ds = ds.map(format_example, remove_columns=ds.column_names)
splits = ds.train_test_split(test_size=0.1, seed=42)

# splits["train"] and splits["test"] each have a "text" column

From a CSV file

from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-4b-pt")

ds = load_dataset("csv", data_files="data.csv", split="train")

def format_example(example):
    messages = [
        {"role": "user", "content": example["prompt"]},
        {"role": "assistant", "content": example["response"]},
    ]
    return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}

ds = ds.map(format_example, remove_columns=ds.column_names)
splits = ds.train_test_split(test_size=0.1, seed=42)

API Reference

finetune(model, datasets, hyperparameters) -> dict

One-call convenience function. Returns test-set metrics if a "test" split was provided, otherwise an empty dict.

from neotune import finetune

results = finetune(
    model="google/gemma-3-4b-pt",
    datasets={"train": train_ds, "validation": val_ds, "test": test_ds},
    hyperparameters={"learning_rate": 2e-4},
)

NeoTune(model, datasets, hyperparameters)

Class-based API.

from neotune import NeoTune

nt = NeoTune(
    model="google/gemma-3-4b-pt",
    datasets={"train": train_ds, "validation": val_ds},
    hyperparameters={"num_train_epochs": 5, "output_dir": "./my-adapter"},
)

results = nt.train()

Parameters

model -- str HuggingFace model ID or local path.

datasets -- dict[str, Dataset] A dict of HuggingFace Dataset objects. "train" is required. "validation" and "test" are optional.

hyperparameters -- dict, optional Override any default. All keys are optional:

Methods

• .train() -> dict -- Fine-tune and return test metrics (if test split provided).
• .tokenizer -- Access the underlying tokenizer.

Compute Environments

neotune auto-detects your hardware and launches training accordingly — same finetune() call everywhere:

Multi-GPU works from scripts, notebooks, and the Python REPL — no need for torchrun or deepspeed CLI. neotune handles the distributed launch automatically.

Override with ds_config:

# Force DeepSpeed off
finetune(model, datasets, {"ds_config": False})

# Force DeepSpeed on (even on single GPU)
finetune(model, datasets, {"ds_config": "auto"})

# Custom DeepSpeed config
finetune(model, datasets, {"ds_config": "my_ds_config.json"})

Advanced Usage

Generative evaluation

from neotune.eval import generate_and_evaluate

results = generate_and_evaluate(
    model_id="google/gemma-3-4b-pt",
    adapter_dir="./my-adapter",
    test_ds=test_ds,
    prompt_col="instruction",
    label_col="expected_output",
)

Distributed training with DeepSpeed (CLI)

deepspeed --num_gpus 4 -m neotune.train --config config.yaml --mode train

Distributed training with Ray

python -m neotune.ray_train --config config.yaml --num_workers 4

Kubernetes (KubeRay)

See k8s/rayjob-lora-sft.yaml for a KubeRay RayJob template.

Environment Variables

Create a .env file in your working directory:

HF_TOKEN=your_token_here
WANDB_API_KEY=your_wandb_key_here

License

MIT