已合并
【R1-Zero-qwen复现 Part1】GRPO+多reward #2216
AtomGit-Bot创建于 2025年2月13日
【R1-Zero-qwen复现 Part1】GRPO+多reward #2216
已合并
从refs/pull/2216/head合入到master
共 28 个文件变更+720-86
| @@ -32,6 +32,10 @@ MindSpeed-LLM是基于昇腾生态的大语言模型分布式训练框架，旨 | |||
| 32 | 32 | ||
| 33 | **[DeepSeek-R1-Distill-LLaMA](./examples/mcore/deepseek_r1_distill_llama/)** 😊 | 33 | **[DeepSeek-R1-Distill-LLaMA](./examples/mcore/deepseek_r1_distill_llama/)** 😊 |
| 34 | 34 | ||
| 35 | +🚀🚀🚀**DeepSeek-R1** 系列功能逐步上线！！🚀🚀🚀 | ||
| 36 | + | ||
| 37 | +**[DeepSeek-R1-ZERO](./examples/mcore/deepseek_r1_recipes/)** 😊 | ||
| 38 | + | ||
| 35 | 包含数据处理、权重转换、在线推理、全参微调 | 39 | 包含数据处理、权重转换、在线推理、全参微调 |
| 36 | 40 | ||
| 37 | ## COMING !!! | 41 | ## COMING !!! |
| @@ -300,6 +300,62 @@ | |||
| 300 | ], | 300 | ], |
| 301 | "replace_eos": true | 301 | "replace_eos": true |
| 302 | }, | 302 | }, |
| 303 | + { | ||
| 304 | + "name": "qwen_r1", | ||
| 305 | + "format_user": { | ||
| 306 | + "slots": [ | ||
| 307 | + "<|im_start|>user\nA conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>Put your final answer within \\\\boxed{}. {{content}}<|im_end|>\n<|im_start|>assistant\n" | ||
| 308 | + ] | ||
| 309 | + }, | ||
| 310 | + "format_system": { | ||
| 311 | + "slots": [ | ||
| 312 | + "<|im_start|>system\n{{content}}<|im_end|>\n" | ||
| 313 | + ] | ||
| 314 | + }, | ||
| 315 | + "format_observation": { | ||
| 316 | + "slots": [ | ||
| 317 | + "<|im_start|>tool\n{{content}}<|im_end|>\n<|im_start|>assistant\n" | ||
| 318 | + ] | ||
| 319 | + }, | ||
| 320 | + "format_separator": { | ||
| 321 | + "slots": [ | ||
| 322 | + "\n" | ||
| 323 | + ] | ||
| 324 | + }, | ||
| 325 | + "default_system": "You are a helpful assistant.", | ||
| 326 | + "stop_words": [ | ||
| 327 | + "<|im_end|>" | ||
| 328 | + ], | ||
| 329 | + "replace_eos": true | ||
| 330 | + }, | ||
| 331 | + { | ||
| 332 | + "name": "qwen_math_r1", | ||
| 333 | + "format_user": { | ||
| 334 | + "slots": [ | ||
| 335 | + "<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n" | ||
| 336 | + ] | ||
| 337 | + }, | ||
| 338 | + "format_system": { | ||
| 339 | + "slots": [ | ||
| 340 | + "<|im_start|>system\n{{content}}<|im_end|>\n" | ||
| 341 | + ] | ||
| 342 | + }, | ||
| 343 | + "format_observation": { | ||
| 344 | + "slots": [ | ||
| 345 | + "<|im_start|>tool\n{{content}}<|im_end|>\n<|im_start|>assistant\n" | ||
| 346 | + ] | ||
| 347 | + }, | ||
| 348 | + "format_separator": { | ||
| 349 | + "slots": [ | ||
| 350 | + "\n" | ||
| 351 | + ] | ||
| 352 | + }, | ||
| 353 | + "default_system": "Please reason step by step, and put your final answer within \\boxed{}.", | ||
| 354 | + "stop_words": [ | ||
| 355 | + "<|im_end|>" | ||
| 356 | + ], | ||
| 357 | + "replace_eos": true | ||
| 358 | + }, | ||
| 303 | { | 359 | { |
| 304 | "name": "llama3", | 360 | "name": "llama3", |
| 305 | "format_user": { | 361 | "format_user": { |
| @@ -0,0 +1,87 @@ | |||
| 1 | +defaults: | ||
| 2 | + - model: | ||
| 3 | + - qwen25-7b | ||
| 4 | + | ||
| 5 | +training: | ||
| 6 | + global_batch_size: 12 | ||
| 7 | + seq_length: 4096 | ||
| 8 | + tokenizer_type: PretrainedFromHF | ||
| 9 | + tokenizer_name_or_path: ./models/Qwen2.5-Math-7B | ||
| 10 | + train_iters: 100 | ||
| 11 | + distributed_backend: nccl | ||
| 12 | + no_shared_storage: true | ||
| 13 | + save_interval: 10 | ||
| 14 | + no_load_optim: true | ||
| 15 | + no_load_rng: true | ||
| 16 | + bf16: true | ||
| 17 | + is_instruction_dataset: true | ||
| 18 | + variable_seq_lengths: true | ||
| 19 | + no_shuffle: false | ||
| 20 | + stage: ray_grpo | ||
| 21 | + sequence_parallel: true | ||
| 22 | + dataset_with_labels: true | ||
| 23 | + | ||
| 24 | +actor_rollout_ref: | ||
| 25 | + actor_rollout: | ||
| 26 | + model: qwen25-7b | ||
| 27 | + do_sample: true | ||
| 28 | + micro_batch_size: 2 | ||
| 29 | + ppo_mini_batch_size: 12 | ||
| 30 | + num_samples_per_step: 2 | ||
| 31 | + temperature: 1.0 | ||
| 32 | + max_prompt_length: 1536 | ||
| 33 | + max-tokens-to-oom: 1000000 | ||
| 34 | + ppo_epochs: 1 | ||
| 35 | + clip_ratio: 0.2 | ||
| 36 | + entropy_coeff: 0.001 | ||
| 37 | + shuffle_minibatch: false | ||
| 38 | + use_kv_cache: true | ||
| 39 | + tensor_model_parallel_size: 2 | ||
| 40 | + pipeline_model_parallel_size: 4 | ||
| 41 | + lr: 1e-6 | ||
| 42 | + lr_decay_style: constant | ||
| 43 | + min_lr: 0.0 | ||
| 44 | + weight_decay: 0.0 | ||
| 45 | + lr_warmup_fraction: 0.0 | ||
| 46 | + clip_grad: 10000.0 | ||
| 47 | + initial_loss_scale: 4096 | ||
| 48 | + finetune: true | ||
| 49 | + load: ./ckpt | ||
| 50 | + save: ./ckpt | ||
| 51 | + num_gpus_for_train: 8 | ||
| 52 | + num_gpus_for_infer: 6 | ||
| 53 | + data_path: ./dataset/pe-nlp/data | ||
| 54 | + split: 100,0,0 | ||
| 55 | + n_samples_per_prompt: 4 | ||
| 56 | + | ||
| 57 | + ref: | ||
| 58 | + model: qwen25-7b | ||
| 59 | + tensor_model_parallel_size: 2 | ||
| 60 | + pipeline_model_parallel_size: 1 | ||
| 61 | + micro_batch_size: 4 | ||
| 62 | + load: ./ckpt | ||
| 63 | + | ||
| 64 | +reward: | ||
| 65 | + model: qwen25-7b | ||
| 66 | + verifier: true | ||
| 67 | + tensor_model_parallel_size: 2 | ||
| 68 | + pipeline_model_parallel_size: 1 | ||
P | |||
| 69 | + micro_batch_size: 4 | ||
| 70 | + load: ./ckpt | ||
| 71 | + | ||
| 72 | +algorithm: | ||
| 73 | + gamma: 1.0 | ||
| 74 | + lam: 0.95 | ||
| 75 | + adv_estimator: group_norm | ||
| 76 | + kl_penalty: kl | ||
| 77 | + kl_ctrl: | ||
| 78 | + type: fixed | ||
| 79 | + kl_coef: 0.05 | ||
| 80 | + missing_eos_penalty: 0.0 | ||
| 81 | + verifier_function: ["acc", "format"] | ||
| 82 | + verifier_weight: [0.5, 0.5] | ||
| 83 | + | ||
| 84 | +resource_pool: | ||
| 85 | + actor_rollout: [6, 8] | ||
| 86 | + ref: [2] | ||
| 87 | + reward: [] | ||
| @@ -0,0 +1,33 @@ | |||
| 1 | +qwen25-7b: | ||
| 2 | + use_mcore_models: true | ||
| 3 | + num_layers: 28 | ||
| 4 | + hidden_size: 3584 | ||
| 5 | + ffn_hidden_size: 18944 | ||
| 6 | + num_attention_heads: 28 | ||
| 7 | + seq_length: 4096 | ||
| 8 | + rotary_base: 1000000 | ||
| 9 | + max_position_embeddings: 32768 | ||
| 10 | + make_vocab_size_divisible_by: 1 | ||
| 11 | + padded_vocab_size: 152064 | ||
| 12 | + untie_embeddings_and_output_weights: true | ||
| 13 | + add_qkv_bias: true | ||
| 14 | + disable_bias_linear: true | ||
| 15 | + group_query_attention: true | ||
| 16 | + num_query_groups: 4 | ||
| 17 | + attention_dropout: 0.0 | ||
| 18 | + init_method_std: 0.01 | ||
| 19 | + hidden_dropout: 0.0 | ||
| 20 | + adam_beta1: 0.9 | ||
| 21 | + adam_beta2: 0.95 | ||
| 22 | + position_embedding_type: rope | ||
| 23 | + normalization: RMSNorm | ||
| 24 | + use_fused_rmsnorm: true | ||
| 25 | + swiglu: true | ||
| 26 | + use_flash_attn: true | ||
| 27 | + use_mc2: true | ||
| 28 | + no_masked_softmax_fusion: true | ||
| 29 | + attention_softmax_in_fp32: true | ||
| 30 | + no_gradient_accumulation_fusion: true | ||
| 31 | + use_fused_swiglu: true | ||
| 32 | + use_fused_rotary_pos_emb: true | ||
| 33 | + bf16: true | ||
| @@ -89,7 +89,7 @@ def model_provider(pre_process=True, post_process=True) -> Union[GPTModelInfer, | |||
| 89 | pre_process=pre_process, | 89 | pre_process=pre_process, |
| 90 | post_process=post_process, | 90 | post_process=post_process, |
| 91 | fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, | 91 | fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, |
| 92 | - parallel_output=True, | 92 | + parallel_output=True if args.sequence_parallel else False, |
| 93 | share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, | 93 | share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, |
| 94 | position_embedding_type=args.position_embedding_type, | 94 | position_embedding_type=args.position_embedding_type, |
| 95 | rotary_percent=args.rotary_percent, | 95 | rotary_percent=args.rotary_percent, |
| @@ -101,7 +101,7 @@ def model_provider(pre_process=True, post_process=True) -> Union[GPTModelInfer, | |||
| 101 | 101 | ||
| 102 | model = GPTModel( | 102 | model = GPTModel( |
| 103 | config, | 103 | config, |
| 104 | - parallel_output=True, | 104 | + parallel_output=True if args.sequence_parallel else False, |
| 105 | pre_process=pre_process, | 105 | pre_process=pre_process, |
| 106 | post_process=post_process | 106 | post_process=post_process |
| 107 | ) | 107 | ) |
| @@ -0,0 +1,156 @@ | |||
| 1 | +# DeepSeek-R1-Zero(qwen) | ||
| 2 | +该特性是基于GPRO+ORM基础上扩展打分器进行混合打分而来，可用于复现DeepSeek-R1-Zero的工作 | ||
| 3 | + | ||
| 4 | +## 整体流程示意图 | ||
| 5 | + | ||
| 6 | + | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +## **模型** | ||
| 10 | + | ||
| 11 | +##### 模型选择 | ||
| 12 | +* Qwen2.5-Math-7B | ||
| 13 | +我们先尝试了该模型，发现模型在初始prompt上对数学类问题的回复存在代码文字混用，中英文混杂的情况；当加了回复格式的约束指令后，没法引导出模型按`<think></think><answer></answer>`格式回复，且回复变得容易重复且难以读懂 | ||
| 14 | +* Qwen2.5-7B-Instruct | ||
| 15 | +该模型指令遵从度高，有一定概率能引导模型输出`<think>...</think><answer>...$\boxed{}</answer>`格式回复 | ||
| 16 | + | ||
| 17 | +为了模型达到更好的效果，我们以下以Qwen2.5-7B-Instruct为例 | ||
| 18 | + | ||
| 19 | +##### 权重转换 | ||
| 20 | +同MindSpeed-LLM仓一样，模型需要从HuggingFace权重转换为MindSpeed权重，可参考[**这里**](../../../docs/features/checkpoint.md) | ||
| 21 | + | ||
| 22 | + | ||
| 23 | +## 模板构造 | ||
| 24 | +R1-Zero是基于一个base模型训练而来，由于其基底模型本身并没有激发长思维的能力，因此我们 | ||
| 25 | +需要根据选择的模型进行调试，编写prompt模板以激发`<think>...</think><answer>...$\boxed{}</answer>`，以Qwen2.5-7B-Instruct为例： | ||
| 26 | + | ||
| 27 | +``` | ||
| 28 | +<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nA conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>Put your final answer within \\boxed{}.\n{你真正的问题}<|im_end|>\n<|im_start|>assistant\n{模型真正的回答} | ||
| 29 | +``` | ||
| 30 | + | ||
| 31 | +用户可以在`configs/finetune/templates.json`添加自己的自定义模板，添加的新数据模板需要在`preprocess_data.py`与`mindspeed_llm/training/arguments.py`中的`prompt-type`参数choices里加上响应的自定义模板名字 | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +## 数据集 | ||
| 35 | +使用SimpleRL仓默认的8K数据集 | ||
| 36 | +* [pe-nlp/math_level3to5_data_processed_with_qwen_prompt](https://huggingface.co/datasets/pe-nlp/math_level3to5_data_processed_with_qwen_prompt?row=0) | ||
| 37 | + | ||
| 38 | +处理的时候要使用qwen_r1的模板（或者自行参考上一节进行构造，自定义数据集需要设置--map-keys映射，具体参考[**这里**](../../../docs/features/alpaca_dataset.md)） | ||
| 39 | + | ||
| 40 | +``` | ||
| 41 | +python ./preprocess_data.py \ | ||
| 42 | + --input ./dataset/pe-nlp/train-00000-of-00001.parquet \ | ||
| 43 | + --tokenizer-name-or-path ./models/Qwen2.5-7B-Instruct \ | ||
| 44 | + --output-prefix ./dataset/pe-nlp/data \ | ||
| 45 | + --handler-name AlpacaStyleInstructionHandler \ | ||
| 46 | + --tokenizer-type PretrainedFromHF \ | ||
| 47 | + --workers 4 \ | ||
| 48 | + --log-interval 1000 \ | ||
| 49 | + --prompt-type qwen_r1 \ | ||
| 50 | + --map-keys '{"prompt":"question", "query":"", "response": "ground_truth_answer", "system":""}' \ | ||
| 51 | + | ||
| 52 | +``` | ||
| 53 | + | ||
| 54 | +## 打分器 | ||
| 55 | +DeepSeek-R1-Zero训练的过程中仅使用了基于程序的打分器而没有使用ORM，我们在数学领域上参考OpenR1的打分逻辑分为以下几个部分： | ||
| 56 | +* [权重: 0.5] 回答风格格式分：完全遵从格式`<think> reasoning process here </think><answer> answer here with final answer within \\boxed{}</answer>`的得1分，否则得0分 | ||
| 57 | +* [权重: 0.5] 准确性得分：模型回复中boxed中提取的与标准答案一致得1分，否则-1分 | ||
| 58 | + | ||
| 59 | +## 训练 | ||
| 60 | +##### 背景 | ||
| 61 | +``` | ||
| 62 | +传统的PPO中需要一个通过广义优势估计（Generalized Advantage Estimation）计算得到的advantage，并依赖于和reward model同结构的需要同步训练的critic model计算得到价值函数(V) | ||
| 63 | +``` | ||
| 64 | + | ||
| 65 | +``` | ||
| 66 | +GRPO通过分组采样n个输出，利用组内的平均奖励作为基线计算每个输出在组内的相对奖励，并基于相对奖励计算优势值，从而避免了引入额外的价值网络（critic model） | ||
| 67 | +``` | ||
| 68 | + | ||
| 69 | + | ||
| 70 | +DeepSeek-R1-Zero的训练过程使用GRPO算法，将ORM（结果奖励模型）替换为基于规则的打分器。 | ||
| 71 | + | ||
| 72 | +##### 配置准备 | ||
| 73 | + | ||
| 74 | +模型结构的配置文件位于configs/model下，训练配置文件位于configs/rlxf目录下，我们以qwen2.5-7b为例[grpo_trainer_qwen25_7b.yaml]，以下为参数配置： | ||
| 75 | + | ||
| 76 | + | ||
| 77 | + | ||
| 78 | +``` | ||
| 79 | +defaults: | ||
| 80 | + - model: | ||
| 81 | + - qwen25-7b <-- 网络结构需要定义在model目录的yaml文件下 | ||
| 82 | + | ||
| 83 | +training: | ||
| 84 | + global_batch_size: 120 <-- 经过多少样本后acotr-train和rollout权重同步 | ||
| 85 | + ... | ||
| 86 | + has-labels: true <-- 使用打分器时需要labels,因此需要开启 | ||
| 87 | + | ||
| 88 | +actor_rollout_ref: | ||
| 89 | + | ||
| 90 | + actor_rollout: | ||
| 91 | + model: qwen25-7b | ||
| 92 | + do_sample: true <-- 一定要开启采样 | ||
| 93 | + micro_batch_size: 2 <-- 训练的mbs | ||
| 94 | + ppo_mini_batch_size: 60 <-- 一个gbs内经过多少样本后acotr-train更新 | ||
| 95 | + num_samples_per_step: 2 <-- 推理的mbs | ||
| 96 | + temperature: 1.0 <-- 推理温度，建议0.8~1.2，越大越鼓励模型探索 | ||
| 97 | + max_prompt_length: 2048 <-- 输入prompt的长度限制，输出长度=seq-length-max_prompt_length | ||
| 98 | + max-tokens-to-oom: 1000000 <-- 同一次推理最大可输入的token数 | ||
| 99 | + ... | ||
| 100 | + | ||
| 101 | + num_gpus_for_train: 8 <-- 用多少卡进行训练，注意训练和推理的TP需要保持一致 | ||
| 102 | + num_gpus_for_infer: 6 <-- 用多少卡进行推理，注意ppo_mini_batch_size % （推理节点 * num_samples_per_step）= 0 | ||
| 103 | + n_samples_per_prompt: 4 | ||
| 104 | + | ||
| 105 | + ref: | ||
| 106 | + model: qwen25-7b | ||
| 107 | + ... | ||
| 108 | + | ||
| 109 | +reward: | ||
| 110 | + model: qwen25-7b | ||
| 111 | + verifier: true <-- 是否开启打分器 | ||
| 112 | + | ||
| 113 | +algorithm: | ||
| 114 | + gamma: 1.0 | ||
| 115 | + lam: 0.95 | ||
| 116 | + adv_estimator: group_norm | ||
| 117 | + kl_penalty: kl | ||
| 118 | + kl_ctrl: | ||
| 119 | + type: fixed | ||
| 120 | + kl_coef: 0.05 | ||
| 121 | + missing_eos_penalty: 0.0 | ||
| 122 | + verifier_function: ["acc", "format"] <-- 打分器子类组合，准确性得分、格式得分 | ||
| 123 | + verifier_weight: [0.5, 0.5] <-- 打分器子类组合权重 | ||
| 124 | + | ||
| 125 | +resource_pool: | ||
| 126 | + actor_rollout: [6,8] <-- actor_rollout worker资源分配 | ||
| 127 | + ref: [2] <-- ref worker资源分配 | ||
| 128 | + reward: [] <-- rm worker资源分配，填空为不需要ORM | ||
| 129 | +``` | ||
| 130 | +##### 启动训练 | ||
| 131 | +与基于ray的其他强化训练一样，我们多机需要先在主节点初始化ray： | ||
| 132 | + | ||
| 133 | +``` | ||
| 134 | +# 创建一个集群，端口6344，dashboard端口8260，有8个NPU | ||
| 135 | +ray start --head --port 6344 --dashboard-host=0.0.0.0 --dashboard-port=8260 --resources='{"NPU": 8}' | ||
| 136 | +``` | ||
| 137 | + | ||
| 138 | +随后，在其他节点加入主节点的集群： | ||
| 139 | +``` | ||
| 140 | +# IP_ADDRESS 处填写主节点 IP 地址 | ||
| 141 | +ray start --address="IP_ADDRESS:6344" --resources='{"NPU": 8}' | ||
| 142 | +``` | ||
| 143 | + | ||
| 144 | +最后，在主节点上启动训练： | ||
| 145 | +``` | ||
| 146 | +export HCCL_CONNECT_TIMEOUT=1800 | ||
| 147 | +export CUDA_DEVICE_MAX_CONNECTIONS=1 | ||
| 148 | + | ||
| 149 | +python ray_gpt.py --config-name grpo_trainer_qwen25_7b | tee logs/r1_zero_qwen25_7b_full.log | ||
| 150 | +``` | ||
| 151 | + | ||
| 152 | +***注意：所有节点的代码、权重、数据等路径的层级要保持一致，且启动ray的时候都位于MindSpeed-LLM目录下*** | ||
| 153 | + | ||
| 154 | + | ||
| 155 | +## VLLM接入、支持DeepSeek-V3-671B-R1复现 | ||
| 156 | +comming soon ... | ||
| @@ -0,0 +1,12 @@ | |||
| 1 | +export HCCL_CONNECT_TIMEOUT=1800 | ||
| 2 | +export CUDA_DEVICE_MAX_CONNECTIONS=1 | ||
| 3 | + | ||
| 4 | +# 主节点初始化ray | ||
| 5 | +ray start --head --port 6344 --dashboard-host=0.0.0.0 --dashboard-port=8260 --resources='{"NPU": 8}' | ||
| 6 | + | ||
| 7 | +# 子节点全部注册上ray后，查看是否状态正常 | ||
| 8 | +ray status | ||
| 9 | + | ||
| 10 | +# 启动训练 | ||
| 11 | +python ray_gpt.py --config-name grpo_trainer_qwen25_7b | tee logs/r1_zero_qwen25_7b_full.log | ||
| 12 | + | ||
| @@ -0,0 +1,6 @@ | |||
| 1 | +export HCCL_CONNECT_TIMEOUT=1800 | ||
| 2 | +export CUDA_DEVICE_MAX_CONNECTIONS=1 | ||
| 3 | + | ||
| 4 | +# 子节点往主节点注册ray | ||
| 5 | +# IP_ADDRESS 处填写主节点 IP 地址 | ||
| 6 | +ray start --address="IP_ADDRESS:6344" --resources='{"NPU": 8}' | ||
| @@ -45,6 +45,9 @@ def model_provider(pre_process=True, post_process=True) -> Union[GPTModelInfer, | |||
| 45 | args = get_args() | 45 | args = get_args() |
| 46 | use_te = args.transformer_impl == "transformer_engine" | 46 | use_te = args.transformer_impl == "transformer_engine" |
| 47 | 47 | ||
| 48 | + if args.sequence_parallel and args.use_kv_cache: | ||
| 49 | + raise AssertionError('Use_kv_cache can not be true in sequence_parallel mode.') | ||
| 50 | + | ||
| 48 | print_rank_0('building GPT model ...') | 51 | print_rank_0('building GPT model ...') |
| 49 | # Experimental loading arguments from yaml | 52 | # Experimental loading arguments from yaml |
| 50 | if args.yaml_cfg is not None: | 53 | if args.yaml_cfg is not None: |
| @@ -69,7 +72,7 @@ def model_provider(pre_process=True, post_process=True) -> Union[GPTModelInfer, | |||
| 69 | pre_process=pre_process, | 72 | pre_process=pre_process, |
| 70 | post_process=post_process, | 73 | post_process=post_process, |
| 71 | fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, | 74 | fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, |
| 72 | - parallel_output=True, | 75 | + parallel_output=True if args.sequence_parallel else False, |
| 73 | share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, | 76 | share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, |
| 74 | position_embedding_type=args.position_embedding_type, | 77 | position_embedding_type=args.position_embedding_type, |
| 75 | rotary_percent=args.rotary_percent, | 78 | rotary_percent=args.rotary_percent, |
| @@ -81,7 +84,7 @@ def model_provider(pre_process=True, post_process=True) -> Union[GPTModelInfer, | |||
| 81 | 84 | ||
| 82 | model = GPTModel( | 85 | model = GPTModel( |
| 83 | config, | 86 | config, |
| 84 | - parallel_output=True, | 87 | + parallel_output=True if args.sequence_parallel else False, |
| 85 | pre_process=pre_process, | 88 | pre_process=pre_process, |
| 86 | post_process=post_process | 89 | post_process=post_process |
| 87 | ) | 90 | ) |
| @@ -87,6 +87,7 @@ def _with_pipelining_forward_step_wrapper(_with_pipelining_forward_step): | |||
| 87 | 87 | ||
| 88 | def wrapper(self, tokens, position_ids, attention_mask, micro_batch_size): | 88 | def wrapper(self, tokens, position_ids, attention_mask, micro_batch_size): |
| 89 | """No interleaving is supported.""" | 89 | """No interleaving is supported.""" |
| 90 | + args = get_args() | ||
| 90 | sequence_length = tokens.size(1) | 91 | sequence_length = tokens.size(1) |
| 91 | batch_size = tokens.size(0) | 92 | batch_size = tokens.size(0) |
| 92 | 93 | ||
| @@ -131,6 +132,7 @@ def _with_pipelining_forward_step_wrapper(_with_pipelining_forward_step): | |||
| 131 | # Here for multi batches generation. | 132 | # Here for multi batches generation. |
| 132 | if args.sequence_parallel: | 133 | if args.sequence_parallel: |
| 133 | output = gather_from_tensor_model_parallel_region(output) | 134 | output = gather_from_tensor_model_parallel_region(output) |
| 135 | + | ||
| 134 | logits[start:end, ...] = output | 136 | logits[start:end, ...] = output |
| 135 | 137 | ||
| 136 | if self.inference_params: | 138 | if self.inference_params: |
| @@ -21,7 +21,6 @@ import torch.nn.functional as F | |||
| 21 | 21 | ||
| 22 | from megatron.training import get_args, get_tokenizer | 22 | from megatron.training import get_args, get_tokenizer |
| 23 | from megatron.core import mpu | 23 | from megatron.core import mpu |
| 24 | -from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region | ||
| 25 | from megatron.inference.text_generation.communication import ( | 24 | from megatron.inference.text_generation.communication import ( |
| 26 | copy_from_last_to_first_pipeline_stage, | 25 | copy_from_last_to_first_pipeline_stage, |
| 27 | broadcast_from_last_pipeline_stage, | 26 | broadcast_from_last_pipeline_stage, |
| @@ -132,8 +131,6 @@ def generate_tokens_probs_and_return_on_first_stage( | |||
| 132 | logits = forward_step(tokens2use, positions2use, attention_mask2use) | 131 | logits = forward_step(tokens2use, positions2use, attention_mask2use) |
| 133 | 132 | ||
| 134 | if mpu.is_pipeline_last_stage(): | 133 | if mpu.is_pipeline_last_stage(): |
| 135 | - # gather along the last dim. | ||
| 136 | - logits = gather_from_tensor_model_parallel_region(logits) | ||
| 137 | # Always the last stage should have an output. | 134 | # Always the last stage should have an output. |
| 138 | assert logits is not None | 135 | assert logits is not None |
| 139 | 136 | ||
| @@ -284,8 +281,6 @@ def beam_search_and_return_on_first_stage( | |||
| 284 | logits = forward_step(tokens2use, positions2use, attention_mask2use) | 281 | logits = forward_step(tokens2use, positions2use, attention_mask2use) |
| 285 | 282 | ||
| 286 | if mpu.is_pipeline_last_stage(): | 283 | if mpu.is_pipeline_last_stage(): |
| 287 | - # gather along the last dim. | ||
| 288 | - logits = gather_from_tensor_model_parallel_region(logits) | ||
| 289 | vocab_size = logits.size(2) | 284 | vocab_size = logits.size(2) |
| 290 | 285 | ||
| 291 | if args.use_kv_cache: | 286 | if args.use_kv_cache: |
| @@ -52,6 +52,7 @@ class MegatronModuleForCausalLMABC(torch.nn.Module, abc.ABC): | |||
| 52 | self.include_input = False | 52 | self.include_input = False |
| 53 | self.stream = False | 53 | self.stream = False |
| 54 | self.return_output_log_probs = False | 54 | self.return_output_log_probs = False |
| 55 | + self.truncate = False | ||
| 55 | 56 | ||
| 56 | 57 | ||
| 57 | def from_pretrained( | 58 | def from_pretrained( |
| @@ -152,6 +153,7 @@ class MegatronModuleForCausalLMABC(torch.nn.Module, abc.ABC): | |||
| 152 | self.include_input = kwargs.pop("include_input", False) | 153 | self.include_input = kwargs.pop("include_input", False) |
| 153 | self.stream = kwargs.pop("stream", False) | 154 | self.stream = kwargs.pop("stream", False) |
| 154 | self.return_output_log_probs = kwargs.pop("return_output_log_probs", False) | 155 | self.return_output_log_probs = kwargs.pop("return_output_log_probs", False) |
| 156 | + self.truncate = kwargs.pop("truncate", False) | ||
| 155 | 157 | ||
| 156 | 158 | ||
| 157 | class MegatronModuleForCausalLM(MegatronModuleForCausalLMABC): | 159 | class MegatronModuleForCausalLM(MegatronModuleForCausalLMABC): |
| @@ -428,7 +430,7 @@ class MegatronModuleForCausalLM(MegatronModuleForCausalLMABC): | |||
| 428 | output = [val[context_lengths[i]:] for i, val in enumerate(output)] | 430 | output = [val[context_lengths[i]:] for i, val in enumerate(output)] |
| 429 | 431 | ||
| 430 | # When batch size > 1, you need truncate the tokens after eos_token_id | 432 | # When batch size > 1, you need truncate the tokens after eos_token_id |
| 431 | - self._truncate_in_multi_batch(output) | 433 | + output = self._truncate_in_multi_batch(output) |
| 432 | 434 | ||
| 433 | if self.detokenize: | 435 | if self.detokenize: |
| 434 | try: | 436 | try: |
| @@ -455,12 +457,23 @@ class MegatronModuleForCausalLM(MegatronModuleForCausalLMABC): | |||
| 455 | 457 | ||
| 456 | def _truncate_in_multi_batch(self, output): | 458 | def _truncate_in_multi_batch(self, output): |
| 457 | if len(output) > 1: | 459 | if len(output) > 1: |
| 460 | + truncated_output = [] | ||
| 458 | for idx, batch in enumerate(output): | 461 | for idx, batch in enumerate(output): |
| 459 | output[idx] = output[idx][:self.max_new_tokens] if self.max_new_tokens else output[idx] | 462 | output[idx] = output[idx][:self.max_new_tokens] if self.max_new_tokens else output[idx] |
| 460 | trunc_index = torch.nonzero(batch == self.tokenizer.eos_token_id) | 463 | trunc_index = torch.nonzero(batch == self.tokenizer.eos_token_id) |
| 461 | 464 | ||
| 462 | if min(trunc_index.shape): | 465 | if min(trunc_index.shape): |
| 463 | - output[idx][trunc_index.min():] = self.tokenizer.eos_token_id | 466 | + if self.truncate: |
| 467 | + truncated_output.append(output[idx][:trunc_index.min()]) | ||
| 468 | + else: | ||
| 469 | + output[idx][trunc_index.min():] = self.tokenizer.eos_token_id | ||
| 470 | + else: | ||
| 471 | + truncated_output.append(output[idx]) | ||
| 472 | + | ||
| 473 | + if self.truncate: | ||
| 474 | + output = [val.tolist() if torch.is_tensor(val) else val for val in truncated_output] | ||
| 475 | + | ||
| 476 | + return output | ||
| 464 | 477 | ||
| 465 | def _yield(self, token_stream): | 478 | def _yield(self, token_stream): |
| 466 | output, context_lengths, log_probs = None, None, None | 479 | output, context_lengths, log_probs = None, None, None |
| @@ -1,20 +1,19 @@ | |||
| 1 | from typing import Type | 1 | from typing import Type |
| 2 | - | ||
| 3 | from codetiming import Timer | 2 | from codetiming import Timer |
| 4 | 3 | ||
| 5 | from mindspeed_llm.tasks.posttrain.rlxf.ray_trainer.ppo_trainer import ResourcePoolManager, Role | 4 | from mindspeed_llm.tasks.posttrain.rlxf.ray_trainer.ppo_trainer import ResourcePoolManager, Role |
| 6 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.ray.megatron import NVMegatronRayWorkerGroup | 5 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.ray.megatron import NVMegatronRayWorkerGroup |
| 7 | -from mindspeed_llm.tasks.posttrain.rlxf.training.core_algos import compute_grpo_data_metrics, reduce_metrics, compute_advantage, \ | 6 | +from mindspeed_llm.tasks.posttrain.rlxf.training.core_algos import compute_grpo_data_metrics, reduce_metrics, \ |
| 8 | - get_last_reward, FixedKLController, AdaptiveKLController | 7 | + compute_advantage, compute_score, FixedKLController, AdaptiveKLController |
| 9 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base import Worker | 8 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base import Worker |
| 10 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.ray.base import create_colocated_worker_cls, \ | 9 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.ray.base import create_colocated_worker_cls, \ |
| 11 | - set_actor_infer_world_size, set_actor_train_world_size, RayResourcePool, RayClassWithInitArgs | 10 | + set_actor_infer_world_size, set_actor_train_world_size, RayClassWithInitArgs |
| 12 | from mindspeed_llm.tasks.posttrain.rlxf.utils.loggers import Loggers | 11 | from mindspeed_llm.tasks.posttrain.rlxf.utils.loggers import Loggers |
| 13 | -from mindspeed_llm.tasks.posttrain.rlxf.workers.critic import CriticWorker | ||
| 14 | from mindspeed_llm.tasks.posttrain.rlxf.workers.actor_train_infer import PPOActorWorker | 12 | from mindspeed_llm.tasks.posttrain.rlxf.workers.actor_train_infer import PPOActorWorker |
| 15 | from mindspeed_llm.tasks.posttrain.rlxf.workers.reference import ReferenceWorker | 13 | from mindspeed_llm.tasks.posttrain.rlxf.workers.reference import ReferenceWorker |
| 16 | from mindspeed_llm.tasks.posttrain.rlxf.workers.reward import RewardWorker | 14 | from mindspeed_llm.tasks.posttrain.rlxf.workers.reward import RewardWorker |
| 17 | 15 | ||
| 16 | + | ||
| 18 | WorkerType = Type[Worker] | 17 | WorkerType = Type[Worker] |
| 19 | 18 | ||
| 20 | 19 | ||
| @@ -35,17 +34,27 @@ class RayGRPOTrainer(object): | |||
| 35 | ref_pool_id = 'ref_pool' | 34 | ref_pool_id = 'ref_pool' |
| 36 | reward_pool_id = 'reward_pool' | 35 | reward_pool_id = 'reward_pool' |
| 37 | 36 | ||
| 38 | - resource_pool_spec = { | 37 | + if config.resource_pool.reward: |
| 39 | - actor_pool_id: config.resource_pool.actor_rollout, | 38 | + resource_pool_spec = { |
| 40 | - ref_pool_id: config.resource_pool.ref, | 39 | + actor_pool_id: config.resource_pool.actor_rollout, |
| 41 | - reward_pool_id: config.resource_pool.reward, | 40 | + ref_pool_id: config.resource_pool.ref, |
| 42 | - } | 41 | + reward_pool_id: config.resource_pool.reward, |
| 42 | + } | ||
| 43 | + mapping = { | ||
| 44 | + Role.ActorRollout: actor_pool_id, | ||
| 45 | + Role.RefPolicy: ref_pool_id, | ||
| 46 | + Role.RewardModel: reward_pool_id, | ||
| 47 | + } | ||
| 48 | + else: | ||
| 49 | + resource_pool_spec = { | ||
| 50 | + actor_pool_id: config.resource_pool.actor_rollout, | ||
| 51 | + ref_pool_id: config.resource_pool.ref | ||
| 52 | + } | ||
| 43 | 53 | ||
| 44 | - mapping = { | 54 | + mapping = { |
| 45 | - Role.ActorRollout: actor_pool_id, | 55 | + Role.ActorRollout: actor_pool_id, |
| 46 | - Role.RefPolicy: ref_pool_id, | 56 | + Role.RefPolicy: ref_pool_id |
| 47 | - Role.RewardModel: reward_pool_id, | 57 | + } |
| 48 | - } | ||
| 49 | 58 | ||
| 50 | self.resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) | 59 | self.resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) |
| 51 | self.use_reference_policy = Role.RefPolicy in self.role_worker_mapping | 60 | self.use_reference_policy = Role.RefPolicy in self.role_worker_mapping |
| @@ -86,11 +95,12 @@ class RayGRPOTrainer(object): | |||
| 86 | role='ref') | 95 | role='ref') |
| 87 | self.resource_pool_to_cls[resource_pool]['ref'] = ref_policy_cls | 96 | self.resource_pool_to_cls[resource_pool]['ref'] = ref_policy_cls |
| 88 | 97 | ||
| 89 | - resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) | 98 | + if self.config.resource_pool.reward: |
| 90 | - reward_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.RewardModel], | 99 | + resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) |
| 91 | - config=self.config, | 100 | + reward_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.RewardModel], |
| 92 | - role='reward') | 101 | + config=self.config, |
| 93 | - self.resource_pool_to_cls[resource_pool]['reward'] = reward_cls | 102 | + role='reward') |
| 103 | + self.resource_pool_to_cls[resource_pool]['reward'] = reward_cls | ||
| 94 | 104 | ||
| 95 | # initialize WorkerGroup | 105 | # initialize WorkerGroup |
| 96 | all_wg = {} | 106 | all_wg = {} |
| @@ -106,8 +116,11 @@ class RayGRPOTrainer(object): | |||
| 106 | self.actor_rollout_wg = all_wg.get('actor_rollout') | 116 | self.actor_rollout_wg = all_wg.get('actor_rollout') |
| 107 | self.actor_rollout_wg.initialize() | 117 | self.actor_rollout_wg.initialize() |
| 108 | 118 | ||
| 109 | - self.reward_wg = all_wg.get('reward') | 119 | + if self.config.resource_pool.reward: |
| 110 | - self.reward_wg.initialize() | 120 | + self.reward_wg = all_wg.get('reward') |
| 121 | + self.reward_wg.initialize() | ||
| 122 | + else: | ||
| 123 | + self.reward_wg = None | ||
| 111 | 124 | ||
| 112 | def train(self): | 125 | def train(self): |
| 113 | """ | 126 | """ |
| @@ -136,17 +149,21 @@ class RayGRPOTrainer(object): | |||
| 136 | 149 | ||
| 137 | with Timer(name='adv', logger=None) as timer: | 150 | with Timer(name='adv', logger=None) as timer: |
| 138 | # compute rm scores. | 151 | # compute rm scores. |
| 139 | - reward_tensor = self.reward_wg.compute_rm_score(batch) | 152 | + batch = compute_score( |
| 140 | - batch = batch.union(reward_tensor) | 153 | + self.reward_wg, |
| 141 | - batch = get_last_reward(batch, | 154 | + batch, |
| 142 | - n_sample_batch=self.config.actor_rollout_ref.actor_rollout.n_samples_per_prompt) | 155 | + metrics, |
| 156 | + self.config | ||
| 157 | + ) | ||
| 143 | 158 | ||
| 144 | # compute advantages, executed on the driver process | 159 | # compute advantages, executed on the driver process |
| 145 | - batch = compute_advantage(batch, | 160 | + batch = compute_advantage( |
| 146 | - self.config.algorithm.gamma, | 161 | + batch, |
| 147 | - self.config.algorithm.lam, | 162 | + self.config.algorithm.gamma, |
| 148 | - adv_estimator=self.config.algorithm.adv_estimator, | 163 | + self.config.algorithm.lam, |
| 149 | - ) | 164 | + adv_estimator=self.config.algorithm.adv_estimator |
| 165 | + ) | ||
| 166 | + | ||
| 150 | metrics['timing/adv'] = timer.last | 167 | metrics['timing/adv'] = timer.last |
| 151 | kl_info = {'kl_ctrl': self.kl_ctrl} | 168 | kl_info = {'kl_ctrl': self.kl_ctrl} |
| 152 | batch.meta_info.update(kl_info) | 169 | batch.meta_info.update(kl_info) |
| @@ -167,4 +184,3 @@ class RayGRPOTrainer(object): | |||
| 167 | 184 | ||
| 168 | if iteration % self.config.training.save_interval == 0: | 185 | if iteration % self.config.training.save_interval == 0: |
| 169 | self.actor_rollout_wg.save_checkpoint(iteration) | 186 | self.actor_rollout_wg.save_checkpoint(iteration) |
| 170 | - | ||
| @@ -20,9 +20,14 @@ implement PPO | |||
| 20 | import numpy as np | 20 | import numpy as np |
| 21 | import torch | 21 | import torch |
| 22 | from copy import deepcopy | 22 | from copy import deepcopy |
| 23 | +from transformers import AutoTokenizer | ||
| 23 | 24 | ||
| 24 | import mindspeed_llm.tasks.posttrain.rlxf.utils.torch_functional as F | 25 | import mindspeed_llm.tasks.posttrain.rlxf.utils.torch_functional as F |
| 26 | +from mindspeed_llm.tasks.posttrain.rlxf.utils.loggers import Loggers | ||
| 25 | from mindspeed_llm.tasks.posttrain.rlxf.utils.protocol import DataProto | 27 | from mindspeed_llm.tasks.posttrain.rlxf.utils.protocol import DataProto |
| 28 | +from mindspeed_llm.tasks.posttrain.verifier.rule_verifier import preprocess_box_response_for_qwen_prompt, format_reward, reasoning_steps_reward | ||
| 29 | + | ||
| 30 | +logger = Loggers() | ||
| 26 | 31 | ||
| 27 | 32 | ||
| 28 | class AdaptiveKLController: | 33 | class AdaptiveKLController: |
| @@ -94,8 +99,12 @@ def compute_gae_advantage_return(token_level_rewards: torch.Tensor, values: torc | |||
| 94 | return advantages, returns | 99 | return advantages, returns |
| 95 | 100 | ||
| 96 | 101 | ||
| 97 | -def compute_group_norm_advantage_return(token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, | 102 | +def compute_group_norm_advantage_return( |
| 98 | - gamma: torch.Tensor, lam: torch.Tensor): | 103 | + token_level_rewards: torch.Tensor, |
| 104 | + eos_mask: torch.Tensor, | ||
| 105 | + gamma: torch.Tensor, | ||
| 106 | + lam: torch.Tensor | ||
| 107 | +): | ||
| 99 | """ | 108 | """ |
| 100 | Args: | 109 | Args: |
| 101 | token_level_rewards: `(torch.Tensor)` | 110 | token_level_rewards: `(torch.Tensor)` |
| @@ -310,9 +319,9 @@ def compute_advantage(data: DataProto, gamma, lam, adv_estimator): | |||
| 310 | data.batch['returns'] = returns | 319 | data.batch['returns'] = returns |
| 311 | elif adv_estimator == 'group_norm': | 320 | elif adv_estimator == 'group_norm': |
| 312 | advantages, returns = compute_group_norm_advantage_return(token_level_rewards=token_level_rewards, | 321 | advantages, returns = compute_group_norm_advantage_return(token_level_rewards=token_level_rewards, |
| 313 | - eos_mask=response_mask, | 322 | + eos_mask=response_mask, |
| 314 | - gamma=gamma, | 323 | + gamma=gamma, |
| 315 | - lam=lam) | 324 | + lam=lam) |
| 316 | data.batch['advantages'] = advantages | 325 | data.batch['advantages'] = advantages |
| 317 | data.batch['returns'] = returns | 326 | data.batch['returns'] = returns |
| 318 | else: | 327 | else: |
| @@ -320,20 +329,117 @@ def compute_advantage(data: DataProto, gamma, lam, adv_estimator): | |||
| 320 | return data | 329 | return data |
| 321 | 330 | ||
| 322 | 331 | ||
| 323 | -def get_last_reward(data, n_sample_batch): | 332 | +def compute_score(reward_wg, batch, metrics, config): |
| 333 | + token_level_rewards = torch.zeros_like(batch.batch["responses"], dtype=torch.float32) | ||
| 334 | + | ||
| 335 | + if reward_wg: | ||
| 336 | + score_tensor = reward_wg.compute_rm_score(batch).batch['rm_scores'] | ||
| 337 | + | ||
| 338 | + rm_token_level_rewards = get_last_reward( | ||
| 339 | + batch, | ||
| 340 | + rm_scores=score_tensor, | ||
| 341 | + n_sample_batch=config.actor_rollout_ref.actor_rollout.n_samples_per_prompt | ||
| 342 | + ) | ||
| 343 | + | ||
| 344 | + token_level_rewards += rm_token_level_rewards | ||
| 345 | + | ||
| 346 | + | ||
| 347 | + if hasattr(config.reward, "verifier") and config.reward.verifier: | ||
| 348 | + verifier_token_level_rewards = compute_verifier_score(batch, metrics, config) | ||
| 349 | + token_level_rewards += verifier_token_level_rewards | ||
| 350 | + | ||
| 351 | + rewards = DataProto.from_dict( | ||
| 352 | + tensors={ | ||
| 353 | + 'token_level_rewards': token_level_rewards, | ||
| 354 | + 'rm_scores': token_level_rewards | ||
| 355 | + } | ||
| 356 | + ) | ||
| 357 | + | ||
| 358 | + return batch.union(rewards) | ||
| 359 | + | ||
| 360 | + | ||
| 361 | +def compute_verifier_score(batch, metrics, config): | ||
| 362 | + tokenizer = AutoTokenizer.from_pretrained(config.training.tokenizer_name_or_path, trust_remote_code=True) | ||
| 363 | + tokenizer.eos_token = tokenizer.decode(tokenizer.eos_token_id, skip_special_tokens=True) | ||
| 364 | + | ||
| 365 | + responses = batch.batch["responses"] | ||
| 366 | + reward_index = batch.batch["responses_ori_length"].unsqueeze(1) - 1 | ||
| 367 | + | ||
| 368 | + str_labels = tokenizer.batch_decode(batch.batch["labels"], skip_special_tokens=True) | ||
| 369 | + str_responses = tokenizer.batch_decode(responses, skip_special_tokens=True) | ||
| 370 | + | ||
| 371 | + scores = verifier(str_responses, str_labels, config, metrics, infos=None) | ||
| 372 | + | ||
| 373 | + scores = torch.tensor( | ||
| 374 | + scores, | ||
| 375 | + dtype=torch.float32, | ||
| 376 | + device=reward_index.device | ||
| 377 | + ).unsqueeze(1) | ||
| 378 | + | ||
| 379 | + logger.logger.info("=" * 50) | ||
| 380 | + logger.logger.info(str_responses[0]) | ||
| 381 | + logger.logger.info(">" * 50) | ||
| 382 | + logger.logger.info(str_labels[0]) | ||
| 383 | + logger.logger.info("=" * 50) | ||
| 384 | + | ||
| 385 | + token_level_rewards = torch.zeros_like(responses, dtype=torch.float32) | ||
| 386 | + token_level_rewards.scatter_(1, reward_index, scores) | ||
| 387 | + | ||
| 388 | + return token_level_rewards | ||
| 389 | + | ||
| 390 | + | ||
| 391 | +def verifier(responses, labels, config, metrics, infos=None): | ||
| 392 | + """ | ||
| 393 | + User-defined verifier scoring process. | ||
| 394 | + | ||
| 395 | + Parameters: | ||
| 396 | + ---------- | ||
| 397 | + responses(List[`str`]): | ||
| 398 | + Actor rollout answers. | ||
| 399 | + labels(List[`str`]): | ||
| 400 | + Ground Truth. | ||
| 401 | + infos(List[`str`], *optional*): | ||
| 402 | + Additional usable information loaded from the dataset. | ||
| 403 | + | ||
| 404 | + Return: | ||
| 405 | + scores(List[`float`]): Final scores. | ||
| 406 | + """ | ||
| 407 | + rule_verifier_function = { | ||
| 408 | + "acc": preprocess_box_response_for_qwen_prompt, | ||
| 409 | + "format": format_reward, | ||
| 410 | + "step": reasoning_steps_reward | ||
| 411 | + } | ||
| 412 | + | ||
| 413 | + scores = [0.0] * len(labels) | ||
| 414 | + | ||
| 415 | + verifier_function = config.algorithm.verifier_function if hasattr( | ||
| 416 | + config.algorithm, "verifier_function") else ["acc"] | ||
| 417 | + verifier_weight = config.algorithm.verifier_weight if hasattr( | ||
| 418 | + config.algorithm, "verifier_weight") else [1.0] | ||
| 419 | + | ||
| 420 | + for idx, fun_verifier in enumerate(verifier_function): | ||
| 421 | + if fun_verifier not in rule_verifier_function: | ||
| 422 | + continue | ||
| 423 | + score = rule_verifier_function[fun_verifier](sequences=responses, answers=labels) | ||
| 424 | + metrics[f"grpo/{fun_verifier}_rewards/mean"] = sum(score) / max(len(score), 1) | ||
| 425 | + scores = [all_score + tmp_score * verifier_weight[idx] | ||
| 426 | + for all_score, tmp_score in zip(scores, score)] | ||
| 427 | + | ||
| 428 | + return scores | ||
| 429 | + | ||
| 430 | + | ||
| 431 | +def get_last_reward(data, rm_scores, n_sample_batch): | ||
| 324 | responses = data.batch['responses'] | 432 | responses = data.batch['responses'] |
| 325 | response_length = responses.size(1) | 433 | response_length = responses.size(1) |
| 326 | attention_mask = data.batch['attention_mask'] | 434 | attention_mask = data.batch['attention_mask'] |
| 327 | response_mask = attention_mask[:, -response_length:] | 435 | response_mask = attention_mask[:, -response_length:] |
| 328 | - rm_scores = data.batch['rm_scores'] | ||
| 329 | eos_indices = response_mask.size(1) - 1 - response_mask.long().fliplr().argmax(dim=1, keepdim=True) | 436 | eos_indices = response_mask.size(1) - 1 - response_mask.long().fliplr().argmax(dim=1, keepdim=True) |
| 330 | reward = rm_scores.gather(dim=1, index=eos_indices).squeeze(1) | 437 | reward = rm_scores.gather(dim=1, index=eos_indices).squeeze(1) |
| 331 | reward = reward.reshape(-1, n_sample_batch) | 438 | reward = reward.reshape(-1, n_sample_batch) |
| 332 | reward = (reward - reward.mean(dim=1, keepdim=True)) / (reward.std(dim=1, keepdim=True) + 1e-8) | 439 | reward = (reward - reward.mean(dim=1, keepdim=True)) / (reward.std(dim=1, keepdim=True) + 1e-8) |
| 333 | reward = reward.reshape(-1) | 440 | reward = reward.reshape(-1) |
| 334 | - last_reward = torch.zeros_like(rm_scores).scatter_(dim=1, index=eos_indices, src=reward.unsqueeze(1).to(rm_scores.dtype)) | 441 | + token_level_rewards = torch.zeros_like(rm_scores).scatter_(dim=1, index=eos_indices, src=reward.unsqueeze(1).to(rm_scores.dtype)) |
| 335 | - data.batch['token_level_rewards'] = last_reward | 442 | + return token_level_rewards |
| 336 | - return data | ||
| 337 | 443 | ||
| 338 | 444 | ||
| 339 | def reduce_metrics(metrics: dict): | 445 | def reduce_metrics(metrics: dict): |
| @@ -14,6 +14,10 @@ from megatron.training import get_args, initialize_megatron, get_timers, get_tok | |||
| 14 | from megatron.core.pipeline_parallel.schedules import get_forward_backward_func | 14 | from megatron.core.pipeline_parallel.schedules import get_forward_backward_func |
| 15 | from megatron.core import parallel_state as mpu, tensor_parallel | 15 | from megatron.core import parallel_state as mpu, tensor_parallel |
| 16 | from megatron.training.training import append_to_progress_log, build_train_valid_test_data_iterators, print_datetime | 16 | from megatron.training.training import append_to_progress_log, build_train_valid_test_data_iterators, print_datetime |
| 17 | +from megatron.training import get_model | ||
| 18 | +from megatron.training.utils import unwrap_model | ||
| 19 | +from megatron.training.checkpointing import save_checkpoint | ||
| 20 | +from megatron.training.training import num_floating_point_operations | ||
| 17 | from mindspeed_llm.tasks.posttrain.rlxf.training.core_algos import compute_policy_loss, find_first_eos_index, compute_grpo_policy_loss | 21 | from mindspeed_llm.tasks.posttrain.rlxf.training.core_algos import compute_policy_loss, find_first_eos_index, compute_grpo_policy_loss |
| 18 | from mindspeed_llm.tasks.posttrain.rlxf.utils.torch_functional import split_dict_tensor_into_batches | 22 | from mindspeed_llm.tasks.posttrain.rlxf.utils.torch_functional import split_dict_tensor_into_batches |
| 19 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base.megatron.worker import MegatronWorker | 23 | from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base.megatron.worker import MegatronWorker |
| @@ -24,14 +28,11 @@ from mindspeed_llm.tasks.posttrain.rlxf.utils.protocol import DataProto, make_ba | |||
| 24 | from mindspeed_llm.tasks.posttrain.base import BaseTrainer | 28 | from mindspeed_llm.tasks.posttrain.base import BaseTrainer |
| 25 | import mindspeed_llm.tasks.posttrain.rlxf.training.parallel_state as ps | 29 | import mindspeed_llm.tasks.posttrain.rlxf.training.parallel_state as ps |
| 26 | from mindspeed_llm.tasks.inference.module import MegatronModuleForCausalLM | 30 | from mindspeed_llm.tasks.inference.module import MegatronModuleForCausalLM |
| 27 | -from inference import model_provider | ||
| 28 | from mindspeed_llm.tasks.preprocess.decoder_packed_mtf_dataset import \ | 31 | from mindspeed_llm.tasks.preprocess.decoder_packed_mtf_dataset import \ |
| 29 | build_train_valid_test_datasets as build_instruction_dataset | 32 | build_train_valid_test_datasets as build_instruction_dataset |
| 30 | from mindspeed_llm.training.initialize import set_jit_fusion_options | 33 | from mindspeed_llm.training.initialize import set_jit_fusion_options |
| 31 | from mindspeed_llm.training.utils import get_finetune_data_on_this_tp_rank, get_tune_attention_mask | 34 | from mindspeed_llm.training.utils import get_finetune_data_on_this_tp_rank, get_tune_attention_mask |
| 32 | from mindspeed_llm.tasks.posttrain.utils import compute_log_probs, append_to_dict | 35 | from mindspeed_llm.tasks.posttrain.utils import compute_log_probs, append_to_dict |
| 33 | -from megatron.training.checkpointing import save_checkpoint | ||
| 34 | -from megatron.training.training import num_floating_point_operations | ||
| 35 | 36 | ||
| 36 | 37 | ||
| 37 | def train_valid_test_datasets_provider(train_val_test_num_samples): | 38 | def train_valid_test_datasets_provider(train_val_test_num_samples): |
| @@ -202,6 +203,66 @@ class PPOActorInferWorker(BaseTrainer): | |||
| 202 | def __init__(self): | 203 | def __init__(self): |
| 203 | super().__init__() | 204 | super().__init__() |
| 204 | self.count = 0 | 205 | self.count = 0 |
| 206 | + self.keys = None | ||
| 207 | + | ||
| 208 | + def model_provider(self, pre_process=True, post_process=True): | ||
| 209 | + """Builds the inference model. | ||
| 210 | + | ||
| 211 | + If you set the use_mcore_models to True, it will return the mcore GPT model. | ||
| 212 | + | ||
| 213 | + Args: | ||
| 214 | + pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True. | ||
| 215 | + post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True. | ||
| 216 | + | ||
| 217 | + | ||
| 218 | + Returns: | ||
| 219 | + Union[GPTModelInfer, GPTModel]: The returned model | ||
| 220 | + """ | ||
| 221 | + | ||
| 222 | + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, \ | ||
| 223 | + get_gpt_layer_local_spec | ||
| 224 | + from megatron.core.transformer.spec_utils import import_module | ||
| 225 | + from megatron.training import get_args, print_rank_0 | ||
| 226 | + from megatron.training.arguments import core_transformer_config_from_args | ||
| 227 | + from megatron.training.yaml_arguments import core_transformer_config_from_yaml | ||
| 228 | + | ||
| 229 | + from mindspeed_llm.tasks.inference.module import GPTModelInfer | ||
| 230 | + | ||
| 231 | + args = get_args() | ||
| 232 | + use_te = args.transformer_impl == "transformer_engine" | ||
| 233 | + | ||
| 234 | + print_rank_0('building GPT Rollout model ...') | ||
| 235 | + # Experimental loading arguments from yaml | ||
| 236 | + if args.yaml_cfg is not None: | ||
| 237 | + config = core_transformer_config_from_yaml(args, "language_model") | ||
| 238 | + else: | ||
| 239 | + config = core_transformer_config_from_args(args) | ||
| 240 | + | ||
| 241 | + if args.spec is not None: | ||
| 242 | + transformer_layer_spec = import_module(args.spec) | ||
| 243 | + else: | ||
| 244 | + if use_te: | ||
| 245 | + transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(args.num_experts, | ||
| 246 | + args.moe_grouped_gemm) | ||
| 247 | + else: | ||
| 248 | + transformer_layer_spec = get_gpt_layer_local_spec(args.num_experts, args.moe_grouped_gemm) | ||
| 249 | + | ||
| 250 | + model = GPTModelInfer( | ||
| 251 | + config=config, | ||
| 252 | + transformer_layer_spec=transformer_layer_spec, | ||
| 253 | + vocab_size=args.padded_vocab_size, | ||
| 254 | + max_sequence_length=args.max_position_embeddings, | ||
| 255 | + pre_process=pre_process, | ||
| 256 | + post_process=post_process, | ||
| 257 | + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, | ||
| 258 | + parallel_output=False, | ||
| 259 | + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, | ||
| 260 | + position_embedding_type=args.position_embedding_type, | ||
| 261 | + rotary_percent=args.rotary_percent, | ||
| 262 | + seq_len_interpolation_factor=args.rotary_seq_len_interpolation_factor | ||
| 263 | + ) | ||
| 264 | + | ||
| 265 | + return model | ||
| 205 | 266 | ||
| 206 | def initialize(self): | 267 | def initialize(self): |
| 207 | train_valid_test_datasets_provider.is_distributed = True | 268 | train_valid_test_datasets_provider.is_distributed = True |
| @@ -220,9 +281,10 @@ class PPOActorInferWorker(BaseTrainer): | |||
| 220 | self.args.num_layer_list = None | 281 | self.args.num_layer_list = None |
| 221 | self.args.micro_batch_size = 1 | 282 | self.args.micro_batch_size = 1 |
| 222 | self.args.sequence_parallel = False | 283 | self.args.sequence_parallel = False |
| 223 | - self.inf_model = MegatronModuleForCausalLM.from_pretrained( | 284 | + |
| 224 | - model_provider=model_provider, | 285 | + self.args.model = unwrap_model(get_model(self.model_provider, wrap_with_ddp=False)) |
| 225 | - ) | 286 | + self.inf_model = self.args.model[0] |
| 287 | + | ||
| 226 | sync_param_nums(self.inf_model) | 288 | sync_param_nums(self.inf_model) |
| 227 | true_pad_to_multiple_of = self.args.pad_to_multiple_of | 289 | true_pad_to_multiple_of = self.args.pad_to_multiple_of |
| 228 | self.args.pad_to_multiple_of = 1 # we don't want to pad data here | 290 | self.args.pad_to_multiple_of = 1 # we don't want to pad data here |
| @@ -237,11 +299,15 @@ class PPOActorInferWorker(BaseTrainer): | |||
| 237 | print('done with setup ...') | 299 | print('done with setup ...') |
| 238 | self.timers.log(['model-setup', 'train/valid/test-data-iterators-setup'], barrier=True) | 300 | self.timers.log(['model-setup', 'train/valid/test-data-iterators-setup'], barrier=True) |
| 239 | 301 | ||
| 240 | - @staticmethod | 302 | + def get_batch(self, data_iterator): |
| 241 | - def get_batch(data_iterator): | ||
| 242 | """Generate a batch identical to Llama factory""" | 303 | """Generate a batch identical to Llama factory""" |
| 243 | args = get_args() | 304 | args = get_args() |
| 244 | 305 | ||
| 306 | + if self.args.dataset_with_labels: | ||
| 307 | + self.keys = ['input_ids', 'labels'] | ||
| 308 | + else: | ||
| 309 | + self.keys = ['input_ids'] | ||
| 310 | + | ||
| 245 | if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()): | 311 | if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()): |
| 246 | if args.variable_seq_lengths and args.pipeline_model_parallel_size > 2: | 312 | if args.variable_seq_lengths and args.pipeline_model_parallel_size > 2: |
| 247 | tokens, _ = get_finetune_data_on_this_tp_rank(data_iterator) | 313 | tokens, _ = get_finetune_data_on_this_tp_rank(data_iterator) |
| @@ -250,48 +316,91 @@ class PPOActorInferWorker(BaseTrainer): | |||
| 250 | return tokens | 316 | return tokens |
| 251 | 317 | ||
| 252 | # Items and their type. | 318 | # Items and their type. |
| 253 | - keys = ['input_ids'] | 319 | + |
| 254 | data_type = torch.int64 | 320 | data_type = torch.int64 |
| 255 | 321 | ||
| 256 | # Broadcast data. | 322 | # Broadcast data. |
| 257 | - data_b = tensor_parallel.broadcast_data(keys, next(data_iterator), data_type) | 323 | + data_b = tensor_parallel.broadcast_data(self.keys, next(data_iterator), data_type) |
| 258 | 324 | ||
| 259 | # Unpack | 325 | # Unpack |
| 260 | - tokens = data_b.get('input_ids').long() | 326 | + batch = {} |
| 261 | - return tokens | 327 | + for key in self.keys: |
| 328 | + batch[key] = data_b.get(key).long() | ||
| 329 | + | ||
| 330 | + return batch | ||
| 262 | 331 | ||
| 263 | def run_inference(self): | 332 | def run_inference(self): |
| 264 | args = get_args() | 333 | args = get_args() |
| 265 | num_infer_steps = args.global_batch_size // (args.data_parallel_size * args.num_samples_per_step) | 334 | num_infer_steps = args.global_batch_size // (args.data_parallel_size * args.num_samples_per_step) |
| 266 | responses = [] | 335 | responses = [] |
| 267 | idx_list = [] | 336 | idx_list = [] |
| 337 | + label_list = [] | ||
| 268 | idx_list_per_step = [] | 338 | idx_list_per_step = [] |
| 339 | + label_list_per_step = [] | ||
| 340 | + | ||
| 269 | max_new_tokens = args.seq_length - args.max_prompt_length | 341 | max_new_tokens = args.seq_length - args.max_prompt_length |
| 270 | assert max_new_tokens % args.pad_to_multiple_of == 0, "please adjust pad_to_multiple_of so that \ | 342 | assert max_new_tokens % args.pad_to_multiple_of == 0, "please adjust pad_to_multiple_of so that \ |
| 271 | max_new_tokens % args.pad_to_multiple_of == 0" | 343 | max_new_tokens % args.pad_to_multiple_of == 0" |
| 272 | for i in range(num_infer_steps): | 344 | for i in range(num_infer_steps): |
| 273 | for j in range(args.num_samples_per_step): | 345 | for j in range(args.num_samples_per_step): |
| 274 | - tokens = self.get_batch(self.train_data_iterator) | 346 | + batch = self.get_batch(self.train_data_iterator) |
| 347 | + | ||
| 348 | + tokens = batch["input_ids"] | ||
| 275 | tokens_list = tokens.view(-1).cpu().numpy().tolist() | 349 | tokens_list = tokens.view(-1).cpu().numpy().tolist() |
| 350 | + labels = batch.get("labels", None) | ||
| 351 | + | ||
| 352 | + if labels is not None: | ||
| 353 | + labels = labels.view(-1).cpu().numpy().tolist()[::-1] | ||
| 354 | + if -100 in labels: | ||
| 355 | + answer_begin_indices = len(tokens_list) - 1 - labels.index(-100) | ||
| 356 | + answer_list = tokens_list[answer_begin_indices + 1:] | ||
| 357 | + tokens_list = tokens_list[:answer_begin_indices + 1] | ||
| 358 | + else: | ||
| 359 | + answer_list = tokens_list[-1:] | ||
| 360 | + tokens_list = tokens_list[:-1] | ||
| 361 | + | ||
| 362 | + for _ in range(args.n_samples_per_prompt): | ||
| 363 | + label_list_per_step.append(copy.deepcopy(answer_list)) | ||
| 364 | + | ||
| 276 | for _ in range(args.n_samples_per_prompt): | 365 | for _ in range(args.n_samples_per_prompt): |
| 277 | idx_list_per_step.append(copy.deepcopy(tokens_list)) | 366 | idx_list_per_step.append(copy.deepcopy(tokens_list)) |
| 367 | + | ||
| 278 | if args.stage == "ray_online_dpo": | 368 | if args.stage == "ray_online_dpo": |
| 279 | idx_list_per_step.append(copy.deepcopy(tokens_list)) | 369 | idx_list_per_step.append(copy.deepcopy(tokens_list)) |
| 280 | 370 | ||
| 281 | - responses_per_step = self.inf_model.generate(copy.deepcopy(idx_list_per_step), | 371 | + responses_per_step = self.inf_model.generate( |
| 282 | - max_new_tokens=max_new_tokens, | 372 | + copy.deepcopy(idx_list_per_step), |
| 283 | - detokenize=False, broadcast=False, do_sample=args.do_sample) | 373 | + max_new_tokens=max_new_tokens, |
| 374 | + temperature=args.temperature, | ||
| 375 | + do_sample=args.do_sample, | ||
| 376 | + detokenize=False, | ||
| 377 | + broadcast=False, | ||
| 378 | + truncate=True | ||
| 379 | + ) | ||
| 380 | + | ||
| 284 | if not isinstance(responses_per_step, list): | 381 | if not isinstance(responses_per_step, list): |
| 285 | responses_per_step = [responses_per_step] | 382 | responses_per_step = [responses_per_step] |
| 286 | 383 | ||
| 287 | responses.extend(responses_per_step) | 384 | responses.extend(responses_per_step) |
| 288 | idx_list.extend(idx_list_per_step) | 385 | idx_list.extend(idx_list_per_step) |
| 386 | + label_list.extend(label_list_per_step) | ||
| 289 | idx_list_per_step = [] | 387 | idx_list_per_step = [] |
| 388 | + label_list_per_step = [] | ||
| 290 | 389 | ||
| 291 | - responses = [response.cpu().numpy().tolist() for response in responses] | 390 | + responses_ori_length, responses_pad_length = pad_to_tensor_dict( |
| 391 | + responses, | ||
| 392 | + pad_multi_of=args.pad_to_multiple_of | ||
| 393 | + ) | ||
| 394 | + prompts_ori_length, prompts_pad_length = pad_to_tensor_dict( | ||
| 395 | + idx_list, "left", | ||
| 396 | + pad_multi_of=args.pad_to_multiple_of | ||
| 397 | + ) | ||
| 292 | 398 | ||
| 293 | - responses_ori_length, responses_pad_length = pad_to_tensor_dict(responses, pad_multi_of=args.pad_to_multiple_of) | 399 | + if label_list: |
| 294 | - prompts_ori_length, prompts_pad_length = pad_to_tensor_dict(idx_list, "left", pad_multi_of=args.pad_to_multiple_of) | 400 | + pad_to_tensor_dict( |
| 401 | + label_list, | ||
| 402 | + pad_multi_of=args.pad_to_multiple_of | ||
| 403 | + ) | ||
| 295 | 404 | ||
| 296 | input_ids = [prompt + response for prompt, response in zip(idx_list, responses)] | 405 | input_ids = [prompt + response for prompt, response in zip(idx_list, responses)] |
| 297 | 406 | ||
| @@ -304,15 +413,32 @@ class PPOActorInferWorker(BaseTrainer): | |||
| 304 | else: | 413 | else: |
| 305 | batch_size = args.global_batch_size // args.data_parallel_size * args.n_samples_per_prompt | 414 | batch_size = args.global_batch_size // args.data_parallel_size * args.n_samples_per_prompt |
| 306 | 415 | ||
| 416 | + if label_list: | ||
| 417 | + batch = TensorDict( | ||
| 418 | + { | ||
| 419 | + "prompts": idx_list, | ||
| 420 | + "labels": label_list, | ||
| 421 | + "responses": responses, | ||
| 422 | + "input_ids": input_ids, | ||
| 423 | + "attention_mask": attention_mask, | ||
| 424 | + "position_ids": position_ids, | ||
| 425 | + "responses_ori_length": responses_ori_length | ||
| 426 | + }, | ||
| 427 | + batch_size=batch_size | ||
| 428 | + ) | ||
| 429 | + else: | ||
| 430 | + batch = TensorDict( | ||
| 431 | + { | ||
| 432 | + "prompts": idx_list, | ||
| 433 | + "responses": responses, | ||
| 434 | + "input_ids": input_ids, | ||
| 435 | + "attention_mask": attention_mask, | ||
| 436 | + "position_ids": position_ids, | ||
| 437 | + "responses_ori_length": responses_ori_length | ||
| 438 | + }, | ||
| 439 | + batch_size=batch_size | ||
| 440 | + ) | ||
| 307 | 441 | ||
| 308 | - batch = TensorDict({ | ||
| 309 | - "prompts": idx_list, | ||
| 310 | - "responses": responses, | ||
| 311 | - "input_ids": input_ids, | ||
| 312 | - "attention_mask": attention_mask, | ||
| 313 | - "position_ids": position_ids, | ||
| 314 | - }, | ||
| 315 | - batch_size=batch_size) | ||
| 316 | return DataProto(batch=batch) | 442 | return DataProto(batch=batch) |
| 317 | 443 | ||
| 318 | def loss_func(self, input_tensor, output_tensor): | 444 | def loss_func(self, input_tensor, output_tensor): |
| @@ -99,7 +99,7 @@ class MegatronPPORM(ORMTrainer): | |||
| 99 | if self.args.load is not None or self.args.pretrained_checkpoint is not None: | 99 | if self.args.load is not None or self.args.pretrained_checkpoint is not None: |
| 100 | self.timers('load-checkpoint', log_level=0).start(barrier=True) | 100 | self.timers('load-checkpoint', log_level=0).start(barrier=True) |
| 101 | self.args.iteration, self.args.num_floating_point_operations_so_far = load_checkpoint( | 101 | self.args.iteration, self.args.num_floating_point_operations_so_far = load_checkpoint( |
| 102 | - model, None, None) | 102 | + model, None, None, strict=False) |
| 103 | self.timers('load-checkpoint').stop(barrier=True) | 103 | self.timers('load-checkpoint').stop(barrier=True) |
| 104 | self.timers.log(['load-checkpoint']) | 104 | self.timers.log(['load-checkpoint']) |
| 105 | else: | 105 | else: |
| @@ -705,6 +705,10 @@ def _add_rl_args(parser): | |||
| 705 | default=False, | 705 | default=False, |
| 706 | help="Path to the reference model used for the PPO training." | 706 | help="Path to the reference model used for the PPO training." |
| 707 | ) | 707 | ) |
| 708 | + group.add_argument( | ||
| 709 | + "--verifier", action='store_true', | ||
| 710 | + help="Enable verifier in cal scores." | ||
| 711 | + ) | ||
| 708 | group.add_argument( | 712 | group.add_argument( |
| 709 | '--kl-coef', | 713 | '--kl-coef', |
| 710 | default=0.3, | 714 | default=0.3, |
| @@ -786,7 +790,7 @@ def _add_training_args(parser): | |||
| 786 | help='Setting jit compile mode to True') | 790 | help='Setting jit compile mode to True') |
| 787 | group.add_argument('--prompt-type', type=str, default=None, | 791 | group.add_argument('--prompt-type', type=str, default=None, |
| 788 | choices=['default', 'empty', 'trl', 'chatglm2', 'chatglm3', 'chatglm3_system', 'glm4', 'chatml', | 792 | choices=['default', 'empty', 'trl', 'chatglm2', 'chatglm3', 'chatglm3_system', 'glm4', 'chatml', |
| 789 | - 'chatml_de', 'qwen', 'llama3', 'llama2', 'mistral', 'mixtral', 'gemma', 'alpaca', | 793 | + 'chatml_de', 'qwen', 'qwen_r1', "qwen_math_r1", 'llama3', 'llama2', 'mistral', 'mixtral', 'gemma', 'alpaca', |
| 790 | 'deepseek2', 'deepseek2-lite', 'minicpm3', 'cpm', 'baichuan2', 'deepseek3', 'intern2'], | 794 | 'deepseek2', 'deepseek2-lite', 'minicpm3', 'cpm', 'baichuan2', 'deepseek3', 'intern2'], |
| 791 | help='Which template to use for constructing prompts in training/inference.' 'e.g., "qwen"') | 795 | help='Which template to use for constructing prompts in training/inference.' 'e.g., "qwen"') |
| 792 | group.add_argument('--prompt-type-path', type=str, default=TEMPLATES_DIR, | 796 | group.add_argument('--prompt-type-path', type=str, default=TEMPLATES_DIR, |
| @@ -901,6 +905,12 @@ def _add_dataset_args(parser): | |||
| 901 | default=False, | 905 | default=False, |
| 902 | help='if no shared storage, set it' | 906 | help='if no shared storage, set it' |
| 903 | ) | 907 | ) |
| 908 | + group.add_argument('--dataset-with-labels', | ||
| 909 | + action='store_true', | ||
| 910 | + default=False, | ||
| 911 | + help='Whether the labels included in dataset.' | ||
| 912 | + ) | ||
| 913 | + | ||
| 904 | return parser | 914 | return parser |
| 905 | 915 | ||
| 906 | 916 | ||
| @@ -112,7 +112,7 @@ def add_data_args(parser): | |||
| 112 | # LlamaFactory | 112 | # LlamaFactory |
| 113 | group.add_argument('--prompt-type', type=str, default=None, | 113 | group.add_argument('--prompt-type', type=str, default=None, |
| 114 | choices=['default', 'empty', 'trl', 'chatglm2', 'chatglm3', 'chatglm3_system', 'glm4', 'chatml', | 114 | choices=['default', 'empty', 'trl', 'chatglm2', 'chatglm3', 'chatglm3_system', 'glm4', 'chatml', |
| 115 | - 'chatml_de', 'qwen', 'llama3', 'llama2', 'mistral', 'mixtral', 'gemma', 'alpaca', | 115 | + 'chatml_de', 'qwen', 'qwen_r1', "qwen_math_r1", 'llama3', 'llama2', 'mistral', 'mixtral', 'gemma', 'alpaca', |
| 116 | 'deepseek2', 'deepseek2-lite', 'cpm', 'baichuan2', 'deepseek3', 'intern2'], | 116 | 'deepseek2', 'deepseek2-lite', 'cpm', 'baichuan2', 'deepseek3', 'intern2'], |
| 117 | help='Which template to use for constructing prompts in training.' | 117 | help='Which template to use for constructing prompts in training.' |
| 118 | 'e.g., "qwen"') | 118 | 'e.g., "qwen"') |
P 图片需要按照Qwen2.5-7B-instruct重新更新一下   | |||
| @@ -18,3 +18,5 @@ tensordict==0.1.2 | |||
| 18 | hydra-core==1.3.2 | 18 | hydra-core==1.3.2 |
| 19 | codetiming | 19 | codetiming |
| 20 | bitsandbytes-npu-beta==0.45.2 | 20 | bitsandbytes-npu-beta==0.45.2 |
| 21 | +word2number | ||
| 22 | +latex2sympy2 | ||
| @@ -16,6 +16,7 @@ | |||
| 16 | 0.5493276, | 16 | 0.5493276, |
| 17 | 0.5800903 | 17 | 0.5800903 |
| 18 | ], | 18 | ], |
| 19 | + "throughput margin": 0.1, | ||
| 19 | "throughput": [ | 20 | "throughput": [ |
| 20 | 4.6, | 21 | 4.6, |
| 21 | 10.9, | 22 | 10.9, |
| @@ -16,6 +16,7 @@ | |||
| 16 | 0.7187351, | 16 | 0.7187351, |
| 17 | 0.659517 | 17 | 0.659517 |
| 18 | ], | 18 | ], |
| 19 | + "throughput margin": 0.1, | ||
| 19 | "throughput": [ | 20 | "throughput": [ |
| 20 | 3.5, | 21 | 3.5, |
| 21 | 5.5, | 22 | 5.5, |
| @@ -16,6 +16,7 @@ | |||
| 16 | 5.516537, | 16 | 5.516537, |
| 17 | 5.494636 | 17 | 5.494636 |
| 18 | ], | 18 | ], |
| 19 | + "throughput margin": 0.1, | ||
| 19 | "throughput": [ | 20 | "throughput": [ |
| 20 | 15.9, | 21 | 15.9, |
| 21 | 44.6, | 22 | 44.6, |
| @@ -16,6 +16,7 @@ | |||
| 16 | 5.696367, | 16 | 5.696367, |
| 17 | 5.602035 | 17 | 5.602035 |
| 18 | ], | 18 | ], |
| 19 | + "throughput margin": 0.1, | ||
| 19 | "throughput": [ | 20 | "throughput": [ |
| 20 | 10.4, | 21 | 10.4, |
| 21 | 101.7, | 22 | 101.7, |
| @@ -1,4 +1,5 @@ | |||
| 1 | { | 1 | { |
| 2 | + "lm loss margin": 0.03, | ||
| 2 | "lm loss": [ | 3 | "lm loss": [ |
| 3 | 0.6931, | 4 | 0.6931, |
| 4 | 0.6925, | 5 | 0.6925, |
| @@ -1,4 +1,5 @@ | |||
| 1 | { | 1 | { |
| 2 | + "lm loss margin": 0.01, | ||
| 2 | "lm loss": [ | 3 | "lm loss": [ |
| 3 | 0.0214, | 4 | 0.0214, |
| 4 | 0.0412, | 5 | 0.0412, |
| @@ -7,14 +8,14 @@ | |||
| 7 | 0.0176, | 8 | 0.0176, |
| 8 | 0.0005, | 9 | 0.0005, |
| 9 | 0.0685, | 10 | 0.0685, |
| 10 | - 0.1176, | 11 | + 0.0608, |
| 11 | - 0.0985, | 12 | + 0.0895, |
| 12 | 0.2175, | 13 | 0.2175, |
| 13 | 0.1695, | 14 | 0.1695, |
| 14 | 0.2064, | 15 | 0.2064, |
| 15 | 0.0822, | 16 | 0.0822, |
| 16 | 0.0318, | 17 | 0.0318, |
| 17 | - 0.0557 | 18 | + 0.0624 |
| 18 | ], | 19 | ], |
| 19 | "throughput": [ | 20 | "throughput": [ |
| 20 | ], | 21 | ], |
| @@ -16,6 +16,7 @@ | |||
| 16 | 1.560994, | 16 | 1.560994, |
| 17 | 1.477175 | 17 | 1.477175 |
| 18 | ], | 18 | ], |
| 19 | + "throughput margin": 0.1, | ||
| 19 | "throughput": [ | 20 | "throughput": [ |
| 20 | 4.2, | 21 | 4.2, |
| 21 | 19.3, | 22 | 19.3, |
要和R1的复现参数格式协调