This is a conventional full finetune distillation on logits.
It was accomplished using my project LLM-Distillery, specifically the v2.0.0 experimental branch.
Conceptually the way to accomplish this full model-agnostic architecture was by incorporating byte marginalization of logits.
So every model would be casted to a unified 256 byte distribution format + accompanying action distribution derived from special tokens, sampled at start and end of each content segment.
To make this have practical performance, for both inference and training, custom CUDA kernels were written, where training ones fuzed forward+marginalize+loss+backward.
Following teacher models were used:
Each teacher had an even contribution to the merged logit dataset:
qwen3.5_9b: 0.333
gemma4_31b: 0.333
qwen3_8b: 0.333
The loss used is a custom mix of forward+backward KL + CE diff + alpha amplification of largest losses.
Everything solely based on personal knowledge and heuristics.
Context length for collected data: 4096
Training ctx: 2048
training log:
═══ Training ═════════════════════════════════════════════════════════
Target: _merged
Student: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
Loss: abomination | Optimizer: apollo | Scheduler: wsd
Precision: bf16 | Strategy: naive_layer_split | GPUs: 1
Pre-compiling CUDA kernels...
Preparing tokenizer and byte vocab index for TinyLlama-1.1B-intermediate-step-1431k-3T...
Preprocessing training data...
Loaded dataset (format=auto)
Eligibility filter: dropped 298 samples (student incompatible or no eligible teacher).
Preprocessing validation data...
Loaded dataset (format=auto)
Eligibility filter: dropped 10 samples (student incompatible or no eligible teacher).
Loading model: TinyLlama-1.1B-intermediate-step-1431k-3T...
Compiling model with torch.compile...
Optimizer: apollo, Scheduler: wsd, LR: 1e-05
Validation - loss: 71.9039, CE: 12.8284, KL: 12.1642 <- new best achieved!
Validation - loss: 44.3783, CE: 8.6418, KL: 7.9567 <- new best achieved!
Validation - loss: 43.5121, CE: 8.5096, KL: 7.8592 <- new best achieved!
Validation - loss: 43.1144, CE: 8.4676, KL: 7.8150 <- new best achieved!
Validation - loss: 42.8548, CE: 8.4247, KL: 7.7834 <- new best achieved!
Validation - loss: 42.6769, CE: 8.3816, KL: 7.7591 <- new best achieved!
Validation - loss: 42.5271, CE: 8.3686, KL: 7.7448 <- new best achieved!
Validation - loss: 42.3926, CE: 8.3509, KL: 7.7274 <- new best achieved!
Validation - loss: 42.2945, CE: 8.3352, KL: 7.7139 <- new best achieved!
Validation - loss: 42.1791, CE: 8.3277, KL: 7.7029 <- new best achieved!
Validation - loss: 42.1273, CE: 8.3194, KL: 7.6947 <- new best achieved!
Validation - loss: 42.0750, CE: 8.3014, KL: 7.6884 <- new best achieved!
Validation - loss: 42.0191, CE: 8.2936, KL: 7.6823 <- new best achieved!
Validation - loss: 41.9545, CE: 8.2907, KL: 7.6750 <- new best achieved!
Validation - loss: 41.9155, CE: 8.2814, KL: 7.6665 <- new best achieved!
Validation - loss: 41.8525, CE: 8.2999, KL: 7.6642 <- new best achieved!
Validation - loss: 41.8450, CE: 8.2736, KL: 7.6609 <- new best achieved!
Validation - loss: 41.7732, CE: 8.2703, KL: 7.6512 <- new best achieved!
Validation - loss: 41.7341, CE: 8.2730, KL: 7.6462 <- new best achieved!
Validation - loss: 41.7573, CE: 8.2676, KL: 7.6474
Validation - loss: 41.7429, CE: 8.2667, KL: 7.6457
Training Summary:
Steps trained: 9404
Best validation loss: 41.7341
Final model saved to: data/distillery_experiment/teacher_cache\student\trained\TinyLlama-1.1B-intermediate-step-1431k-3T_step_final
Training complete.
training config:
# Paths
cache_folder = "data/distillery_experiment/teacher_cache"
dataset_path = "data/distillery_experiment/train.jsonl"
validation_dataset_path = "data/distillery_experiment/val.jsonl"
teacher_configs_path = "teacher_configs/"
student_config_path = "student_configs/tinyllama_1.1b.toml"
# General model settings
save_roles = ["assistant"]
auto_approve = true
# Collection settings
marg_chunk_inference = 256
marg_chunk_training = 256
# Training settings
num_epochs = 2
num_warmup_steps = 50
batch_size = 8
grad_accum_batches = 1
grad_checkpointing = true
max_grad_norm = 15.0
torch_compile = true
torch_compile_mode = "default"
torch_compile_backend = "inductor"
liger_kernel = false
lr = 1e-5
adam_betas = [0.9, 0.95]
adam_decay = 2e-3
lr_decay_start = 0.9
alpha = 0.5
loss_type = "abomination"
lr_scheduler = "wsd"
optimizer = "apollo"
data_order = "random"
training_precision = "bf16"
train_on = ["qwen3.5_9b", "gemma4_31b", "qwen3_8b"]
training_strategy = "naive"
seed = 42
# Validation & Saving
validate_every_n_epochs = 0.1
best_metric = "train_loss"
save_student_every_n_epochs = 1.0
save_best_model = true
save_best_state = false
keep_last_n_checkpoints = 2
# Multi-GPU / device settings
device_map = "balanced"
multi_gpu = false
wandb_comment = "byte_training"
wandb_project = "student_training"
[max_memory]
0 = "24GB"
1 = "24GB"
cpu = "60GB"
student config:
model_path = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
freeze_layers = []
attn_implementation = "flash_attention_2"
save_final_training_state = false
context_len = 2048
supports_reasoning = false
supports_tool_calls = false
chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}{{ '<|user|>\n' + message['content'] + eos_token }}{% elif message['role'] == 'system' %}{{ '<|system|>\n' + message['content'] + eos_token }}{% elif message['role'] == 'assistant' %}{{ '<|assistant|>\n' + message['content'] + eos_token }}{% endif %}{% if loop.last and add_generation_prompt %}{{ '<|assistant|>' }}{% endif %}{% endfor %}"