Views
No views yet
0.16.0.dev01# ══════════════════════════════════════════════════════════════════════════════
2# Axolotl — Full Fine-Tuning Continued Pre-Training
3# Model: Gemma 4 31B Dense (google/gemma-4-31B) — all parameters trainable
4# GPUs: 8× A100 80GB SXM (NVLink) — DeepSpeed ZeRO-3
5# Data: 321,196 chunks | 75% domain (Vedic/SPH) + 25% FineWeb-Edu
6# Tokens: ~1.013B | Sequence: 4096 | 1 epoch
7# Cost: 6× $1.49/hr = $8.94/hr → est. 55-75 hrs → $490-670
8#
9# Launch:
10# PYTORCH_ALLOC_CONF=expandable_segments:True accelerate launch --num_processes 8 -m axolotl.cli.train axolotl_cpt.yml > train.log 2>&1
11#
12# References (verified):
13# - MEDITRON-70B (EPFL, arXiv:2311.16079): FFT CPT, LR=1.5e-4, 48B tok
14# - Me-LLaMA-70B (UF, PMC/11142305): FFT CPT, LR=8e-6, 129B tok
15# - Biderman et al. (TMLR 2024, arXiv:2405.09673): FFT > LoRA for CPT
16# ══════════════════════════════════════════════════════════════════════════════
17
18# ── Model (FFT — no adapter, no quantization) ─────────────────────────────
19# Use the BASE (pre-trained) model, NOT instruction-tuned (-it).
20base_model: google/gemma-4-31B
21model_type: AutoModelForCausalLM
22tokenizer_type: AutoTokenizer
23trust_remote_code: true
24hf_use_auth_token: true
25
26# ── DeepSpeed ZeRO-3 ──────────────────────────────────────────────────────
27# Shards weights, gradients, and optimizer states across 6 GPUs.
28# Per-GPU: ~62GB sharded model state + activations → fits 80GB with grad ckpt.
29# Config ships with Axolotl — no custom JSON needed.
30deepspeed: /workspace/axolotl/deepspeed_configs/zero3_bf16.json
31
32# ── Dataset ────────────────────────────────────────────────────────────────
33# Loaded from HuggingFace Hub — pre-shuffled: 75% domain (Vedic/SPH) + 25% FineWeb-Edu.
34# Pre-chunked to ≤4096 Gemma tokens with 256-tok intra-doc overlap.
35# type: completion → each {"text": "..."} line is one sample, loss on all tokens.
36# (Do NOT use type: pretrain — that re-concatenates, breaking our chunking.)
37datasets:
38 - path: AiForgeMaster/gemma4-31b-cpt-data
39 type: completion
40 split: train
41
42dataset_prepared_path: /workspace/axolotl/axolotl_cache/cpt
43
44# ── Sequence & Packing ─────────────────────────────────────────────────────
45sequence_len: 4096
46sample_packing: true # packs shorter chunks together — eliminates padding waste
47pad_to_sequence_len: true
48gemma4_hybrid_attn_impl: true # FA2 on sliding (head_dim=256) layers, SDPA on global (head_dim=512) layers; sets flash_attention internally
49
50# ── Training Hyperparameters ───────────────────────────────────────────────
51num_epochs: 1 # one epoch — ~1,932 steps over 1.013B tokens
52
53# Effective batch = micro_batch × grad_accum × 8 GPUs = 1 × 16 × 8 = 128 samples
54# → ~524K tokens/step
55micro_batch_size: 1 # mbs=2 OOM'd on transient all-gather (8 GB failed alloc) — stay at mbs=1
56gradient_accumulation_steps: 16
57
58chunked_cross_entropy: true # avoids materializing full (B,S,V) logits tensor
59
60plugins:
61 - axolotl.integrations.liger.LigerPlugin
62liger_glu_activation: true # fused GEGLU MLP for Gemma 4 (Triton)
63liger_rms_norm: false # keep existing fused_attn.py RMSNorm patch
64liger_rope: false # Gemma 4 incompatible (separate q/k)
65liger_cross_entropy: false # chunked_cross_entropy handles this
66liger_fused_linear_cross_entropy: false # Gemma 4 incompatible
67
68optimizer: adamw_bnb_8bit # 8-bit Adam — ~6 bytes/param opt state instead of 12; saves ~23 GB/GPU
69lr_scheduler: cosine
70learning_rate: 5e-5 # conservative for FFT CPT on 31B with 1B tokens
71 # verified range: 5e-6 (Biderman) to 1.5e-4 (MEDITRON)
72 # 5e-5 balances learning vs forgetting for our data scale
73weight_decay: 0.1 # standard for FFT with AdamW (MEDITRON used 0.1)
74max_grad_norm: 1.0
75warmup_ratio: 0.01 # ~25-33 warmup steps before cosine decay
76
77# ── Precision & Memory ─────────────────────────────────────────────────────
78bf16: true
79tf32: true
80
81gradient_checkpointing: true
82gradient_checkpointing_kwargs:
83 use_reentrant: true # DeepSpeed compatibility (per Axolotl docs)
84
85# ── Output & Checkpointing ─────────────────────────────────────────────────
86# /workspace/data is on a 16 TB volume — full DS checkpoints (~310 GB each) fit fine.
87output_dir: /workspace/data/axolotl_output/gemma4-31b-cpt
88logging_steps: 10
89
90save_only_model: false # save optimizer + scheduler + RNG for exact resume
91saves_per_epoch: 4 # every ~25% of epoch (~every 483 steps / ~14 hrs)
92save_total_limit: 2 # keep latest 2 (briefly 3 during write) — ~900 GB peak on disk
93# Resume: accelerate launch ... axolotl_cpt.yml --resume_from_checkpoint <path>
94# Infra switch (different GPU count): run zero_to_fp32.py on old checkpoint,
95# then start fresh — optimizer state resets, loss wobbles briefly then recovers.
96
97val_set_size: 0 # no eval split — CPT trains on all data
98load_best_model_at_end: false
99
100# ── Weights & Biases ──────────────────────────────────────────────────────
101wandb_project: virtual_agama
102wandb_run_id: gemma4-31b-fft-stage1
103
104# ── Benchmark First! ──────────────────────────────────────────────────────
105# Before committing full budget, run a quick throughput test:
106# 1. Set max_steps: 50
107# 2. Launch training, note tokens/sec from logs
108# 3. Calculate: 1,013,000,000 / tok_per_sec / 3600 * 8.94 = total cost
109# 4. If over $500, options:
110# a) Train on domain only (760M tok) — skip GK mix, add at SFT stage
111# b) Stretch budget $50-100 — worth it for FFT quality over QLoRA
112