# ─── Model ────────────────────────────────────────────────────────────────────
# base_model: allenai/OLMoE-1B-7B-0924
base_model: nlp-projects/almo-OLMoE-1B-7B-0924-wglobalcopy-b3-layerbalancing
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer
trust_remote_code: true
output_dir: /data2/mohsenfayyaz/projects/ALMo/ALMo/src/sft/artifacts/outputs/almo-olmoe-wglobalcopy-b3-layerbalancing-ftgatewb-nopeft
hub_model_id: nlp-projects/almo-OLMoE-1B-7B-0924-wglobalcopy-b3-layerbalancing-nopeft
push_to_hub: false
hub_private_repo: false
# resume_from_checkpoint: /data2/mohsenfayyaz/projects/ALMo/ALMo/src/sft/outputs/almo-olmoe-router-only-5k/checkpoint-1000

# Limit tokenization/training data for smoke test
# max_samples: 100

# ─── Logging ──────────────────────────────────────────────────────────────────
###### WandB
wandb_project: almo
wandb_entity: mohsenfayyaz
wandb_name: almo-olmoe-wglobalcopy-b3-layerbalancing-ftgatewb-nopeft

# ─── Training mode / Frozen Parameters ────────────────────────────────────────────────────────────
# adapter: lora  # find sth lora for auto adjusting alpha
# peft_use_rslora: true
# lora_r: 64  # 16 --> 64
# lora_alpha: 64  # Change Accordingly
# lora_dropout: 0.05  # 0.05

# lora_target_modules:
  # - q_proj
  # - k_proj
  # - v_proj
  # - o_proj
  # - gate_down_proj
  # - gate_up_proj
  # - lm_head

# lora_target_parameters:          # experts (fused nn.Parameter) via parameter-LoRA --> lora_dropout: 0
#   - expert_pool.gate_up_proj
#   - expert_pool.down_proj
# ddp_find_unused_parameters: true

# lora_modules_to_save:
  # - gate          # matches every layer's mlp.gate (weight + bias), full-trainable router
  # - model.layers.*.mlp.gate.weight
  # - model.layers.*.mlp.gate.bias
  # - model.expert_pool.gate_up_proj
  # - model.expert_pool.down_proj

# Router-only training.
# Your custom model uses names like:
#   model.layers.0.mlp.gate.weight
#   model.layers.1.mlp.gate.weight
# This freezes everything except router/gate weights.

unfrozen_parameters:
  - "^model.layers.[0-9]+.mlp.gate.weight$"
  - "^model.layers.[0-9]+.mlp.gate.bias$"


# ─── Training ─────────────────────────────────────────────────────────────────
# Training
bf16: true
fp16: false
tf32: true

# gpu * micro_batch_size * gradient_accumulation_steps = 4*2*16 = 128
# 4*2*16 = 128
# 2*2*32 = 128
# 1*2*64 = 128
# 4*1*32 = 128
# last run: 4gpus*8*8=256
micro_batch_size: 8  # 8 for 1 gpu
gradient_accumulation_steps: 8  # 16 for 1 gpu
num_epochs: 1

# Required when streaming pretraining datasets
max_steps: 3000

learning_rate: 1.0e-4  # 1.0e-4
lr_scheduler: cosine
warmup_steps: 100

# optimizer: adamw_torch
optimizer: adamw_torch_fused         # fastest; switch to paged_adamw_8bit if OOM
gradient_checkpointing: false
gradient_checkpointing_kwargs:
  use_reentrant: false

# Useful because almost all params are frozen
layer_offloading: false

# dataloader_num_workers: 8
# dataloader_prefetch_factor: 8

logging_steps: 10
save_steps: 500
save_total_limit: 2
save_strategy: steps
eval_strategy: steps
load_best_model_at_end: true
metric_for_best_model: eval_loss
greater_is_better: false

# Keep this false unless you know your custom model supports cache during training
use_cache: false

# Common downstream evals, close to OLMoE/OLMES set
# plugins:
#   - axolotl.integrations.lm_eval.LMEvalPlugin

# lm_eval_tasks:
#   - mmlu
#   - hellaswag
#   - arc_easy
#   - arc_challenge
#   - boolq
#   - openbookqa
#   - piqa
#   - sciq
#   - winogrande
#   - social_iqa

