Training Guide¶

LORA_EXCLUDED_PATHS = ("switch_mlp", "lm_head")
LORA_EXCLUDED_NAMES = ("x_proj", "dt_proj", "gate", "shared_expert_gate")

from bit_axon.training.lora import apply_lora_to_model

# Apply DoRA (default) — returns list of wrapped layer paths
wrapped = apply_lora_to_model(model, rank=8, dropout=0.0, scale=20.0, use_dora=True)

# Apply plain LoRA
wrapped = apply_lora_to_model(model, rank=8, dropout=0.0, scale=20.0, use_dora=False)

from bit_axon.training.merging import merge_adapters

# LoRA fuse: W_fused = W + scale · B^T · A^T
# DoRA fuse: W_fused = (m / ||W + delta||) · (W + delta)
merge_adapters(model)

from bit_axon.training.packing import SequencePacker

packer = SequencePacker(max_seq_len=2048, eos_token_id=151645)

for token_ids, loss_mask in dataset:
    packed_batches = packer.add_example(token_ids, loss_mask)
    for batch in packed_batches:
        process(batch)  # PackedBatch(token_ids, loss_mask)

final = packer.flush()  # Remaining tokens, padded to max_seq_len

from bit_axon.training.cooling import CoolingScheduler, ThermalPolicy, ThermalShutdownError

policy = ThermalPolicy(
    max_speed_temp=75.0,   # Batch reduction zone starts
    pause_temp=85.0,       # Training pauses above this
    stop_temp=95.0,        # Training stops above this
    pause_duration=0.5,    # Sleep interval during pause (seconds)
)

cooling = CoolingScheduler(monitor, policy)

# Called before each training step:
cooling.check_before_step(step)  # Pauses or raises ThermalShutdownError

# Check total time spent in thermal pauses:
print(f"Paused {cooling.total_pause_time:.1f}s for cooling")

from bit_axon.training.config import TrainingConfig
from bit_axon.training.data import ORPODataset
from bit_axon.training.orpo_trainer import ORPOTrainer

config = TrainingConfig(
    training_mode="orpo",
    beta=0.1,
    max_steps=2000,
)

dataset = ORPODataset("prefs.jsonl", tokenizer, max_seq_len=2048)
trainer = ORPOTrainer(model, config, dataset, cooling_scheduler=cooling)
result = trainer.train()

# Result keys:
# step, loss, grad_norm, chosen_reward, rejected_reward, reward_margin, reward_accuracy

from bit_axon.training.config import TrainingConfig

config = TrainingConfig(
    # Optimizer
    learning_rate=1e-4,       # Peak LR after warmup
    weight_decay=0.01,        # AdamW weight decay
    warmup_steps=100,         # Linear warmup steps
    max_steps=10_000,         # Total training steps
    max_grad_norm=1.0,        # Gradient clipping threshold
    grad_accum_steps=4,       # Gradient accumulation steps

    # LoRA / DoRA
    lora_rank=8,              # Low-rank decomposition rank
    lora_dropout=0.0,         # Dropout on adapter path
    lora_scale=20.0,          # Adapter output scaling
    use_dora=True,            # Use DoRA (weight-decomposed LoRA)

    # ORPO alignment
    beta=0.1,                 # ORPO preference strength
    training_mode="sft",      # "sft" or "orpo"

    # Quantization
    quantize_bits=4,          # Base weight bit-width (NF4)
    quantize_group_size=64,   # Quantization group size

    # Data
    batch_size=1,             # Sequences per batch
    max_seq_len=2048,         # Packing target length

    # Checkpointing
    save_every=500,           # Save checkpoint every N steps
    eval_every=500,           # Evaluate every N steps
    output_dir="checkpoints", # Checkpoint directory

    # Thermal thresholds (°C)
    temp_max_speed=75.0,      # Batch reduction zone
    temp_pause=85.0,          # Pause training
    temp_stop=95.0,           # Stop training
    temp_poll_interval=1.0,   # Temperature poll interval (seconds)

    # Misc
    seed=42,                  # Random seed
)

