See axolotl config

axolotl version: 0.18.0

# Gemma-4-E2B 全參數 SFT（單機 8x H100，FSDP2）— Stage 2 / 2
#
# 兩階段訓練的第二階段：接續 gemma4-e2b-sft-stage1-embed.yaml（只訓練
# input/output embedding）之後，對全部參數做完整 SFT。base_model 預設為原始
# pretrained checkpoint，方便單獨測試這份 config；實際兩階段流程由
# run_sft.sh / train_sft_stage2.sub 在 sbatch 提交時用
# `base_model=<stage1 output_dir>` 覆寫，從 stage1 訓練完的 embedding 權重接續。
#
# 模型: AlexHung29629/gemma-4-E2B (Gemma4ForConditionalGeneration，text+vision+audio，
#       safetensors ~10.2GB bf16，全參數量約 5B)。多模態模型即使做純文字訓練也走
#       multimodal pipeline（見 axolotl docs/multimodal.qmd 的 Gemma-4 章節）。
# 資料: AlexHung29629/pi_traces（tools + messages 欄位，messages content 內含
#       {"type":"image","base64":...} 圖片，transformers ProcessorMixin.apply_chat_template
#       原生支援 base64 key，不需另外轉換）。
# 目標: 全參數 SFT（無 LoRA），僅訓練語言骨幹。DDP 下每張卡要放完整模型+
#       optimizer state，記憶體吃緊（見下方 optimizer 段落），改用
#       freeze_mm_modules 凍結 vision_tower/audio_tower/embed_vision/
#       embed_audio，縮小 trainable 參數量與 optimizer state 大小。

base_model: AlexHung29629/gemma-4-E2B
hub_model_id: AlexHung29629/gemma-4-E2B-sft
# --- Multimodal ---
processor_type: AutoProcessor
skip_prepare_dataset: true
remove_unused_columns: false
sample_packing: false

# 凍結 vision/audio encoder 與 multimodal projector（embed_vision/embed_audio），
# 只訓練語言骨幹，縮小 DDP 下每張卡要維護的 optimizer state。
freeze_mm_modules: true

# --- Plugins ---
plugins:
  - axolotl.integrations.liger.LigerPlugin
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin

# Cut Cross Entropy (Apple ML)：axolotl fork 對 gemma4 有專門實作
# (cut_cross_entropy/transformers/gemma4.py 的 patch_gemma4)，直接把
# Gemma4ForConditionalGeneration.forward 換成 cce_forward_multimodal，loss
# 用融合 kernel 從 hidden_states + lm_head.weight 算，不會 materialize
# [seq_len, vocab_size] 的 fp32 logits——這正是 job 175513/175515 OOM
# 的來源。需先安裝 axolotl 的 fork：
#   pip install "cut-cross-entropy[transformers] @ \
#     git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@5f0c7a7"
# 跟 liger_fused_linear_cross_entropy（對多模態 gemma4 無效，見下方）是兩條
# 不同的路，互不衝突。
cut_cross_entropy: true

# Liger: 僅開 RMSNorm 加速。liger_glu_activation 必須關閉——axolotl 0.18.0 的
# Liger gemma4 plugin（_LigerGemma4MLP）在 patch Gemma4TextMLP 時捨棄了
# layer_idx，一律用 config.intermediate_size 建構 MLP，沒有處理 Gemma4 的
# kv-shared layer（num_kv_shared_layers=20，即 layer 15-34）use_double_wide_mlp
# 邏輯（該範圍應為 2x intermediate_size）。實測會導致 load_state_dict 時 20 層
# mlp.{up,down,gate}_proj 全部 shape mismatch 並被隨機初始化，等於報廢預訓練權重。
# liger_rope 對 Gemma4 本身不相容（plugin 會自動跳過並警告），fused_linear_cross_entropy
# 在 liger-kernel 0.8.0 對 multimodal gemma4 亦未支援，故一併關閉。
liger_rms_norm: true
liger_glu_activation: false
liger_rope: true
liger_layer_norm: false
liger_fused_linear_cross_entropy: false

strict: false

