README.md

metadata

title: Zero-GPU Training Space
emoji: 🧠
colorFrom: blue
colorTo: purple
sdk: gradio
sdk_version: 5.9.1
app_file: app.py
pinned: false
short_description: Train SFT/DPO/KTO one step at a time, resumable via bucket
tags:
  - gradio
  - zero-gpu
  - trl
  - training
  - lora
  - sft
  - dpo
  - kto
hardware:
  accelerator: zero-gpu
  cpu: 2
  memory: 16GB
startup_duration_timeout: 1h

🧠 Zero-GPU Training Space

Train SFT / DPO / KTO, checkpointing weights + optimizer/scheduler state to an Hugging Face Storage Bucket. The uploaded training data is never persisted — only the trained LoRA adapter and optimizer state are saved, so you can shut the Space down and resume training later without restarting the optimizer from scratch.

How a request becomes training steps

You upload a trace (a few JSON rows). The number of optimizer steps is derived from the trace, not chosen by hand:

steps = ceil(num_examples / grad_accum)        (SFT / DPO, per-device batch = 1)

Uploaded examples	GRAD_ACCUM	Optimizer steps
1 (a single trace)	4	1
10	4	3 (4 + 4 + 2)
5	4	2 (4 + 1)

So a single trace trains a single step; 10 traces train 3 steps (4 + 4 + 2). GRAD_ACCUM is set once per deployment (env). Each call continues from the last checkpoint (max_steps = last_step + steps), so the optimizer state is never restarted.

Queue: requests are not lost to a timeout

Requests are queued, not run inline — so a client-side timeout cannot drop them.

• Submit parses the trace and caches it in an in-RAM queue, then returns immediately (no GPU, no training). The request is safe the moment the call returns, even if the client then times out.
• Train drains the queue under @spaces.GPU, running jobs in method-priority order: all SFT first, then all DPO, then all KTO, FIFO within a method. If 5 SFT and 6 DPO requests arrive, every SFT job runs before any DPO job starts.
• Set QUEUE_PERSIST=true to also mirror the queue to the bucket, so queued requests survive a Space restart (the RAM queue already survives a client timeout — the Space process keeps running while a client waits).

Design: one bucket + one model per deployment

MODEL_ID and BUCKET_ID are set once via the Space's environment variables — they are not re-entered per request. The base model is loaded once and kept resident in RAM/VRAM across calls, so the expensive download

• load happens a single time. Only the LoRA adapter + optimizer state are re-applied per call (from the resumed checkpoint). This is far more efficient than re-loading the model on every request.

Why it works on a Zero-GPU / CPU Space

• GPU when available, CPU otherwise. The device is chosen from torch.cuda.is_available(). On a ZeroGPU Space @spaces.GPU attaches an A10G for the call (so it trains on GPU); on a dedicated-GPU or CPU Space the decorator is a no-op and it trains on that device.
• Resident model cache (persistent devices). On a dedicated-GPU or CPU Space the base model stays cached in memory across calls. On ZeroGPU the GPU is transient (detached between calls), so the model is loaded per-call from the local HF cache instead (no network re-download).
• LoRA. Only a small adapter is trained, so even a 2-core / 16 GB CPU container is enough for a small base model (default HuggingFaceTB/SmolLM2-135M).
• Incremental. Each call advances the run by a few steps; the HF Trainer checkpoint (adapter weights + optimizer.pt + scheduler.pt + rng + trainer_state.json) is persisted to the bucket and resumed on the next call.
• No data persistence. The trace is parsed into an in-memory datasets object for the call and discarded afterwards. Only weights + state are saved.

Configuration (Space environment variables / secrets)

Set these in the Space settings (or as secrets). Everything below is set once per deployment — the UI never asks for it per request.

