ROLL/examples/qwen2.5-7B-distill_megatron/distill_megatron.yaml at af34922be08b8c0093ea414b5f57118a807fc099 · alibaba/ROLL · 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
hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "distill_megatron"
seed: 42
logging_dir: ./output/logs
output_dir: ./output

checkpoint_config:
  type: file_system
  output_dir: /var/tmp/ckpt


save_steps: 100
logging_steps: 1
resume_from_checkpoint: false

student_pretrain: Qwen/Qwen2.5-7B-Instruct
teacher_pretrain: Qwen/Qwen2.5-14B-Instruct

# distill config
logits_topk: 64
distill_loss_weight: 0.85
kd_objective: forward_kl
distill_on_prompt: False

logits_transfer_backend: "nccl-only" # support "ipc+nccl", "nccl_only" and "ray"

sequence_length: 2048
max_grad_norm: 1.0

question_key: question_zh
answer_key: answer_zh

student:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: false
    dtype: bf16
    model_type: ~
  training_args:
    learning_rate: 2.0e-5
    lr_scheduler_type: constant
    per_device_train_batch_size: 8
    gradient_accumulation_steps: 4
    warmup_steps: 0
    num_train_epochs: 1

  data_args:
    template: qwen2_5
    file_name:
      - data/GSM8K_zh.json  #https://huggingface.co/datasets/meta-math/GSM8K_zh
    preprocessing_num_workers: 16
  strategy_args:
    strategy_name: megatron_train
    strategy_config:
      tensor_model_parallel_size: 2
      sequence_parallel: True
      pipeline_model_parallel_size: 2
      context_parallel_size: 2
      use_distributed_optimizer: true
      recompute_granularity: full
  use_sequence_packing: True
  device_mapping: list(range(0,8))

teacher:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: true
    dtype: bf16
  data_args:
    template: qwen2_5
  training_args:
    # teacher forward micro_batch_size
    per_device_train_batch_size: 8
  strategy_args:
    strategy_name: megatron_infer
    strategy_config:
      tensor_model_parallel_size: 2
      sequence_parallel: True
      pipeline_model_parallel_size: 2
      context_parallel_size: 2
      bf16: true
  use_sequence_packing: True
  device_mapping: list(range(0,8))

system_envs:
  RAY_PROFILING: "0"