forked from oumi-ai/oumi
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain.yaml
85 lines (77 loc) · 2.38 KB
/
train.yaml
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
# 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
#
# 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/**/pretraining/, configs/**/sft/, configs/**/dpo/
model:
model_name: "meta-llama/Meta-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"
use_async_dataset: True
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_dropout: 0.0
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"