Variable	Meaning	Default	Example
HF_TOKEN	token used to read/write the bucket (set as a secret)	—	hf_…
MODEL_ID	base model, one per deployment	HuggingFaceTB/SmolLM2-135M	Qwen/Qwen2.5-27B
BUCKET_ID	HF Storage Bucket, one per deployment	—	flozi00/training-checkpoints
LORA_R	LoRA rank	8	16
LORA_ALPHA	LoRA alpha	16	32
LORA_DROPOUT	LoRA dropout	0.0	0.1
LORA_TARGET_MODULES	comma-separated target modules	q_proj,k_proj,v_proj,o_proj	q_proj,v_proj
LEARNING_RATE	learning rate	1e-4	2e-5
MAX_LENGTH	max sequence length	256	512
GRAD_ACCUM	gradient-accumulation size (steps = ceil(n / G))	4	8
QUEUE_PERSIST	mirror the queue to the bucket (survives a restart)	false	true

LoRA hyper-params are fixed at first load (the base model is cached with its LoRA wrapper). To change them mid-deployment, restart the Space (which re-runs the cache).

The three methods

Method	Trainer	Trace row format
SFT	trl.SFTTrainer	{"messages":[{"role":"user","content":"..."},{"role":"assistant","content":"..."}]} or {"text":"..."}
DPO	trl.DPOTrainer	{"prompt":"...","chosen":"...","rejected":"..."}
KTO	trl.KTOTrainer	{"prompt":"...","completion":"...","desirable":true}

Accept a JSON array or JSONL (one object per line).

How it persists

For one experiment (base model + method) the bucket holds:

<bucket_id>/
  state.json            # tiny index: {global_step, method, model_id, loss}
  checkpoint-<step>/
    adapter_config.json
    adapter_model.safetensors   # LoRA weights  <-- persisted
    tokenizer.json / ...
    trainer_state.json          # global_step, LR, log   <-- persisted
    optimizer.pt                # Adam moments          <-- persisted (the key!)
    scheduler.pt                # LR scheduler state    <-- persisted
    rng.pt                      # RNG seed state        <-- persisted

On the next call the latest checkpoint is downloaded and trainer.train(resume_from_checkpoint=...) continues exactly where it left off — the optimizer does not restart from scratch.

Local development

python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
export MODEL_ID=HuggingFaceTB/SmolLM2-135M   # optional; defaults to this
export BUCKET_ID=flozi00/training-checkpoints  # optional
python app.py          # Gradio on http://localhost:7860

Deploy as a Space (Zero-GPU)

# 1) create the Space (zero-a10g flavor) + a bucket
hf repos create flozi00/zero-gpu-training-space --type space --sdk gradio --flavor zero-a10g
hf repos create flozi00/training-checkpoints --type bucket

# 2) push the code
hf upload flozi00/zero-gpu-training-space --repo-type space . . \
  --include "*.py" --include "*.md" --include "*.txt" \
  --include "examples/*.jsonl" --exclude ".venv" --exclude ".cache" --exclude "__pycache__"

# 3) set env: model + bucket (variables) and the token (a secret)
hf spaces env add flozi00/zero-gpu-training-space MODEL_ID HuggingFaceTB/SmolLM2-135M
hf spaces env add flozi00/zero-gpu-training-space BUCKET_ID flozi00/training-checkpoints
hf spaces env secret add flozi00/zero-gpu-training-space HF_TOKEN hf_...

Then open the UI: paste a trace, pick a method, click ① Submit to queue (the request is cached — safe against a client timeout), then ② Train queue (runs all SFT, then DPO, then KTO). The number of optimizer steps is derived from the trace (see above). Each restart of the Space resumes from the latest checkpoint.

Files

• app.py — Gradio UI (submit → queue → train, method-priority, resume)
• train.py — SFT/DPO/KTO trainer + dynamic grad-accum + resident model cache + resume
• job_queue.py — timeout-resilient, method-priority job queue (RAM + optional bucket)
• store.py — HF Storage Bucket persistence (weights + state only)
• examples/ — example traces for each method
• requirements.txt — ZeroGPU-compatible pinned deps