almo-OLMoE-1B-7B-0924-wglobalcopy-b0-layerbalancing

<!-- This model card has been generated automatically according to the information the Trainer had access to. You should probably proofread and complete it, then remove this comment. -->

<img src="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>

<details><summary>See axolotl config</summary>

axolotl version: 0.16.1

# ─── Model ────────────────────────────────────────────────────────────────────
# base_model: allenai/OLMoE-1B-7B-0924
base_model: nlp-projects/almo-OLMoE-1B-7B-0924-wglobalcopy-b0-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-b0-layerbalancing-ftgatewb
hub_model_id: nlp-projects/almo-OLMoE-1B-7B-0924-wglobalcopy-b0-layerbalancing
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-b0-layerbalancing-ftgatewb

# ─── 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
gradient_accumulation_steps: 16  # 16
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

</details><br>

almo-OLMoE-1B-7B-0924-wglobalcopy-b0-layerbalancing

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

• Loss: 2.7755
• Ppl: 16.0473
• Memory/max Active (gib): 18.06
• Memory/max Allocated (gib): 18.06
• Memory/device Reserved (gib): 64.08

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
• gradient_accumulation_steps: 16
• total_train_batch_size: 128
• 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

Framework versions

• PEFT 0.19.1
• Transformers 5.5.0
• Pytorch 2.8.0+cu129
• Datasets 4.5.0
• Tokenizers 0.22.2