1Config file

Every training run is driven by a YAML config file. The full config for LoRA on the last 12 of 20 kept layers looks like this:

model:
  name: Qwen/Qwen3-0.6B
  num_layers: 20        # only the first 20 decoder layers are built

data:
  max_state_tokens: 1024
  max_choice_tokens: 64
  batch_size: 32        # per micro-batch
  max_train_questions: null  # null = use all

lora:
  last_k_layers: 12     # layers 8-19 get LoRA adapters
  r: 8
  alpha: 16

optim:
  epochs: 1
  lr_lora: 1.0e-4
  lr_head: 1.0e-3
  grad_accum_steps: 2   # effective batch = 32 x 2 = 64

The head section sets new_dim: 512 and num_layers: 2 for the ChoiceHead. Config values are the only thing that changes between runs; the code is config-driven throughout.

2Two LR groups

The LoRA adapter weights and the ChoiceHead weights are in separate parameter groups with different learning rates:

lora_params = [p for p in network.backbone.parameters() if p.requires_grad]
optimizer = AdamW([
    {"params": lora_params,             "lr": lr_lora},  # e.g. 1e-4
    {"params": network.head.parameters(), "lr": lr_head},  # e.g. 1e-3
], weight_decay=0.01)

3Gradient accumulation

When grad_accum_steps > 1, the optimizer step is deferred until accum micro-batches have been processed. Each micro-batch computes loss and backward normally, but the loss is divided by accum so the gradients accumulate to the correct scale:

for micro, batch in enumerate(train_loader, 1):
    logits = network(**batch)
    loss = F.cross_entropy(logits, batch["labels"])
    (loss / accum).backward()       # accumulate scaled gradients
    if micro % accum:               # not yet full — skip optimizer
        continue
    clip_grad_norm_(network.parameters(), grad_clip)
    optimizer.step()                # actual weight update
    scheduler.step()
    optimizer.zero_grad()

An incomplete final group is silently dropped: if the dataset has 65 batches and accum=2, the 65th batch is skipped. This keeps step counts exact and logging/eval intervals consistent.

4Learning rate schedule

A cosine schedule with linear warmup is applied to both parameter groups via a single scheduler. Warmup runs for warmup_ratio * total_steps steps (default 5%), after which the LR decays along a cosine curve to zero.

5Checkpointing and evaluation

Checkpoints are written under runs/<name>-<timestamp>/checkpoints/ and consist of the ChoiceHead state dict plus, when LoRA is active, the LoRA adapter weights as a safetensors file. The base Qwen weights are never stored in checkpoints — they are always loaded separately.

Validation runs every eval_every steps on the test split, reporting loss, overall accuracy, per-task-type accuracy, and yes/no AUC (when both classes appear). Metrics are logged to a JSON file in the run directory; no external logging service is required.