1# === Model Configuration ===
2base_model: mistralai/Mistral-Nemo-Instruct-2407 # e.g. "mistralai/Mistral-Small-24B-Instruct-2501"
3load_in_8bit: false
4load_in_4bit: false
5
6# === Training Setup ===
7num_epochs: 2
8micro_batch_size: 16
9gradient_accumulation_steps: 1
10sequence_len: 8192
11sample_packing: true
12pad_to_sequence_len: true
13
14# === Hyperparameter Configuration ===
15optimizer: apollo_adamw
16# Apollo-mini configuration:
17optim_args: "proj=random,rank=1,scale=128.0,scale_type=tensor,update_proj_gap=200"
18# Regular Apollo configuration:
19# optim_args:
20optim_target_modules: all_linear
21learning_rate: 1e-5
22lr_scheduler: rex
23weight_decay: 0.01
24warmup_ratio: 0.05
25
26# === Data Configuration ===
27datasets:
28 - path: allura-org/inkmix-v3.0
29 type: chat_template
30 split: train
31 field_messages: conversations
32 message_field_role: from
33 message_field_content: value
34
35dataset_prepared_path: last_run_prepared
36chat_template: jinja
37chat_template_jinja: |
38 {{- bos_token }}{%- for message in messages %}
39 {%- if message['role'] == 'system' %}
40 {{- '[SYSTEM_PROMPT]' + message['content'] + '[/SYSTEM_PROMPT]' }}
41 {%- elif message['role'] == 'user' %}
42 {{- '[INST]' + message['content'] + '[/INST]' }}
43 {%- elif message['role'] == 'assistant' %}
44 {{- message['content'] + eos_token }}
45 {%- endif %}
46 {%- endfor %}
47
48# === Plugins ===
49plugins:
50 - axolotl.integrations.liger.LigerPlugin
51 - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
52
53# === Hardware Optimization ===
54gradient_checkpointing: unsloth
55gradient_checkpointing_kwargs:
56 use_reentrant: false
57liger_rope: true
58liger_rms_norm: true
59liger_glu_activation: true
60cut_cross_entropy: true
61torch_compile: true
62# Only if using multiple GPUs:
63# deepspeed: [DEEPSPEED_CONFIG_PATH] # e.g. "deepspeed_configs/zero3_bf16.json"
64
65# === Wandb Tracking ===
66wandb_project: nemo12b-inkmix-v3
67
68# === Checkpointing ===
69saves_per_epoch: 2
70save_total_limit: 3
71
72# === Advanced Settings ===
73output_dir: offload
74bf16: auto
75flash_attention: true
76train_on_inputs: false
77group_by_length: false
78logging_steps: 1
79trust_remote_code: true
80# nemo doesnt support system prompt ootb
81tokens:
82 - "[SYSTEM_PROMPT]"
83 - "[/SYSTEM_PROMPT]"
84special_tokens:
85 pad_token: "<pad>"
86