Views
No views yet
Okay, let's see after <|channel>thought does make it bit more concise usually.1# =============================================================================
2# BASE MODEL
3# =============================================================================
4base_model: /home/arli/models/gemma-4-26B-A4B-it
5
6
7# =============================================================================
8# PLUGINS & KERNEL OPTIMIZATIONS
9# =============================================================================
10plugins:
11 - axolotl.integrations.liger.LigerPlugin # not sure if it works with Gemma 4 but it doesn't crash at least
12 - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin # must have! KV cache is too expensive otherwise
13 - axolotl.integrations.kernels.KernelsPlugin # required for scattermoe and batched_mm for efficient MoE training
14cut_cross_entropy: true
15liger_rope: true
16liger_rms_norm: true
17liger_layer_norm: true
18liger_glu_activation: true
19liger_rms_norm_gated: true
20
21use_kernels: true
22use_scattermoe: true
23experts_implementation: grouped_mm
24
25# =============================================================================
26# QUANTIZATION
27# =============================================================================
28load_in_8bit: false
29load_in_4bit: false
30
31
32# =============================================================================
33# DATASET
34# =============================================================================
35shuffle_merged_datasets: true
36datasets:
37 - path: allura-forge/musica-sft-v1-gemma4-pretok # finally, pretokenized datasets
38 ds_type: parquet
39 type:
40
41dataset_prepared_path: ./last_run_prepared
42val_set_size: 0
43
44
45# =============================================================================
46# OUTPUT & ADAPTER
47# =============================================================================
48output_dir: ./outputs/v1
49adapter: lora
50save_safetensors: true
51
52
53# =============================================================================
54# SEQUENCE & SAMPLE PACKING
55# =============================================================================
56sequence_len: 8192 # ideally 16384 but Gemma 4 31B has too expensive KV cache
57sample_packing: true # DOES in fact work with SDPA
58pad_to_sequence_len: false
59
60
61# =============================================================================
62# LORA
63# =============================================================================
64lora_r: 64
65lora_alpha: 64
66lora_dropout: 0.0
67lora_target_modules: 'model.language_model.layers.[\d]+.(_checkpoint_wrapped_module.)?(mlp|self_attn).(up|down|gate|q|k|v|o)_proj'
68lora_target_parameters:
69 - experts.gate_up_proj
70 - experts.down_proj
71
72lora_mlp_kernel: false
73lora_qkv_kernel: false
74lora_o_kernel: false
75
76
77# =============================================================================
78# TRAINING HYPERPARAMETERS
79# =============================================================================
80gradient_accumulation_steps: 4
81micro_batch_size: 2
82num_epochs: 1
83optimizer: adamw_torch_fused
84lr_scheduler: constant_with_warmup
85learning_rate: 1e-5
86warmup_ratio: 0.05
87max_grad_norm: 0.5
88weight_decay: 0.05
89
90# =============================================================================
91# PRECISION
92# =============================================================================
93bf16: auto
94
95
96# =============================================================================
97# ATTENTION
98# =============================================================================
99sdp_attention: true
100#flash_attention: true # Doesn't work on Gemma 4 currently
101#flex_attention: true # up to 40% less memory use with compile, but slower than SDPA
102#torch_compile: true # speed up, but unreliable and breaks often
103#gemma4_hybrid_attn_impl: true
104
105# =============================================================================
106# LOGGING & MONITORING
107# =============================================================================
108use_comet: true # install comet-ml with pip and do comet login before starting
109comet_project_name: musica-26b-a4b
110logging_steps: 1
111
112
113# =============================================================================
114# CHECKPOINTING & SAVING
115# =============================================================================
116auto_resume_from_checkpoints: false
117evals_per_epoch: 0
118saves_per_epoch: 4
119save_total_limit: 4
120
121gradient_checkpointing: false
122gradient_checkpointing_kwargs:
123 use_reentrant: false
124
125
126# =============================================================================
127# FSDP
128# =============================================================================
129fsdp_config:
130 fsdp_version: 2
131 offload_params: false
132 cpu_ram_efficient_loading: false
133 auto_wrap_policy: TRANSFORMER_BASED_WRAP
134 transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer
135 state_dict_type: FULL_STATE_DICT
136 sharding_strategy: FULL_SHARD
137 reshard_after_forward: true
138 activation_checkpointing: true