SocioReasoner/examples/train/rlvr_megatron.yaml at main · AMAP-ML/SocioReasoner · 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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
defaults:
  - ../config/deepspeed_zero@_here_
  - ../config/deepspeed_zero2@_here_
  - ../config/deepspeed_zero3@_here_
  - ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "qwen2_5_vl_3B_socioseg"
seed: 42
logging_dir: ./output/train/logs
output_dir: ./output/train

pretrain: Qwen/Qwen2.5-VL-3B-Instruct

checkpoint_config:
  type: file_system
  output_dir: ./output/train/checkpoint


track_with: tensorboard
tracker_kwargs:
  log_dir: ./output/train/tensorboard

save_steps: 20
logging_steps: 1
eval_steps: 20
resume_from_checkpoint: false

rollout_batch_size: 128
num_return_sequences_in_group: 8
is_num_return_sequences_expand: true
prompt_length: 4096
response_length: 2048
generate_opt_level: 0

ppo_epochs: 1
value_clip: 0.5
reward_clip: 10
advantage_clip: 10.0
# dual_clip_loss: true
# reward_shift: true
whiten_advantages: false
init_kl_coef: 0.0
adv_estimator: "grpo"
use_kl_loss: true
kl_loss_coef: 5.0e-3


validation:
 data_args:
   template: qwen2-vl
   file_name: path/to/dataset/
   dataset_dir:
 generating_args:
   max_new_tokens: ${response_length}
   top_p: 0.99
   top_k: 100
   num_beams: 1
   temperature: 0.99
   num_return_sequences: 1


actor_train:
  model_args:
    disable_gradient_checkpointing: false
    dtype: bf16
    model_type: ~
    max_pixels: 1344 * 1344
    min_pixels: 500 * 500
  training_args:
    learning_rate: 1.0e-6
    weight_decay: 1.0e-2
    per_device_train_batch_size: 2
    gradient_accumulation_steps: 4
    warmup_steps: 0
    num_train_epochs: 10
  data_args:
    template: qwen2-vl
    file_name: path/to/dataset/
    dataset_dir:
    preprocessing_num_workers: 16
  strategy_args:
    strategy_name: megatron_train
    strategy_config:
      sequence_parallel: true
      tensor_model_parallel_size: 2
      context_parallel_size: 1
      expert_model_parallel_size: 1
      pipeline_model_parallel_size: 1
      overlap_grad_reduce: true
      use_distributed_optimizer: true
      bf16: true
  device_mapping: list(range(0,4))
  infer_batch_size: 8

actor_infer:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
  generating_args:
    max_new_tokens: ${response_length}
    top_p: 0.99
    top_k: 100
    num_beams: 1
    temperature: 0.99
    num_return_sequences: ${num_return_sequences_in_group}
  data_args:
    template: qwen2-vl
    response: solution
  strategy_args:
    strategy_name: vllm
    strategy_config:
      gpu_memory_utilization: 0.8
      block_size: 16
      disable_mm_preprocessor_cache: true
      limit_mm_per_prompt:
        image: 2
  num_gpus_per_worker: 1
  device_mapping: list(range(0,4))
  system_envs:
    VLLM_USE_V1: '0'
  infer_batch_size: 24

seg_infer:
  model_args:
    model_name_or_path: facebook/sam2-hiera-large
    disable_gradient_checkpointing: true
    dtype: bf16
  strategy_args:
    strategy_name: seg_infer
  num_gpus_per_worker: 1
  device_mapping: list(range(0,4))
  infer_batch_size: 32

reference:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
    model_type: ~
  data_args:
    template: qwen2-vl
    response: solution
  strategy_args:
    strategy_name: megatron_infer
    strategy_config:
      sequence_parallel: true
      tensor_model_parallel_size: 1
      context_parallel_size: 1
      pipeline_model_parallel_size: 1
      expert_model_parallel_size: 1
      bf16: true
  device_mapping: list(range(0,4))
  infer_batch_size: 8

rewards:
  # vl pipeline support MathRuleRewardWorker only, at present.
  # Support for rewards in other domains will be retained for future implementation.
  math_rule:
    worker_cls: roll.pipeline.rlvr.rewards.socioseg_rule_reward_worker.SocioSegRuleRewardWorker
    model_args:
      model_name_or_path: ${pretrain}
    data_args:
      template: qwen2-vl
      response: solution
    strategy_args:
      strategy_name: hf_infer
      strategy_config: ~
    world_size: 16
    infer_batch_size: 4