# lm_eval_batch_size: 8
# lm_eval_model: hf
# lm_eval_model_args:
#   pretrained: ./outputs/almo-olmoe-router-only
#   trust_remote_code: true

# ─── Dataset ──────────────────────────────────────────────────────────────────
# Raw LM continued pretraining

dataset_prepared_path: /data2/mohsenfayyaz/projects/ALMo/ALMo/src/sft/artifacts/prepared_datasets/dolma-v1_7-3B
datasets:
  - path: emozilla/dolma-v1_7-3B
    type: completion
    text_column: text
    split: train

# pretraining_dataset:
#   - path: emozilla/dolma-v1_7-3B
#     type: pretrain
#     text_column: text
#     split: train
#     streaming: false

###### Eval during training: LM loss / perplexity proxy
# Not "exact OLMoE Paloma" unless you convert OLMo .npy eval files.
eval_sample_packing: false
eval_batch_size: 8
eval_steps: 50

test_datasets:
  - path: allenai/paloma
    type: completion
    name: wikitext_103
    text_column: text
    split: val
    streaming: false

# ─── Sequence ─────────────────────────────────────────────────────────────────
sequence_len: 2048  # 2048, 4096
sample_packing: true
pad_to_sequence_len: true

# ─── Liger kernels: ~25% speedup on Qwen ─────────────────────────────────────
plugins:
  - axolotl.integrations.liger.LigerPlugin
# Keep liger for the speedy ops:
liger_rope: true
liger_rms_norm: true
liger_glu_activation: true

# But disable this — the fused linear CE is the one that doesn't apply during eval anyway,
# and chunked_cross_entropy is the more reliable replacement
liger_fused_linear_cross_entropy: false

chunked_cross_entropy: true
chunked_cross_entropy_num_chunks: 8


# ─── Attention: FA3 via HF kernels (no compile) ──────────────────────────────
flash_attention: true
attn_implementation: kernels-community/vllm-flash-attn3

# ─── DeepSpeed ────────────────────────────────────────────────────────────────
# deepspeed: ./deepspeed_configs/zero2_bf16.json

almo-OLMoE-1B-7B-0924-wglobalcopy-b3-layerbalancing-nopeft

This model is a fine-tuned version of nlp-projects/almo-OLMoE-1B-7B-0924-wglobalcopy-b3-layerbalancing on the emozilla/dolma-v1_7-3B dataset. It achieves the following results on the evaluation set:

Loss: 2.1404
Ppl: 8.5027
Memory/max Active (gib): 18.06
Memory/max Allocated (gib): 18.06
Memory/device Reserved (gib): 64.43

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: 0.0001
train_batch_size: 8
eval_batch_size: 8
seed: 42
distributed_type: multi-GPU
num_devices: 2
gradient_accumulation_steps: 8
total_train_batch_size: 128
total_eval_batch_size: 16
optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
lr_scheduler_type: cosine
lr_scheduler_warmup_steps: 100
training_steps: 3000

Training results

