llama-3.1-8b-lora-othello-snake-replay

llama-3.1-8b-lora-othello-snake-replay

Training Hyperparameters

Training Results

• Total steps: 966
• Best metric: 0.2670501470565796
• Best checkpoint: /scratch/project/eu-26-55/knowledge-ft/axolotl/models/llama-3.1-8b-lora-othello-snake-replay/checkpoint-966

Eval History

Axolotl Config

Note: this is the config file at push time; training parameters above are extracted from the checkpoint and reflect what was actually used.

base_model: meta-llama/Llama-3.1-8B-Instruct
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer

load_in_8bit: false
load_in_4bit: false

# ============================================================================
# LoRA fine-tune of the Othello-snake task WITH a chat-data replay mix.
# LoRA freezes the base weights and trains only a low-rank delta, so it forgets
# far less than full-FT by construction -- but we still add a lighter chat replay
# (~3:1 game:chat trained tokens, built by data/push_replay_subset.py --ratio 3)
# to keep instruction-following crisp. Two data sources are mixed:
#   1. The game, as a raw `input_output` continuation document (masked prefix).
#   2. A general chat replay set as `type: chat_template` (Llama-3 role tokens
#      give the model a context switch between game-mode and chat-mode).
# ============================================================================
chat_template: llama3
datasets:
  # --- Game shards (template-free, loss-masked continuation) ---
  # First shard sliced to hold out its last 500 games for a game-ONLY val set.
  - path: cfierro/othello-snake-llama3-fixed-prefix
    type: input_output
    split: games_0_5k[:-500]
  - path: cfierro/othello-snake-llama3-fixed-prefix
    type: input_output
    split: games_5k_10k
  - path: cfierro/othello-snake-llama3-fixed-prefix
    type: input_output
    split: games_10k_25k
  - path: cfierro/othello-snake-llama3-fixed-prefix
    type: input_output
    split: games_25k_50k
  # --- Replay: lighter (3:1) general chat mix for LoRA ---
  - path: cfierro/tulu3-sft-replay-3to1
    type: chat_template
    field_messages: messages
    split: train
train_on_inputs: false          # game: segment labels apply; chat: mask user, train assistant
dataset_prepared_path: /scratch/project/eu-26-55/knowledge-ft/axolotl/datasets/llama-3.1-8b/lora-othello-snake-replay
output_dir: /scratch/project/eu-26-55/knowledge-ft/axolotl/models/llama-3.1-8b-lora-othello-snake-replay

sequence_len: 1024              # games are short; replay was filtered to <=1024 tokens
sample_packing: true            # pack many short docs per sequence (block-diagonal attn)
eval_sample_packing: false

# --- Validation loss on GAME DATA ONLY ---
# Same scheme as the full-FT replay config: disable the mixed-set carve and point
# eval at a held-out game slice, so eval_loss == game performance and drives
# load_best_model_at_end. Slicing requires Axolotl to forward `split` to
# load_dataset -- verify with `axolotl preprocess ... --debug`; fallback is to hold
# out a whole shard (drop games_25k_50k from datasets and use it as test_datasets).
val_set_size: 0
test_datasets:
  - path: cfierro/othello-snake-llama3-fixed-prefix
    type: input_output
    split: games_0_5k[-500:]

# --- LoRA ---
# Rank 32 over all linear layers gives ample capacity for this narrow, low-entropy
# game skill; bump lora_r to 64 only if game eval_loss plateaus too high. alpha=2*r
# is a standard scaling. A little dropout regularizes the adapter.
adapter: lora
lora_model_dir:
lora_r: 32
lora_alpha: 64
lora_dropout: 0.05
lora_target_linear: true        # q,k,v,o + gate,up,down -- maximum coverage

wandb_project: knowledge-ft
wandb_entity: cfierro
wandb_watch:
wandb_name: llama-3.1-8b-lora-othello-snake-replay
wandb_log_model: "false"

# --- Multi-GPU: 4x A40 (48GB) per node ---
# LoRA is light on memory (frozen base), so micro_batch 4 fits easily at seq 1024.
# Effective batch = 4 GPUs * micro 4 * grad_accum 1 = 16 packed sequences.
# Recompute steps/epoch with `axolotl preprocess` once the replay mix is set.
gradient_accumulation_steps: 1
micro_batch_size: 4
num_epochs: 3                   # LoRA forgets little, so 3 epochs is fine

optimizer: adamw_bnb_8bit
lr_scheduler: cosine            # was 'constant' -- cosine decay is the usual LoRA default
learning_rate: 2e-4             # was 1e-5 (full-FT-scale, far too low for LoRA);
                                # 2e-4 is the standard LoRA LR for 8B. Drop to 1e-4
                                # if training is unstable.

bf16: auto
tf32: false

gradient_checkpointing: true
resume_from_checkpoint:
logging_steps: 1
flash_attention: true

warmup_ratio: 0.03
eval_strategy: epoch            # strategies must match for load_best_model_at_end
save_strategy: epoch
save_total_limit: 3             # keep all per-epoch checkpoints; best is always retained
load_best_model_at_end: true    # best by GAME-ONLY eval_loss (see test_datasets)
metric_for_best_model: eval_loss
greater_is_better: false
weight_decay: 0.0
special_tokens:
   pad_token: <|end_of_text|>

# DeepSpeed ZeRO Stage 2 — good speed/memory balance for LoRA
deepspeed: deepspeed_configs/zero2.json

# Verify the data mix + loss mask before training:
#   axolotl preprocess axolotl_configs/lora-othello-snake-8b.yaml --debug
# Confirm: game rows show the masked prefix (-100) + trained moves/<|end_of_text|>;
# chat rows show masked user/system turns + trained assistant turns.