# --- Dataset ---
# 用 prepare_pi_traces.py 事先轉換過的本地 jsonl，而非直接指向
# AlexHung29629/pi_traces。原因：axolotl 0.18.0 的多模態訓練路徑
# (skip_prepare_dataset: true) 完全跳過 ChatTemplateStrategy 的 .map()
# 轉換，改由 MultiModalChatDataCollator 在 collate 時直接把原始 messages
# 丟給 processor.apply_chat_template()；但該 collator 沒有把
# tool_calls[].function.arguments 從 JSON 字串 parse 成 dict 的邏輯
# （ChatTemplateStrategy._transform_message 有做，但多模態路徑不會呼叫到）。
# pi_traces 的 arguments 是標準 OpenAI 格式字串，直接訓練會在第一個 batch
# 就因 Gemma4 chat_template.jinja 要求 arguments 必須是 dict 而報錯：
#   jinja2.exceptions.TemplateError: chat_template: tool_calls[].function
#   .arguments must be a JSON object (mapping), not a string.
# 執行 `python prepare_pi_traces.py` 產生 data/pi_traces_sft.jsonl 後即可解決；
# 已用 2-step 小樣本訓練實測驗證修正後圖片 token 與 loss 計算皆正常。
# tools/messages 欄位名稱與 axolotl 預設 field_tools="tools"、field_messages="messages"
# 一致，此處明列僅為清楚起見。
datasets:
  - path: data/pi_traces_sft.jsonl
    type: chat_template
    field_messages: messages
    field_tools: tools
    split: train
    message_property_mappings:
      role: role
      content: content
      tool_calls: tool_calls
      tool_call_id: tool_call_id
      name: name
      reasoning_content: reasoning_content

    # Note: conversations in this dataset often end with "tool" role (after last assistant action).
    # Setting to "all" ensures every EOS/EOT (incl. after tool responses) gets unmasked,
    # which avoids the "Last turn is not trainable" warning.
    # If you prefer to only unmask the last *assistant* EOS/EOT, set back to "last" (warning will appear but is usually harmless).
    train_on_eos: all
    train_on_eot: all
    roles_to_train:
      - system
      - assistant

chat_template_kwargs:
  enable_thinking: true
  preserve_thinking: true


# 用很少量資料做驗證，只是用來挑最佳 checkpoint，不是嚴謹的評估。
val_set_size: 16
dataset_num_proc: 8

# 圖片 resize 至 image_size，視覺 token 數固定為
# vision_soft_tokens_per_image=280，交由 processor 決定 patch 數量。
image_size: 512
image_resize_algorithm: bilinear

# --- Sequence length ---
# 資料集文字長度分布：median ~700 tokens，p90 ~10k tokens，少數 agentic trace
# （bash/tool 輸出）可到數萬 tokens。多模態 chat_template 目前不會依 sequence_len
# 截斷/丟棄樣本（axolotl 限制），故上限不能設太低以免整段極端長樣本被吃掉。
# 32768/16384 都曾在 loss 計算（ForCausalLMLoss 內 logits.float()，seq_len x
# vocab_size 的 fp32 logits）OOM（job 175496 附近 32768；job 175513/175515
# 的 DDP+8bit optimizer 仍在 16384 附近 OOM）。job 175519 加上 Cut Cross
# Entropy（見上方 plugins）後這個 fp32 logits materialize 已被繞開
# （device_reserved 從 62.3GiB 壓到 54.5GiB 且不再隨 step 上升），故拉回
# 16384。注意：skip_prepare_dataset=true 讓多模態路徑跳過 axolotl 內建的
# excess_length_strategy 過濾，這裡改 sequence_len 本身不會影響實際資料長度，
# 真正的過濾在 prepare_pi_traces.py 的 MAX_SEQ_LEN（需與此值同步更新後重跑
# 該腳本，才會實際把超長樣本濾掉）。
sequence_len: 16384
pad_to_sequence_len: false

# --- Full fine-tuning (no LoRA) ---
# 未設定 adapter，預設為全參數訓練；未設 freeze_mm_modules，vision/audio encoder
# 亦為可訓練參數。

attn_implementation: flash_attention_2
gemma4_hybrid_attn_impl: true

# --- DDP（放棄 FSDP2）---
# FSDP2 full-shard 在這個環境下，第一個 forward 的 all-gather 會卡死 30 分鐘
# NCCL watchdog timeout（job 175496/175508/175512 皆同一模式：rank5 的
# all-gather NumelIn 與其他 7 個 rank 不一致，懷疑是 auto_wrap_policy 對
# Gemma4 多模態架構切出不均勻的 FSDP unit，但換 wrap policy／freeze projector
# 都沒解決，根因未查清）。模型全參數約 5B，單卡 80GB 放得下完整複製，改用
# 普通 DDP 徹底繞開 FSDP2 sharding，用 gradient_checkpointing 換 activation
# 記憶體。風險：AdamW optimizer state（fp32 master + 2 moment）+ bf16
# 權重/梯度 已接近佔滿 80GB，若遇到 OOM 需再降 sequence_len 或改用
# adamw_bnb_8bit。
gradient_checkpointing: true

# --- Training ---
num_epochs: 20
micro_batch_size: 1
gradient_accumulation_steps: 8
# effective batch = 1 x 8 GPUs x 8 accum = 64 samples/step

