Training and Fine Tuning

Training and Fine-Tuning (HKTrainer)

Status: experimental, Python/PyTorch only. Training lives in python/hk/trainer.py and python/hk/modeling.py. It is not part of the Zig engine, which is inference-only. Two things to know before you start:

1. HKForCausalLM is HK's own small transformer (learned token embedding, multi-head attention without grouped-query or RoPE, LayerNorm, a GELU two-layer MLP mlp_fc1/mlp_fc2). It is not a loader for Llama, Qwen, Mistral or other Hugging Face architectures, and it does not read GGUF files. You can pretrain or fine-tune models of this architecture; you cannot point it at an arbitrary downloaded LLM. (from_pretrained reads .hk or .safetensors files whose tensor names match this architecture.)
2. The models this produces are not loadable by the hk run inference engine unless you export weights with matching names for a supported architecture. They are for experiments and research on the Python side.

The trainer's test coverage is a small-model smoke test (tests/ Python suites). No convergence, speed or memory results are claimed on this page.

---

What HKTrainer does

A plain PyTorch training loop (AdamW, linear warmup, gradient accumulation, grad clipping) with optional extras:

Feature	Flag	What it does
Full fine-tuning	use_qlora=False	All parameters train.
QLoRA-style adapters	use_qlora=True, lora_rank, lora_alpha	enable_qlora freezes the model, replaces linear layers with HKQuantizedLinear (base weights quantized to Q4_0 by default) and attaches trainable LoRA A/B matrices. Only the adapters get gradients.
Plasticity isolation	protect_base_capacity=True	Gradient hooks zero the updates to pre-expansion neurons and vocabulary rows (see Dynamic Architecture Growth).
Plateau growth	enable_adaptive_growth, growth_patience, growth_width_factor	When loss does not improve by >1e-4 for growth_patience checks, widens every MLP by growth_width_factor (model.grow_width, new units seeded with 1e-5 noise) and migrates AdamW moment estimates to the new shapes. Function preservation holds up to that small noise. Without an eval_dataset the training loss is the plateau signal.
Self-play / sandbox flags	enable_self_play, enable_sandbox_eval	Currently only construct a SPINLoss / CodeSandbox object on the trainer (trainer.spin_loss_fn, trainer.sandbox). The built-in training loop does not call them; they are there for custom loops. spin_lambda is not read. See Autonomous Self-Training.
Async evaluation	(with eval_dataset)	Validation runs on a background thread over a snapshot of the weights.

from hk import HKConfig, HKForCausalLM
from hk.trainer import HKTrainer, HKTrainingArguments

config = HKConfig(vocab_size=8192, hidden_size=256, num_hidden_layers=4,
                  num_attention_heads=4, intermediate_size=1024)
model = HKForCausalLM(config)

args = HKTrainingArguments(output_dir="./out", learning_rate=5e-4, batch_size=4,
                           num_train_epochs=3, use_qlora=False)
trainer = HKTrainer(model=model, args=args, train_dataset=train_ds, eval_dataset=eval_ds)
result = trainer.train()        # returns global_step, final_loss, eval_metrics, generation, output_path

Datasets yield (input_ids, labels), a dict with input_ids (and optionally labels), or bare input_ids. Loss comes from the model's labels= argument. For instruction tuning, mask the prompt tokens in labels with -100 yourself; the trainer does no chat templating or masking. If you pass no dataset, the trainer trains on one random dummy batch (a smoke-test mode, not useful training).

---

Checkpoints and the appendix

save_model calls model.save_pretrained(path, alignment=...) and then appends one lora_adapter record to the file's appendix (see In-Container Version Lineage) containing:

• the torch.save bytes of every parameter with requires_grad, plus loss/accuracy/pass-rate metrics;
• or the marker HK_BASE_WEIGHTS_SAVED if nothing is trainable.

Consequences you should know:

• With QLoRA the trainable set is just the adapters, so records are small.
• With full fine-tuning every parameter is trainable, so each appendix record is a complete copy of the model's weights. FFT checkpoints are therefore not smaller than ordinary checkpoints; there is no delta compression.
• The appendix write is wrapped in a broad try/except that ignores failures, so a failed append does not stop training and does not warn. Verify with hk appendix <file> when it matters.
• Appended adapters are not merged or applied on load; reading them back is up to your code.

---

Pretraining and sharding

Define a model through HKConfig, train with HKTrainer, and save with save_pretrained. For large checkpoints hk.torch.save_sharded_file / hk.raw.save_sharded_raw produce shard files (see Storage and Sparsity). Multi-GPU data/tensor parallelism is not implemented in the trainer; use PyTorch's own tooling around it if you need that.

---

Arguments (HKTrainingArguments)

Argument	Default	Description
output_dir	"./output"	Where checkpoints (checkpoint-<step>.hk, final model.hk) are written.
learning_rate	5e-4	Peak AdamW learning rate.
batch_size	4	Per-step batch size.
num_train_epochs	3	Epochs.
warmup_steps	10	Linear warmup steps (constant LR afterwards; no decay schedule).
weight_decay	0.01	AdamW weight decay.
logging_steps	5	Interval for averaging loss and dispatching async eval.
save_steps	50	Checkpoint interval in optimizer steps (0 disables).
gradient_accumulation_steps	1	Micro-batches per optimizer step.
max_grad_norm	1.0	Gradient clipping (0 disables).
enable_adaptive_growth	False	Plateau-triggered MLP widening.
growth_patience	5	Stagnant checks before growing.
growth_width_factor	1.25	Intermediate-size multiplier (rounded up to a multiple of 4).
protect_base_capacity	False	Gradient masks on pre-expansion capacity.
use_qlora	False	Quantize base weights and train low-rank adapters only.
lora_rank / lora_alpha	8 / 16.0	LoRA hyperparameters.
enable_self_play / spin_lambda	False / 0.1	Creates trainer.spin_loss_fn (fixed beta 0.1); not used by train(). spin_lambda is unused.
enable_sandbox_eval / sandbox_timeout_ms	False / 2000	Creates trainer.sandbox (a CodeSandbox); not used by train().
alignment	128	Payload alignment of saved checkpoints (see Raw Storage).