1# =========================
2# Axolotl SFT config (Gemma3 1B PT, full finetune, bf16, no grad ckpt)
3# =========================
4
5# ---- Model ----
6base_model: google/gemma-3-1b-pt
7model_type: AutoModelForCausalLM
8tokenizer_type: AutoTokenizer
9
10trust_remote_code: true
11strict: false
12
13# Quantization OFF (full finetune)
14load_in_8bit: false
15load_in_4bit: false
16
17# ---- Chat formatting ----
18chat_template: gemma3
19
20# ---- Dataset ----
21datasets:
22 - path: RLHFlow/RLHFlow-SFT-Dataset-ver2
23 type: chat_template
24 field_messages: conversations
25 roles_to_train: ["assistant"]
26 split: train
27 train_on_split: train
28
29
30val_set_size: 0.01
31train_on_inputs: false # only learn on assistant tokens
32
33# ---- Tokenization / packing ----
34sequence_len: 8192
35sample_packing: true
36pad_to_sequence_len: true
37
38# Cache the prepared dataset (put on $WORK on Jean Zay)
39dataset_prepared_path: ./prepared/gemma3-1b-8192
40dataset_processes: 32 # adjust to cpus-per-task; 32 is fine on JZ A100 nodes
41dataloader_pin_memory: true
42dataloader_num_workers: 8
43dataloader_prefetch_factor: 2
44
45# ---- Output / logging ----
46output_dir: ./outputs/gemma3-1b-sft
47save_safetensors: true
48
49logging_steps: 10
50save_strategy: "epoch"
51save_total_limit: 5
52
53# Optional W&B
54wandb_project: gemma3-sft
55wandb_name: gemma3-1b-pt_seq8192_lr3e-5_bs128
56wandb_watch:
57wandb_log_model:
58
59# ---- Precision / speed ----
60bf16: true
61fp16: false
62tf32: true
63
64flash_attention: true
65xformers_attention:
66
67# ---- Training hyperparams ----
68num_epochs: 3 # start with 1 epoch to validate end-to-end; then increase
69micro_batch_size: 1
70gradient_accumulation_steps: 16 # with 8 GPUs => effective batch = 1*16*8 = 128
71
72optimizer: adamw_torch_fused
73lr_scheduler: cosine
74learning_rate: 3.0e-5 # good starting point for 1B full finetune
75warmup_ratio: 0.05
76
77weight_decay: 0.0
78max_grad_norm: 1.0
79
80group_by_length: false
81
82# ---- Memory knobs ----
83gradient_checkpointing: false # you asked: OFF for the beginning
84
85# ---- Distributed ----
86# When you launch with torchrun, Axolotl will use DDP.
87# Keep these empty so you do NOT enable ZeRO/FSDP.
88ddp:
89deepspeed:
90fsdp:
91fsdp_config:
92
93# ---- Debug ----
94debug:
95