WhoTookMyAmogusNickname/tinyllama-1.1B-FFT-Distilled

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:

1. gemma-4-31b-it-exl3 | 6.00bpw
2. Qwen3.5-9B-exl3 | 6.00bpw
3. Qwen3-8B-exl2 | 8hb-6.0bpw

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 %}"