Training Loss	Epoch	Step	Validation Loss	Ppl	Active (gib)	Allocated (gib)	Reserved (gib)
No log	0	0	3.9875	53.9181	21.2	21.2	21.58
2.8675	0.0057	50	2.6445	14.0758	18.06	18.06	63.41
2.4690	0.0115	100	2.2957	9.9311	18.06	18.06	63.85
2.4391	0.0172	150	2.2377	9.3719	18.06	18.06	63.87
2.4281	0.0229	200	2.2106	9.1210	18.06	18.06	64.44
2.4412	0.0286	250	2.1946	8.9766	18.06	18.06	64.49
2.3958	0.0344	300	2.1837	8.8792	18.06	18.06	64.33
2.3976	0.0401	350	2.1765	8.8155	18.06	18.06	64.22
2.3725	0.0458	400	2.1708	8.7653	18.06	18.06	61.49
2.3945	0.0515	450	2.1665	8.7274	18.06	18.06	64.44
2.3884	0.0573	500	2.1634	8.7009	18.06	18.06	64.31
2.3758	0.0630	550	2.1601	8.6724	18.06	18.06	61.39
2.3772	0.0687	600	2.1574	8.6487	18.06	18.06	64.49
2.3425	0.0745	650	2.1550	8.6276	18.06	18.06	64.51
2.3803	0.0802	700	2.1529	8.6094	18.06	18.06	64.37
2.3562	0.0859	750	2.1514	8.5970	18.06	18.06	63.98
2.3742	0.0916	800	2.1504	8.5884	18.06	18.06	61.55
2.3716	0.0974	850	2.1494	8.5798	18.06	18.06	61.46
2.3665	0.1031	900	2.1486	8.5725	18.06	18.06	64.1
2.3704	0.1088	950	2.1473	8.5620	18.06	18.06	64.42
2.3440	0.1145	1000	2.1463	8.5529	18.06	18.06	64.43
2.3731	0.1203	1050	2.1458	8.5486	18.06	18.06	64.44
2.3819	0.1260	1100	2.1455	8.5459	18.06	18.06	64.47
2.3528	0.1317	1150	2.1446	8.5387	18.06	18.06	61.46
2.3939	0.1375	1200	2.1441	8.5346	18.06	18.06	64.33
2.3764	0.1432	1250	2.1441	8.5343	18.06	18.06	61.45
2.3747	0.1489	1300	2.1434	8.5287	18.06	18.06	64.3
2.3307	0.1546	1350	2.1432	8.5266	18.06	18.06	64.49
2.3664	0.1604	1400	2.1430	8.5246	18.06	18.06	64.42
2.3397	0.1661	1450	2.1430	8.5247	18.06	18.06	61.53
2.3429	0.1718	1500	2.1422	8.5183	18.06	18.06	64.36
2.3954	0.1775	1550	2.1418	8.5151	18.06	18.06	61.44
2.3647	0.1833	1600	2.1421	8.5172	18.06	18.06	64.32
2.3349	0.1890	1650	2.1415	8.5126	18.06	18.06	64.39
2.3590	0.1947	1700	2.1412	8.5096	18.06	18.06	64.39
2.4027	0.2005	1750	2.1415	8.5123	18.06	18.06	64.42
2.3454	0.2062	1800	2.1415	8.5118	18.06	18.06	61.42
2.3464	0.2119	1850	2.1409	8.5075	18.06	18.06	61.43
2.3420	0.2176	1900	2.1411	8.5086	18.06	18.06	64.39
2.3378	0.2234	1950	2.1406	8.5044	18.06	18.06	64.4
2.3958	0.2291	2000	2.1413	8.5102	18.06	18.06	64.36
2.3613	0.2348	2050	2.1413	8.5103	18.06	18.06	61.39
2.3783	0.2405	2100	2.1413	8.5101	18.06	18.06	61.44
2.3737	0.2463	2150	2.1410	8.5080	18.06	18.06	61.37
2.3738	0.2520	2200	2.1409	8.5068	18.06	18.06	64.43
2.3226	0.2577	2250	2.1407	8.5056	18.06	18.06	61.45
2.3531	0.2635	2300	2.1406	8.5048	18.06	18.06	64.41
2.3689	0.2692	2350	2.1404	8.5030	18.06	18.06	61.4
2.3607	0.2749	2400	2.1403	8.5019	18.06	18.06	61.4
2.3431	0.2806	2450	2.1404	8.5026	18.06	18.06	64.46
2.3787	0.2864	2500	2.1403	8.5018	18.06	18.06	64.5
2.3668	0.2921	2550	2.1405	8.5035	18.06	18.06	61.45
2.3700	0.2978	2600	2.1402	8.5015	18.06	18.06	64.56
2.3793	0.3036	2650	2.1401	8.5005	18.06	18.06	64.38
2.3652	0.3093	2700	2.1404	8.5025	18.06	18.06	61.48
2.3724	0.3150	2750	2.1403	8.5017	18.06	18.06	64.46
2.3623	0.3207	2800	2.1403	8.5021	18.06	18.06	64.4
2.3474	0.3265	2850	2.1403	8.5019	18.06	18.06	61.41
2.3491	0.3322	2900	2.1400	8.4993	18.06	18.06	64.41
2.3488	0.3379	2950	2.1403	8.5021	18.06	18.06	64.4
2.3765	0.3436	3000	2.1404	8.5027	18.06	18.06	64.43

Framework versions

Transformers 5.5.0
Pytorch 2.8.0+cu129
Datasets 4.5.0
Tokenizers 0.22.2