import mlx.core as mx
from bit_axon import BitAxonConfig, BitAxonModel
from bit_axon.tokenizer import QwenTokenizerWrapper
from bit_axon.training import TrainingConfig, Trainer, apply_lora_to_model
from bit_axon.training.data import SFTDataset, CacheDataset
from bit_axon.training.cooling import CoolingScheduler, ThermalPolicy

# 1. Load model
model_config = BitAxonConfig()
model = BitAxonModel(model_config)

# 2. Apply DoRA adapters
wrapped_layers = apply_lora_to_model(
    model,
    rank=8,
    dropout=0.0,
    scale=20.0,
    use_dora=True,
)
mx.eval(model.parameters())

# 3. Prepare data
tokenizer = QwenTokenizerWrapper("Qwen/Qwen2.5-3B")
dataset = CacheDataset(SFTDataset("data.json", tokenizer, max_seq_len=2048))
val_dataset = SFTDataset("val.json", tokenizer, max_seq_len=2048)

# 4. Configure training
config = TrainingConfig(
    learning_rate=1e-4,
    max_steps=5000,
    grad_accum_steps=4,
    save_every=500,
    eval_every=500,
    output_dir="checkpoints/my-run",
)

# 5. Set up thermal monitoring
policy = ThermalPolicy(pause_temp=85.0, stop_temp=95.0)
cooling = CoolingScheduler(thermal_monitor, policy)

# 6. Train
trainer = Trainer(model, config, dataset, val_dataset, cooling_scheduler=cooling)
result = trainer.train()

print(f"Step {result['step']}: loss={result['loss']:.4f}, grad_norm={result['grad_norm']:.4f}")

from bit_axon.training.config import TrainingConfig
from bit_axon.training.data import ORPODataset
from bit_axon.training.orpo_trainer import ORPOTrainer

config = TrainingConfig(
    training_mode="orpo",
    beta=0.1,
    max_steps=2000,
    save_every=500,
)

dataset = ORPODataset("prefs.jsonl", tokenizer, max_seq_len=2048)
trainer = ORPOTrainer(model, config, dataset, cooling_scheduler=cooling)
result = trainer.train()

print(f"Reward margin: {result['reward_margin']:.4f}")
print(f"Reward accuracy: {result['reward_accuracy']:.2f}")

from bit_axon.training.lora import apply_lora_to_model, LoRALinear
from bit_axon.training.dora import DoRALinear

# Apply to all target layers (returns list of wrapped paths)
wrapped = apply_lora_to_model(model, rank=8, scale=20.0, use_dora=True)
print(f"Wrapped {len(wrapped)} layers: {wrapped[:3]}...")

# Individual layer control
dora_layer = DoRALinear.from_base(existing_linear, r=8, scale=20.0)
lora_layer = LoRALinear.from_base(existing_linear, r=8, scale=20.0)

# Fuse adapter back into base weight
fused_linear = dora_layer.fuse()  # Re-normalizes with magnitude vector

from bit_axon.training.checkpoint import (
    save_checkpoint,
    load_checkpoint,
    get_latest_checkpoint,
    save_adapter_only,
)

# Save a checkpoint manually
ckpt_path = save_checkpoint(model, optimizer, step=1500, loss=1.23, output_dir="checkpoints")

# Find the latest checkpoint
latest = get_latest_checkpoint("checkpoints")  # Returns Path or None

# Load a checkpoint (restores model weights + optimizer state)
step, loss = load_checkpoint(model, optimizer, latest)

# Export only adapter weights for sharing or deployment
save_adapter_only(model, "my_adapter.safetensors")

from bit_axon.training.merging import merge_adapters, save_merged_model, load_and_merge

# One-step: load base + adapter, merge, re-quantize, save
load_and_merge(
    base_model_path="./model",
    adapter_path="./checkpoints/final_adapter.safetensors",
    output_dir="./merged-model",
    quantize_after_merge=True,
    bits=4,
    group_size=64,
    lora_rank=8,
)

# Manual step-by-step
merge_adapters(model)
save_merged_model(model, "./merged-model", config=model_config)

from bit_axon.training.data import SFTDataset, CacheDataset

raw = SFTDataset("train.jsonl", tokenizer, max_seq_len=2048)
cached = CacheDataset(raw)

# First access tokenizes and caches; subsequent accesses hit the cache
for token_ids, loss_mask in cached:
    train_step(token_ids, loss_mask)