learning_rate: 1.0e-5
lr_scheduler: constant_with_warmup
warmup_ratio: 0.03
weight_decay: 0.01
adam_beta1: 0.9
adam_beta2: 0.95
adam_epsilon: 1.0e-8
max_grad_norm: 1.0
# exp_avg/exp_avg_sq 量化成 int8（原本 adamw_torch 是 fp32），換 DDP 後每張卡
# 要放完整 optimizer state，這樣可省下可觀的記憶體，避免長序列樣本在 loss
# 計算時 OOM（見 job 175513：GPU 3 在 70.59GiB 已用時，allocate 14.54GiB 的
# logits 失敗）。用 adamw_torch_8bit（torchao 實作）而非 adamw_bnb_8bit——
# 後者在這個環境對某個大 tensor（可能是 vocab embedding）觸發 bitsandbytes
# 自己的 CUDA kernel bug："Error invalid configuration argument at line 98
# in file /src/csrc/ops.cu"（job 175514）。
optimizer: adamw_torch_8bit

bf16: true

# --- Logging / Saving / Eval ---
# 每個 epoch 評估一次 val loss，只保留 val loss 最低的 checkpoint
# （save_total_limit=1 + load_best_model_at_end：訓練結束時把最佳 checkpoint
# 載回並存到 output_dir，其餘 epoch 的 checkpoint 會被刪除）。
eval_strategy: epoch
save_strategy: best
save_total_limit: 1
load_best_model_at_end: true
metric_for_best_model: eval_loss
greater_is_better: false
logging_steps: 1

use_wandb: true
wandb_project: gemma4-e2b-sft
use_tensorboard: true

output_dir: ./outputs/gemma4-e2b-sft

seed: 42
dataloader_num_workers: 8

gemma-4-E2B-sft

This model is a fine-tuned version of AlexHung29629/gemma-4-E2B on the data/pi_traces_sft.jsonl dataset. It achieves the following results on the evaluation set:

• Loss: 0.9445
• Ppl: 2.5714
• Memory/max Active (gib): 28.36
• Memory/max Allocated (gib): 28.36
• Memory/device Reserved (gib): 54.48

Model description

More information needed

Intended uses & limitations

More information needed

Training and evaluation data

More information needed

Training procedure

Training hyperparameters

The following hyperparameters were used during training:

• learning_rate: 1e-05
• train_batch_size: 1
• eval_batch_size: 1
• seed: 42
• distributed_type: multi-GPU
• num_devices: 8
• gradient_accumulation_steps: 8
• total_train_batch_size: 64
• total_eval_batch_size: 8
• optimizer: Use OptimizerNames.ADAMW_TORCH_8BIT with betas=(0.9,0.95) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
• lr_scheduler_type: constant_with_warmup
• lr_scheduler_warmup_steps: 5
• training_steps: 185

Training results

Training Loss	Epoch	Step	Validation Loss	Ppl	Reserved (gib)	Active (gib)	Allocated (gib)
No log	0	0	5.3994	221.2723	19.8	19.57	19.57
2.6857	1.0	10	2.6677	14.4065	57.58	28.36	28.36
2.0359	2.0	20	1.8532	6.3803	56.72	28.36	28.36
1.5290	3.0	30	1.5275	4.6067	55.34	28.36	28.36
1.4167	4.0	40	1.3854	3.9966	55.5	28.36	28.36
1.3123	5.0	50	1.2836	3.6095	55.89	28.36	28.36
1.1989	6.0	60	1.2128	3.3630	54.66	28.36	28.36
1.1018	7.0	70	1.1648	3.2052	54.26	28.36	28.36
0.9369	8.0	80	1.1252	3.0808	55.05	28.36	28.36
0.8972	9.0	90	1.0929	2.9829	55.14	28.36	28.36
0.9646	10.0	100	1.0622	2.8926	55.46	28.36	28.36
0.8758	11.0	110	1.0415	2.8333	54.99	28.36	28.36
0.9764	12.0	120	1.0154	2.7605	57.07	28.36	28.36
0.7265	13.0	130	0.9968	2.7096	55.81	28.36	28.36
1.0216	14.0	140	0.9849	2.6776	56.54	28.36	28.36
0.7794	15.0	150	0.9835	2.6737	55.47	28.36	28.36
0.8172	16.0	160	0.9695	2.6365	55.02	28.36	28.36
0.8172	16.0	160	0.9647	2.6240	28.32	28.32	28.55
0.6855	17.0	170	0.9421	2.5654	28.36	28.36	58.45
0.6007	18.0	180	0.9365	2.5511	28.36	28.36	54.56
0.7379	18.5405	185	0.9445	2.5714	28.36	28.36	54.48

Framework versions

• Transformers 5.14.1
• Pytorch 2.11.0
• Datasets 4.8.4
• Tokenizers 0.22.2