#Fine-tuning LLMs with Training Hub

training_hub is a Python library that wraps several LLM post-training algorithms — Supervised Fine-Tuning (SFT), Orthogonal Subspace Fine-Tuning (OSFT), LoRA / QLoRA, and continued pre-training (CPT) — behind single function calls (sft(...), osft(...), lora_sft(...)) that handle single-GPU, multi-GPU, and multi-node training uniformly.

• Automatic memory management — max_tokens_per_gpu caps GPU memory and auto-computes micro-batch size and gradient accumulation to hit your target effective_batch_size.
• OSFT implements Nayak et al., 2025 (arXiv:2504 .07097) — restricting weight updates to orthogonal subspaces prevents catastrophic forgetting without replay data.
• QLoRA loads the frozen base model in 4-bit and trains only LoRA adapters, fitting large models on a single small GPU slice (Dettmers et al., 2023, arXiv:2305 .14314).
• Continued pre-training (CPT) runs next-token prediction over a raw-text corpus to inject domain knowledge before instruction tuning.
• Built-in checkpointing, experiment tracking, and Liger kernel support.

#Requirements

• Alauda AI Workbench installed in your cluster.
• A workbench with internet (or internal PyPI mirror), at least one NVIDIA GPU, and persistent storage for checkpoints.
• HuggingFace model name or local path.
• Training data in JSONL format (see below).

#Data format

Each line is a conversation:

{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is machine learning?"}, {"role": "assistant", "content": "Machine learning is a subset of AI..."}]}

Roles: system, user, assistant, pretraining. Masking:

• SFT (default) — only assistant content contributes to loss. Add "unmask": true to a sample to include all non-system content.
• OSFT — controlled by unmask_messages (default False).

Pre-tokenized datasets with input_ids / labels are supported via use_processed_dataset=True.

#Run the example notebooks

Download into your workbench and execute cell-by-cell:

Install and configure:

pip install training-hub
# `training-hub` pulls a stray `attr` package that shadows `attrs.attr` and
# breaks `aiohttp`; uninstall it.
pip uninstall -y attr

python -m venv /workspace/venv
source /workspace/venv/bin/activate
pip install training-hub
pip uninstall -y attr

Edit the parameter cells:

model_path = "Qwen/Qwen2.5-7B-Instruct"        # HF name or local path
data_path       = "/path/to/your/training_data.jsonl"
ckpt_output_dir = "/path/to/checkpoints/my_experiment"
selected_distributed = "single_node_8gpu"      # or single_gpu_dev, multi_node_master, ...
# OSFT only:
unfreeze_rank_ratio = 0.25                      # 0.1–0.3 conservative, 0.3–0.5 balanced

Bundled model presets cover Qwen 2.5 7B, Llama 3.1 8B, Phi 4 Mini, and generic 7B / small models.

Run all cells. The final training cell calls:

from training_hub import sft, osft

result = sft(model_path=model_path, data_path=data_path, ckpt_output_dir=ckpt_output_dir,
             effective_batch_size=128, max_tokens_per_gpu=20000, max_seq_len=16384,
             learning_rate=1e-5, num_epochs=3, nproc_per_node=8)

result = osft(model_path=model_path, data_path=data_path, ckpt_output_dir=ckpt_output_dir,
              unfreeze_rank_ratio=0.25,
              effective_batch_size=128, max_tokens_per_gpu=10000, max_seq_len=8192,
              learning_rate=5e-6, num_epochs=1, nproc_per_node=8)

Checkpoints land in ckpt_output_dir at each epoch (controlled by checkpoint_at_epoch).

#Key parameters

Common (SFT and OSFT):

OSFT-only:

#QLoRA (4-bit LoRA)

QLoRA freezes the base model in 4-bit NormalFloat (NF4) precision and trains only small LoRA adapter matrices on top. A 7B model that needs ~60 GiB for full SFT then fits on a single 16–24 GiB GPU (or HAMI vGPU slice). Use it whenever you are GPU-memory bound; the cost is a small quantization quality gap and slightly slower steps.

training_hub exposes QLoRA through lora_sft(...) — LoRA plus the bitsandbytes 4-bit knobs:

from training_hub import lora_sft

result = lora_sft(
    model_path="Qwen/Qwen2.5-7B-Instruct",   # HF name or local path
    data_path="/path/to/training_data.jsonl", # same chat JSONL as SFT
    ckpt_output_dir="/path/to/checkpoints/qlora_run",
    # LoRA adapter
    lora_r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # or "all-linear"
    # 4-bit quantization — this is what makes it QLoRA
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype="bfloat16",
    bnb_4bit_use_double_quant=True,
    # core training (LoRA tolerates higher LRs than full SFT)
    num_epochs=3, effective_batch_size=16, learning_rate=2e-4, max_seq_len=4096,
    nproc_per_node=1,
)

QLoRA-specific parameters:

The output is a LoRA adapter, not a full checkpoint. To serve, load base + adapter with peft, or merge once with merge_and_unload() and export a standalone model.

#Continued pre-training (CPT)

Continued pre-training (CPT) keeps the original next-token prediction objective but runs it over a raw-text corpus from your domain (medical, legal, code, a new language). It injects knowledge and vocabulary — unlike SFT/OSFT, which teach behavior from chat data. A common pipeline is CPT → SFT → alignment.

training_hub runs CPT through the same sft(...) entrypoint with is_pretraining=True: the loss covers all tokens (no assistant-only masking) and documents are packed into fixed-size block_size windows.

from training_hub import sft

result = sft(
    model_path="Qwen/Qwen2.5-7B",             # a BASE checkpoint
    data_path="/path/to/corpus.jsonl",         # raw text, one document per line
    ckpt_output_dir="/path/to/checkpoints/cpt_run",
    # continued pre-training
    is_pretraining=True,
    block_size=4096,                           # packed context-window length
    document_column_name="text",               # JSONL field holding the raw text
    # core training — keep the LR low to limit forgetting
    num_epochs=1, effective_batch_size=128, learning_rate=5e-6,
    max_seq_len=4096, max_tokens_per_gpu=20000,
    nproc_per_node=8,
)

CPT data is raw text, not chat turns — one document per line under document_column_name:

{"text": "A paragraph or document of raw domain text ..."}
{"text": "Another document ..."}

CPT-specific parameters:

CPT updates all weights and writes a full checkpoint. It can cause catastrophic forgetting of general ability — mitigate with a low learning rate, by mixing in some general-domain text, and by following up with an SFT/OSFT pass (point model_path at the CPT checkpoint). If forgetting is the primary concern, prefer OSFT.

#Multi-node

Run the notebook (or script) on every node with the same rdzv_id / rdzv_endpoint and varying node_rank:

nproc_per_node = 8
nnodes         = 2
rdzv_id        = 42
rdzv_endpoint  = "10.0.0.1:29500"
node_rank      = 0   # 1 on the worker

All nodes need network reachability to rdzv_endpoint before training starts.