configs/recipes/llama3_1/sft/8b_qlora/train.yaml

# QLora config for Llama 8B.
# Borrows param values from:
# https://github.com/pytorch/torchtune/blob/main/recipes/configs/llama3_1/8B_qlora_single_device.yaml
#
# Requirements:
#   - Log into WandB (`wandb login`) or disable `enable_wandb`
#   - Log into HF: `hf auth login`
#   - Request access to Llama 3.1: https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct
#
# Usage:
#   oumi train -c configs/recipes/llama3_1/sft/8b_qlora/train.yaml
#
# See Also:
#   - Documentation: https://oumi.ai/docs/en/latest/user_guides/train/train.html
#   - Config class: oumi.core.configs.TrainingConfig
#   - Config source: https://github.com/oumi-ai/oumi/blob/main/src/oumi/core/configs/training_config.py
#   - Other training configs: configs/**/*train.yaml

model:
  model_name: "meta-llama/Llama-3.1-8B-Instruct"
  model_max_length: 8192
  torch_dtype_str: "bfloat16"
  attn_implementation: "sdpa"
  load_pretrained_weights: True
  trust_remote_code: True
  # Enabling Liger slows down training.
  enable_liger_kernel: False

data:
  train:
    datasets:
      - dataset_name: "yahma/alpaca-cleaned" # 51,760 examples
    target_col: "prompt"

training:
  trainer_type: "TRL_SFT"
  use_peft: True
  save_steps: 50
  num_train_epochs: 1
  per_device_train_batch_size: 2
  gradient_accumulation_steps: 32

  enable_gradient_checkpointing: True
  gradient_checkpointing_kwargs:
    use_reentrant: False
  ddp_find_unused_parameters: False
  optimizer: "adamw_torch_fused"
  learning_rate: 3.0e-04
  lr_scheduler_type: "cosine"
  warmup_steps: 100
  weight_decay: 0.01
  compile: False

  dataloader_num_workers: "auto"
  dataloader_prefetch_factor: 32

  logging_steps: 100
  log_model_summary: False
  empty_device_cache_steps: 50
  output_dir: "output/llama8b.qlora"
  include_performance_metrics: True
  enable_wandb: True

peft:
  q_lora: True
  # https://github.com/pytorch/torchtune/blob/37337f71677da69f0967a9cde34b96ad7fec3cb6/torchtune/modules/peft/lora.py#L95
  bnb_4bit_quant_type: "nf4"
  # Must use a float type for quantized data storage. See:
  # https://huggingface.co/docs/bitsandbytes/main/en/fsdp_qlora#quantized-data-storage.
  bnb_4bit_quant_storage: "bfloat16"
  bnb_4bit_compute_dtype: "bfloat16"
  lora_r: 8
  lora_alpha: 16
  lora_target_modules:
    - "q_proj"
    - "k_proj"
    - "v_proj"
    - "o_proj"
    - "gate_proj"
    - "up_proj"
    - "down_proj"

fsdp:
  enable_fsdp: True
  sharding_strategy: "HYBRID_SHARD"
  forward_prefetch: True
  auto_wrap_policy: "TRANSFORMER_BASED_WRAP"
  transformer_layer_cls: "LlamaDecoderLayer"