ssa/3B_sft_SSA_qwen.yaml at main · user074/ssa · GitHub

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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
# Config for multi-node GRPO in dev/grpo_full_finetune_distributed.py
# using a Llama3.2 3B Base model
#
# This config assumes that you've run the following command before launching
# this run:
#   tune download meta-llama/Llama-3.2-3B --output-dir /tmp/Llama-3.2-3B --ignore-patterns "original/consolidated.00.pth"
#
# It can be beneficial to first train the base model with SFT using the 3B_sft recipe.
#
# To launch on 4 devices, run the following command from root:
# tune run --nproc_per_node 4 dev/grpo_full_finetune_distributed --config dev/3B_full_grpo
#
# You can add specific overrides through the command line. For example
# to override the checkpointer directory while launching training
# you can run:
#  tune run --nproc_per_node 4 full_finetune_distributed --config ./1B_full_sft_composer_qwen.yaml
#
# This config works best when the model is being fine-tuned on 2+ GPUs.
#
# Furthermore, you can launch it on multiple nodes by going to recipes/dev/ and using
#  sbatch multinode_grpo.sbatch

name: sft_SSA

output_dir: /your/model/directory/model/checkpoints/${name}
base_model_path: /your/model/directory/model/Qwen2.5-3B  # Use this to train from the slightly trained SFT model

# Tokenizer
tokenizer:
  _component_: torchtune.models.qwen2_5.qwen2_5_tokenizer
  path: /your/model/directory/model/Qwen2.5-3B/vocab.json
  merges_file: /your/model/directory/model/Qwen2.5-3B/merges.txt
  max_seq_len: null

# Dataset
dataset:
  _component_: torchtune.dev.grpo.gsm8k.gsm8k_concat_sft
  source: concat/gsm8k/answers/with/oracle/responses
  name:
seed: null
shuffle: True

# Model Arguments
model:
  _component_: torchtune.models.qwen2_5.qwen2_5_3b

checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: ${base_model_path}
  checkpoint_files: [
    model-00001-of-00002.safetensors,
    model-00002-of-00002.safetensors,
  ]
  recipe_checkpoint: null
  output_dir: ${output_dir}
  model_type: QWEN2


ref_checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: ${base_model_path}
  checkpoint_files: [
    model-00001-of-00002.safetensors,
    model-00002-of-00002.safetensors,
  ]
  recipe_checkpoint: null
  output_dir: ${output_dir}/ref  # shouldn't be used?
  model_type: QWEN2


resume_from_checkpoint: False
batch_size: 2
epochs: 1

optimizer:
  _component_: torch.optim.AdamW
  lr: 2e-5
  fused: True
lr_scheduler:
  _component_: torchtune.training.lr_schedulers.get_cosine_schedule_with_warmup
  num_warmup_steps: 50
loss:
  _component_: torchtune.modules.loss.CEWithChunkedOutputLoss
max_steps_per_epoch: null
clip_grad_norm: null
compile: False  # torch.compile the model + loss, True increases speed + decreases memory
optimizer_in_bwd: False  # True saves memory. Requires gradient_accumulation_steps=1
gradient_accumulation_steps: 1  # Use to increase effective batch size

# Training env
device: cuda

# Memory management
enable_activation_checkpointing: True  # True reduces memory
enable_activation_offloading: False  # True reduces memory

# Reduced precision
dtype: bf16

# Logging
metric_logger:
  _component_: torchtune.training.metric_logging.WandBLogger
  project: torchtune
log_every_n_steps: 1
log_peak_memory_stats: True


# Profiler (disabled)
profiler:
  _component_: torchtune.training.setup_torch_profiler
  enabled: False

  #Output directory of trace artifacts
  output_dir: ${output_dir}/profiling_outputs

  #`torch.profiler.ProfilerActivity` types to trace
  cpu: True
  cuda: True

  #trace options passed to `torch.profiler.profile`
  profile_memory: False
  with_stack: False
  record_shapes: True
  with_flops: False

  # `torch.profiler.schedule` options:
  # wait_steps -> wait, warmup_steps -> warmup, active_steps -> active, num_cycles -> repeat
  wait_steps: 5
  warmup_steps: 3
  active_steps: 2
  num_cycles: 1