【part 2】支持trl ppo，增加PPOEngine和PPOTrainer-MindSpeed-LLM-AtomGit

已合并

【part 2】支持trl ppo，增加PPOEngine和PPOTrainer #2094

AtomGit-Bot创建于 2024年12月24日从refs/pull/2094/head 合入到 master

【part 2】支持trl ppo，增加PPOEngine和PPOTrainer #2094

已合并

AtomGit-Bot创建于 2024年12月24日

从refs/pull/2094/head合入到master

共 11 个文件变更+1104-8

Adocs/features/trl_ppo.md+124-0

@@ -0,0 +1,124 @@
	1	+# Trl PPO算法
	2	+Trl PPO实现参考[huggingface/trl](https://github.com/huggingface/trl)库的PPO进行实现。
	3	+PPO（Proximal Policy Optimization）是一种强化学习算法，它通过引入奖励信号来调整模型的行为，使模型生成的内容更符合人类的偏好。位于RLHF整个系统中的第3阶段。
	4	+
	5	+![](../../sources/images/trl_ppo/ppo_in_rlhf.png)
	6	+
	7	+Trl PPO算法涉及4个模型：actor模型、reference模型、reward模型、critic模型。其中actor模型和critic模型训练中会更新。actor模型和reference模型初始由RLHF第一阶段sft训练产生，reward模型和critic模型初始由RLHF第二阶段reward训练产生。
	8	+Trl PPO输入数据仅需要prompt，response由actor模型在线产生。
	9	+
	10	+## 使用说明
	11	+
	12	+### 数据预处理示例
	13	+
	14	+```shell
	15	+python ./preprocess_data.py \
	16	+ --input ./dataset/train-00000-of-00001-a09b74b3ef9c3b56.parquet \
	17	+ --tokenizer-name-or-path ./model_from_hf/Llama-3.2-1B-Instruct \
	18	+ --output-prefix ./finetune_dataset/llama-3.2-1b-trl/alpaca \
	19	+ --workers 16 \
	20	+ --log-interval 1000 \
	21	+ --tokenizer-type PretrainedFromHF \
	22	+ --handler-name PPOAlpacaStyleInstructionHandler \
	23	+ --prompt-type empty
	24	+```
	25	+
	26	+
	27	+### 训练脚本示例
	28	+
	29	+目前仓上已包含 Llama3.2-1B trl_ppo训练脚本。以 Llama3.2-1B 为例，训练脚本参照：examples/mcore/llama32/trl_ppo_llama32_1b.sh
	30	+
	31	+相较预训练，Trl PPO训练需要额外配置以下参数：
	32	+
	33	+- `--stage trl_ppo`
	34	+
	35	+ 指定进行PPO训练
	36	+
	37	+- `--ref-model`
	38	+
	39	+ 必选，reference模型和actor模型的路径
	40	+
	41	+- `--reward-model`
	42	+
	43	+ 必选，reward模型和critic模型的路径
	44	+
	45	+- `--kl-coef`
	46	+
	47	+ 可选，计算优势函数值时的KL散度系数，默认0.3
	48	+
	49	+- `--gamma`
	50	+
	51	+ 可选，计算优势函数值时的Discount factor，默认1.0
	52	+
	53	+- `--lam`
	54	+
	55	+ 可选，计算优势函数值时的Lambda value，默认0.95
	56	+
	57	+- `--clip-ratio`
	58	+
	59	+ 可选，actor loss的裁剪幅度，默认0.2
	60	+
	61	+- `--cliprange-value`
	62	+
	63	+ 可选，critic loss的裁剪幅度，默认0.2
	64	+
	65	+- `--max-length`
	66	+
	67	+ 可选，prompt+response的最长长度，默认256
	68	+
	69	+- `--max-new-tokens`
	70	+
	71	+ 可选，最大response生成长度，默认128
	72	+
	73	+- `--do-sample`
	74	+
	75	+ 可选，生成样本时采用采样选择或贪婪选择token，默认贪婪选择
	76	+
	77	+## Trl PPO 流程
	78	+
	79	+
	80	+## 日志说明
	81	+单个iteration 输出两个模型的训练日志，第一行是actor模型信息，第二行是critic模型信息
	82	+
	83	+actor 日志说明:
	84	+- 1.pg_loss: policy gradient loss即是actor模型训练loss
	85	+- 2.abs_pg_loss: absolute policy gradient loss 即取每个token的policy gradient loss的绝对值来求平均，观测时使用。
	86	+- 3.ppo_ratio: 新policy的action分布与旧policy的action分布的比值
	87	+- 4.ppo_ratio_clamped: 新policy的action分布与旧policy的action分布的截断比值
	88	+
	89	+critic 日志说明:
	90	+- 1.vf_loss: value function loss 即是 critic 模型训练loss
	91	+
	92	+## 目前局限
	93	+
	94	+- 1、4模型只支持同类同参数模型
	95	+- 2、只支持tp、pp、dp
	96	+- 3、4模型tp、pp、dp需一致
	97	+- 4、梯度累积、rollout_batch_size、ppo_epoch未支持
	98	+- 5、更多观测指标的日志未支持
	99	+- 6、val/test数据集评估未支持
	100	+
	101	+该版本为demo版本，待进一步功能完善。
	102	+
	103	+## MindSpeed-LLM 与 <a href="https://github.com/huggingface/trl">trl</a> loss 对比：
	104	+
	105	+数据集：alpaca
	106	+
	107	+训练参数设置：
	108	+```
	109	+per_device_train_batch_size: 1
	110	+gradient_accumulation_steps: 1
	111	+learning_rate: 1.0e-7
	112	+lr_scheduler_type: constant
	113	+fp32: true
	114	+```
	115	+
	116	+loss对比方法：
	117	+- 由于trl实现中，对advantages做了masked_whiten操作，导致pg_loss均值在0附近，故观测abs_pg_loss进行对比。
	118	+
	119	+
	120	+![alt text](../../sources/images/trl_ppo/llama32_trl_ppo_loss.png)
	121	+
	122	+## 参考文献
	123	+
	124	+- [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347)

Aexamples/mcore/llama32/trl_ppo_llama32_1b_ptd.sh+113-0

@@ -0,0 +1,113 @@
	1	+#!/bin/bash
	2	+
	3	+export CUDA_DEVICE_MAX_CONNECTIONS=1
	4	+
	5	+GPUS_PER_NODE=4
	6	+MASTER_ADDR=localhost
	7	+MASTER_PORT=6001
	8	+NNODES=1
	9	+NODE_RANK=0
	10	+WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
	11	+
	12	+DATA_PATH="your data path"
	13	+SFT_LOAD_DIR="your actor/ref model ckpt path"
	14	+REWARD_LOAD_DIR="your reward/critic model ckpt path"
	15	+TOKENIZER_MODEL="your tokenizer path"
	16	+CKPT_SAVE_DIR="your model save ckpt path"
	17	+TP=2
	18	+PP=2
	19	+
	20	+DISTRIBUTED_ARGS="
	21	+ --nproc_per_node $GPUS_PER_NODE \
	22	+ --nnodes $NNODES \
	23	+ --node_rank $NODE_RANK \
	24	+ --master_addr $MASTER_ADDR \
	25	+ --master_port $MASTER_PORT
	26	+"
	27	+
	28	+GPT_ARGS="
	29	+ --tensor-model-parallel-size ${TP} \
	30	+ --pipeline-model-parallel-size ${PP} \
	31	+ --use-fused-swiglu \
	32	+ --use-mcore-models \
	33	+ --micro-batch-size 4 \
	34	+ --global-batch-size 4 \
	35	+ --use-mc2 \
	36	+ --use-rotary-position-embeddings \
	37	+ --rope-scaling-type llama3 \
	38	+ --rope-scaling-factor 32.0 \
	39	+ --low-freq-factor 1.0 \
	40	+ --high-freq-factor 4.0 \
	41	+ --original-max-position-embeddings 8192 \
	42	+ --tokenizer-type PretrainedFromHF \
	43	+ --tokenizer-name-or-path ${TOKENIZER_MODEL} \
	44	+ --num-layers 16 \
	45	+ --hidden-size 2048 \
	46	+ --ffn-hidden-size 8192 \
	47	+ --num-attention-heads 32 \
	48	+ --group-query-attention \
	49	+ --num-query-groups 8 \
	50	+ --seq-length 131072 \
	51	+ --max-position-embeddings 131072 \
	52	+ --make-vocab-size-divisible-by 1 \
	53	+ --padded-vocab-size 128256 \
	54	+ --disable-bias-linear \
	55	+ --attention-dropout 0.0 \
	56	+ --init-method-std 0.01 \
	57	+ --hidden-dropout 0.0 \
	58	+ --position-embedding-type rope \
	59	+ --rotary-base 500000 \
	60	+ --normalization RMSNorm \
	61	+ --norm-epsilon 1e-5 \
	62	+ --swiglu \
	63	+ --lr 1e-7 \
	64	+ --train-iters 4 \
	65	+ --lr-decay-style constant \
	66	+ --min-lr 1e-7 \
	67	+ --weight-decay 0.0 \
	68	+ --lr-warmup-fraction 0.00 \
	69	+ --clip-grad 1.0 \
	70	+ --adam-beta1 0.9 \
	71	+ --adam-beta2 0.95 \
	72	+ --initial-loss-scale 4096 \
	73	+ --no-gradient-accumulation-fusion \
	74	+ --no-load-optim \
	75	+ --no-load-rng \
	76	+ --finetune \
	77	+ --is-instruction-dataset \
	78	+ --variable-seq-lengths \
	79	+ --attention-softmax-in-fp32 \
	80	+ --no-masked-softmax-fusion \
	81	+ --seq-length 8192 \
	82	+"
	83	+
	84	+DATA_ARGS="
	85	+ --data-path $DATA_PATH \
	86	+ --split 100,0,0
	87	+"
	88	+
	89	+OUTPUT_ARGS="
	90	+ --log-interval 1 \
	91	+ --save-interval 2000 \
	92	+ --eval-interval 2000 \
	93	+ --eval-iters 1 \
	94	+"
	95	+
	96	+RL_ARGS="
	97	+ --stage trl_ppo \
	98	+ --max-new-tokens 256 \
	99	+ --max-length 512
	100	+"
	101	+
	102	+torchrun $DISTRIBUTED_ARGS posttrain_gpt.py \
	103	+ $GPT_ARGS \
	104	+ $OUTPUT_ARGS \
	105	+ $DATA_ARGS \
	106	+ $RL_ARGS \
	107	+ $PROFILE_ARGS \
	108	+ --tokenizer-not-use-fast \
	109	+ --distributed-backend nccl \
	110	+ --ref-model ${SFT_LOAD_DIR} \
	111	+ --reward-model ${REWARD_LOAD_DIR} \
	112	+ --save ${CKPT_SAVE_DIR} \
	113	+ \| tee logs/trl_ppo_llama32_1b.log

Mmindspeed_llm/tasks/posttrain/launcher.py+6-3

	@@ -11,6 +11,7 @@ from mindspeed_llm.tasks.posttrain.dpo import DPOTrainer
11	from mindspeed_llm.tasks.posttrain.orm import ORMTrainer	11	from mindspeed_llm.tasks.posttrain.orm import ORMTrainer
12	from mindspeed_llm.tasks.posttrain.prm import PRMTrainer	12	from mindspeed_llm.tasks.posttrain.prm import PRMTrainer
13	from mindspeed_llm.tasks.posttrain.dpo import SimPOTrainer	13	from mindspeed_llm.tasks.posttrain.dpo import SimPOTrainer
		14	+from mindspeed_llm.tasks.posttrain.trl_ppo import TrlPPOTrainer
14		15
15	logger = logging.getLogger(__name__)	16	logger = logging.getLogger(__name__)
16		17
	@@ -18,7 +19,7 @@ logger = logging.getLogger(__name__)
18	def get_trainer(stage):	19	def get_trainer(stage):
19	"""	20	"""
20	Factory function to select the appropriate trainer based on the 'stage' argument.	21	Factory function to select the appropriate trainer based on the 'stage' argument.
21	-	22	+
22	:param stage: A string representing the stage of the training.	23	:param stage: A string representing the stage of the training.
23	:return: An instance of the appropriate trainer class.	24	:return: An instance of the appropriate trainer class.
24	"""	25	"""
	@@ -36,6 +37,8 @@ def get_trainer(stage):
36	return RayPPOTrainer	37	return RayPPOTrainer
37	elif stage == "ray_online_dpo":	38	elif stage == "ray_online_dpo":
38	return RayOnlineDPOTrainer	39	return RayOnlineDPOTrainer
		40	+ elif stage == "trl_ppo":
		41	+ return TrlPPOTrainer()
39	elif stage == "ray_grpo":	42	elif stage == "ray_grpo":
40	return RayGRPOTrainer	43	return RayGRPOTrainer
41	else:	44	else:
	@@ -49,11 +52,11 @@ class AutoTrainer:
49	It chooses the appropriate trainer (e.g., SFTTrainer, DPOTrainer, ORMTrainer...)	52	It chooses the appropriate trainer (e.g., SFTTrainer, DPOTrainer, ORMTrainer...)
50	based on the 'stage' argument.	53	based on the 'stage' argument.
51	"""	54	"""
52	-	55	+
53	def __init__(self):	56	def __init__(self):
54	"""	57	"""
55	Initializes the AutoTrainer.	58	Initializes the AutoTrainer.
56	-	59	+
57	- Initializes the training system.	60	- Initializes the training system.
58	- Retrieves the 'stage' argument.	61	- Retrieves the 'stage' argument.
59	- Uses the 'stage' to select the correct trainer.	62	- Uses the 'stage' to select the correct trainer.

Mmindspeed_llm/tasks/posttrain/rlxf/utils/torch_functional.py+4-3

	@@ -46,9 +46,10 @@ def masked_var(values, mask, unbiased=True):
46	raise ValueError("At least one element in the mask has to be 1.")	46	raise ValueError("At least one element in the mask has to be 1.")
47	# note that if mask_sum == 1, then there is a division by zero issue	47	# note that if mask_sum == 1, then there is a division by zero issue
48	# to avoid it you just need to use a larger minibatch_size	48	# to avoid it you just need to use a larger minibatch_size
49	- if mask_sum == 1:	49	+ elif mask_sum == 1:
50	- raise ValueError("The sum of the mask is one, which can cause a division by zero.")	50	+ bessel_correction = mask_sum
51	- bessel_correction = mask_sum / (mask_sum - 1)	51	+ else:
		52	+ bessel_correction = mask_sum / (mask_sum - 1)
52	variance = variance * bessel_correction	53	variance = variance * bessel_correction
53	return variance	54	return variance
54		55

Amindspeed_llm/tasks/posttrain/trl_ppo/TrlPPOEngine.py+291-0

@@ -0,0 +1,291 @@
	1	+# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.
	2	+import dataclasses
	3	+from functools import partial
	4	+import torch
	5	+import torch.nn.functional as F
	6	+
	7	+from megatron.core import mpu
	8	+from megatron.core.enums import ModelType
	9	+from megatron.core.optimizer import get_megatron_optimizer, OptimizerConfig
	10	+from megatron.training import get_model
	11	+from megatron.training.checkpointing import load_checkpoint
	12	+from megatron.training.utils import (
	13	+ average_losses_across_data_parallel_group,
	14	+ print_rank_0,
	15	+ unwrap_model)
	16	+from megatron.training.global_vars import (
	17	+ get_args,
	18	+ get_timers,
	19	+ get_tokenizer)
	20	+from megatron.training.training import (
	21	+ get_optimizer_param_scheduler,
	22	+ build_train_valid_test_data_iterators
	23	+)
	24	+from mindspeed_llm.tasks.posttrain.rlxf.utils.torch_functional import masked_mean, masked_whiten
	25	+from mindspeed_llm.training.utils import get_tune_attention_mask
	26	+from mindspeed_llm.tasks.posttrain.utils import train_valid_test_datasets_provider
	27	+from mindspeed_llm.tasks.posttrain.trl_ppo.actor_model import ActorModel
	28	+from mindspeed_llm.tasks.posttrain.trl_ppo.utils import model_provider
	29	+
	30	+
	31	+class TrlPPOEngine():
	32	+ def __init__(self):
	33	+ self.actor_model = None
	34	+ self.ref_model = None
	35	+ self.reward_model = None
	36	+ self.critic_model = None
	37	+ self.actor_optimizer = None
	38	+ self.critic_optimizer = None
	39	+ self.actor_opt_param_scheduler = None
	40	+ self.critic_opt_param_scheduler = None
	41	+ self.tokenizer = None
	42	+ self.model_type = ModelType.encoder_or_decoder
	43	+ self.process_non_loss_data_func = None
	44	+ self.train_data_iterator, self.valid_data_iterator, self.test_data_iterator = None, None, None
	45	+
	46	+ @staticmethod
	47	+ def actor_model_provider(pre_process=True, post_process=True):
	48	+ return model_provider(is_reward_model=False, pre_process=pre_process, post_process=post_process)
	49	+
	50	+ @staticmethod
	51	+ def critic_model_provider(pre_process=True, post_process=True):
	52	+ return model_provider(is_reward_model=True, pre_process=pre_process, post_process=post_process)
	53	+
	54	+ def initialize(self):
	55	+ """
	56	+ This function will run the followings in the order provided:
	57	+ 1) setup model, optimizer and lr schedule using the model_provider.
	58	+ 2) call train_val_test_data_provider to get train/val/test datasets.
	59	+ """
	60	+
	61	+ train_valid_test_datasets_provider.is_distributed = True
	62	+
	63	+ self.actor_model = ActorModel()
	64	+
	65	+ self.actor_model.model, self.actor_optimizer, self.actor_opt_param_scheduler = \
	66	+ self.setup_model_and_optimizer(self.actor_model_provider, self.model_type, load_arg='ref_model')
	67	+
	68	+ self.critic_model, self.critic_optimizer, self.critic_opt_param_scheduler = \
	69	+ self.setup_model_and_optimizer(self.critic_model_provider, self.model_type, load_arg='reward_model')
	70	+
	71	+ self.ref_model = self.setup_model(self.actor_model_provider, self.model_type, load_arg='ref_model')
	72	+ self.reward_model = self.setup_model(self.critic_model_provider, self.model_type, load_arg='reward_model')
	73	+
	74	+ self.train_data_iterator, self.valid_data_iterator, self.test_data_iterator \
	75	+ = build_train_valid_test_data_iterators(train_valid_test_datasets_provider)
	76	+ self.tokenizer = get_tokenizer().tokenizer
	77	+
	78	+ def forward_only_function(self, input_data, model):
	79	+ tokens = input_data["padded_query_responses"].clone()
	80	+ position_ids = input_data["position_ids"]
	81	+ attention_mask_1d = input_data["attention_mask"].to(bool)
	82	+ attention_mask = get_tune_attention_mask(attention_mask_1d)
	83	+ logits = model(tokens, position_ids, attention_mask)
	84	+
	85	+ def loss_func(logits: torch.Tensor):
	86	+ return logits, {'logits': logits}
	87	+
	88	+ return logits, partial(loss_func)
	89	+
	90	+ def get_actor_and_critic_train_data(self, data):
	91	+
	92	+ if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()):
	93	+ return None, None, None, None, None
	94	+
	95	+ tokens = data['padded_query_responses']
	96	+ position_ids = data["position_ids"]
	97	+ attention_mask_1d = data["attention_mask"].to(bool)
	98	+ attention_mask = get_tune_attention_mask(attention_mask_1d)
	99	+ return tokens, None, None, attention_mask, position_ids
	100	+
	101	+ def actor_loss_func(self, data, parallel_logits):
	102	+ args = get_args()
	103	+ context_length = data['context_length']
	104	+ query_responses = data['padded_query_responses']
	105	+ advantages = data["advantages"]
	106	+ prev_log_probs = data["actor_logprobs"]
	107	+ padding_mask = data["padding_mask"]
	108	+
	109	+ curr_log_probs = parallel_logits[:, context_length - 1:-1].to(torch.float32)
	110	+ curr_log_probs = F.log_softmax(curr_log_probs, dim=-1)
	111	+ tokens_indices = torch.unsqueeze(query_responses[:, context_length:], 2)
	112	+ curr_log_probs = torch.gather(curr_log_probs, 2, index=tokens_indices).squeeze(2)
	113	+
	114	+ # Calculate clipped PPO surrogate loss function.
	115	+ ratios = (curr_log_probs - prev_log_probs).exp()
	116	+ ratios_clamped = ratios.clamp(1.0 - args.clip_ratio, 1.0 + args.clip_ratio)
	117	+
	118	+ loss1 = -advantages * ratios
	119	+ loss2 = -advantages * ratios_clamped
	120	+ actor_loss = masked_mean(torch.max(loss1, loss2), ~padding_mask)
	121	+
	122	+ with torch.no_grad():
	123	+ ppo_ratio = masked_mean(ratios.detach(), ~padding_mask)
	124	+ ppo_ratio_clamped = masked_mean(ratios_clamped.detach(), ~padding_mask)
	125	+
	126	+ abs_actor_loss = masked_mean(torch.max(torch.abs(loss1), torch.abs(loss2)), ~padding_mask)
	127	+ reduced_abs_actor_loss = average_losses_across_data_parallel_group([abs_actor_loss])
	128	+
	129	+ reduced_actor_loss = average_losses_across_data_parallel_group([actor_loss])
	130	+
	131	+ return (
	132	+ actor_loss,
	133	+ {
	134	+ "pg_loss": reduced_actor_loss,
	135	+ "abs_pg_loss": reduced_abs_actor_loss,
	136	+ "ppo_ratio": ppo_ratio,
	137	+ "ppo_ratio_clamped": ppo_ratio_clamped,
	138	+ },
	139	+ )
	140	+
	141	+ def get_actor_forward_output_and_loss_func(self):
	142	+
	143	+ def fwd_output_and_loss_func(data, model):
	144	+ tokens, labels, loss_mask, attention_mask, position_ids = self.get_actor_and_critic_train_data(data)
	145	+ parallel_logits = model(input_ids=tokens, position_ids=position_ids, attention_mask=attention_mask)
	146	+
	147	+ return parallel_logits, partial(self.actor_loss_func, data)
	148	+
	149	+ return fwd_output_and_loss_func
	150	+
	151	+
	152	+ def critic_loss_func(self, data, curr_values):
	153	+ args = get_args()
	154	+ context_length = data['context_length']
	155	+ returns = data["returns"]
	156	+ prev_values = data["values"]
	157	+ padding_mask_p1 = data["padding_mask_p1"]
	158	+
	159	+ curr_values = curr_values.squeeze(-1)[:, context_length - 1:-1].to(torch.float32).contiguous()
	160	+ curr_values = torch.masked_fill(curr_values, padding_mask_p1, 0)
	161	+ curr_values_clipped = torch.clamp(
	162	+ curr_values,
	163	+ prev_values - args.cliprange_value,
	164	+ prev_values + args.cliprange_value,
	165	+ )
	166	+ vf_losses1 = torch.square(curr_values - returns)
	167	+ vf_losses2 = torch.square(curr_values_clipped - returns)
	168	+
	169	+ # Critic loss
	170	+ loss = 0.5 * masked_mean(torch.max(vf_losses1, vf_losses2), ~padding_mask_p1)
	171	+ reduced_loss = average_losses_across_data_parallel_group([loss])
	172	+ return loss, {"vf_loss": reduced_loss}
	173	+
	174	+ def get_critic_forward_output_and_loss_func(self):
	175	+
	176	+ def fwd_output_and_loss_func(data, model):
	177	+ tokens, labels, loss_mask, attention_mask, position_ids = self.get_actor_and_critic_train_data(data)
	178	+ curr_values = model(input_ids=tokens, position_ids=position_ids, attention_mask=attention_mask)
	179	+ return curr_values, partial(self.critic_loss_func, data)
	180	+
	181	+ return fwd_output_and_loss_func
	182	+
	183	+ def setup_model_and_optimizer(self,
	184	+ model_provider_func,
	185	+ model_type,
	186	+ no_wd_decay_cond=None,
	187	+ scale_lr_cond=None,
	188	+ lr_mult=1.0,
	189	+ load_arg='ref_model',
	190	+ ):
	191	+ """Setup model and optimizer."""
	192	+ args = get_args()
	193	+ timers = get_timers()
	194	+
	195	+ model = get_model(model_provider_func, model_type)
	196	+ unwrapped_model = unwrap_model(model)
	197	+
	198	+ kwargs = {}
	199	+ for f in dataclasses.fields(OptimizerConfig):
	200	+ if hasattr(args, f.name):
	201	+ kwargs[f.name] = getattr(args, f.name)
	202	+ config = OptimizerConfig(**kwargs)
	203	+ config.timers = timers
	204	+ optimizer = get_megatron_optimizer(config, model, no_wd_decay_cond,
	205	+ scale_lr_cond, lr_mult)
	206	+ opt_param_scheduler = get_optimizer_param_scheduler(optimizer)
	207	+
	208	+ args.iteration, args.num_floating_point_operations_so_far = load_checkpoint(
	209	+ model, optimizer, opt_param_scheduler, load_arg=load_arg)
	210	+
	211	+ # get model without FP16 and/or DDP wrappers
	212	+ if args.iteration == 0 and len(unwrapped_model) == 1 \
	213	+ and hasattr(unwrapped_model[0], 'init_state_dict_from_bert'):
	214	+ print_rank_0("Initializing ICT from pretrained BERT model")
	215	+ unwrapped_model[0].init_state_dict_from_bert()
	216	+ if args.fp16:
	217	+ optimizer.reload_model_params()
	218	+
	219	+ return model, optimizer, opt_param_scheduler
	220	+
	221	+ def setup_model(self, model_provider_func,
	222	+ model_type,
	223	+ load_arg,
	224	+ ):
	225	+ """Setup model."""
	226	+ model = get_model(model_provider_func, model_type)
	227	+ load_checkpoint(model, None, None, load_arg=load_arg)
	228	+ return model
	229	+
	230	+ def compute_advantages_and_returns(self, rollout_batch):
	231	+ INVALID_LOGPROB = 1
	232	+ args = get_args()
	233	+
	234	+ context_length = rollout_batch['context_length']
	235	+ sequence_lengths = rollout_batch['sequence_lengths']
	236	+ logprobs = rollout_batch['actor_logprobs']
	237	+ ref_logprobs = rollout_batch['ref_logprobs']
	238	+ responses = rollout_batch['padded_query_responses'][:, context_length:].contiguous()
	239	+ response_lengths = sequence_lengths - context_length - 1
	240	+
	241	+ response_idxs = torch.arange(responses.shape[1], device=logprobs.device).repeat(responses.shape[0], 1)
	242	+ padding_mask = response_idxs > response_lengths.unsqueeze(1)
	243	+
	244	+ logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB)
	245	+ ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB)
	246	+
	247	+ response_lengths_p1 = response_lengths + 1
	248	+ padding_mask_p1 = response_idxs > (response_lengths_p1.unsqueeze(1))
	249	+
	250	+ rewards = rollout_batch["scores"]
	251	+ values = rollout_batch["values"]
	252	+ values = torch.masked_fill(values, padding_mask_p1, 0)
	253	+ kl = (logprobs - ref_logprobs)
	254	+
	255	+ # compute rewards_with_kl
	256	+ non_score_reward = -args.kl_coef * kl
	257	+ rewards_with_kl = non_score_reward.cpu()
	258	+ actual_start = torch.arange(rewards_with_kl.size(0), device=rewards_with_kl.device)
	259	+ actual_end = torch.where(response_lengths_p1 < rewards_with_kl.size(1), response_lengths_p1, response_lengths)
	260	+ rewards_with_kl[[actual_start.cpu(), actual_end.cpu()]] += rewards.squeeze(1).cpu()
	261	+ rewards_with_kl = rewards_with_kl.contiguous().to(torch.cuda.current_device())
	262	+
	263	+ # compute advantages and returns
	264	+ lastgaelam = 0
	265	+ advantages_reversed = []
	266	+ gen_length = responses.shape[1]
	267	+ for t in reversed(range(gen_length)):
	268	+ nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0
	269	+ delta = rewards_with_kl[:, t] + args.gamma * nextvalues - values[:, t]
	270	+ lastgaelam = delta + args.gamma * args.lam * lastgaelam
	271	+ advantages_reversed.append(lastgaelam)
	272	+ advantages = torch.stack(advantages_reversed[::-1], dim=1)
	273	+ returns = advantages + values
	274	+ advantages = masked_whiten(advantages, ~padding_mask)
	275	+ advantages = torch.masked_fill(advantages, padding_mask, 0)
	276	+
	277	+ if args.empty_unused_memory_level >= 1:
	278	+ torch.cuda.empty_cache()
	279	+
	280	+ return (advantages.contiguous().to(torch.cuda.current_device()),
	281	+ returns.contiguous().to(torch.cuda.current_device()),
	282	+ padding_mask.contiguous().to(torch.cuda.current_device()),
	283	+ padding_mask_p1.contiguous().to(torch.cuda.current_device()))
	284	+
	285	+ def set_model_eval(self, model):
	286	+ for model_module in model:
	287	+ model_module.eval()
	288	+
	289	+ def set_model_train(self, model):
	290	+ for model_module in model:
	291	+ model_module.train()

Amindspeed_llm/tasks/posttrain/trl_ppo/TrlPPOTrainer.py+559-0

@@ -0,0 +1,559 @@
	1	+# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.
	2	+import time
	3	+import gc
	4	+import torch
	5	+import torch.nn.functional as F
	6	+
	7	+from megatron.core import mpu, tensor_parallel
	8	+from megatron.core.pipeline_parallel import get_forward_backward_func
	9	+from megatron.training.initialize import set_jit_fusion_options
	10	+from megatron.training.utils import (
	11	+ get_batch_on_this_cp_rank,
	12	+ calc_params_l2_norm,
	13	+ check_adlr_autoresume_termination,
	14	+ print_rank_0)
	15	+from megatron.training.global_vars import (
	16	+ get_args,
	17	+ get_signal_handler,
	18	+ get_timers,
	19	+ get_num_microbatches)
	20	+from megatron.training.training import (
	21	+ print_datetime,
	22	+ get_one_logger,
	23	+ append_to_progress_log,
	24	+ num_floating_point_operations,
	25	+ training_log
	26	+)
	27	+from megatron.inference.text_generation.communication import broadcast_from_last_pipeline_stage
	28	+from mindspeed_llm.tasks.posttrain.rlxf.workers.actor_train_infer import generate_attention_mask, generate_position_ids_from_attention_mask
	29	+
	30	+from mindspeed_llm.tasks.posttrain.trl_ppo.TrlPPOEngine import TrlPPOEngine
	31	+from mindspeed_llm.training.training import get_profiler, is_profile_enabled
	32	+from mindspeed_llm.training.initialize import set_jit_fusion_options
	33	+from mindspeed_llm.training.utils import get_tune_attention_mask, get_finetune_data_on_this_tp_rank, generate_actual_seq_len
	34	+from mindspeed_llm.tasks.posttrain.trl_ppo.utils import pad_to_tensor_dict, save_checkpoint_and_time, save_checkpoint
	35	+
	36	+_TRAIN_START_TIME = time.time()
	37	+
	38	+
	39	+class TrlPPOTrainer():
	40	+ def __init__(self, process_non_loss_data_func=None):
	41	+ self.trl_ppo_engine = TrlPPOEngine()
	42	+
	43	+ def initialize(self):
	44	+ args = get_args()
	45	+ self.timers = get_timers()
	46	+
	47	+ if args.log_progress:
	48	+ append_to_progress_log("Starting job")
	49	+
	50	+ # Set pytorch JIT layer fusion options and warmup JIT functions.
	51	+ set_jit_fusion_options()
	52	+
	53	+ # Adjust the startup time so it reflects the largest value.
	54	+ # This will be closer to what scheduler will see (outside of
	55	+ # image ... launches.
	56	+ global _TRAIN_START_TIME
	57	+ start_time_tensor = torch.tensor([_TRAIN_START_TIME], dtype=torch.float, device='cuda')
	58	+ torch.distributed.all_reduce(start_time_tensor, op=torch.distributed.ReduceOp.MIN)
	59	+ _TRAIN_START_TIME = start_time_tensor.item()
	60	+ print_rank_0('time to initialize megatron (seconds): {:.3f}'.format(time.time() - _TRAIN_START_TIME))
	61	+ print_datetime('after megatron is initialized')
	62	+
	63	+ one_logger = get_one_logger()
	64	+ if one_logger:
	65	+ one_logger.log_metrics({
	66	+ 'train_iterations_warmup': 5
	67	+ })
	68	+
	69	+ self.trl_ppo_engine.initialize()
	70	+
	71	+ def train_step(self, forward_step_func, exp_data,
	72	+ model, optimizer, opt_param_scheduler):
	73	+ """Single training step."""
	74	+ args = get_args()
	75	+
	76	+ seq_len = exp_data["padded_query_responses"].shape[1]
	77	+
	78	+ # Set grad to zero.
	79	+ for model_chunk in model:
	80	+ model_chunk.zero_grad_buffer()
	81	+ optimizer.zero_grad()
	82	+
	83	+ # Forward pass.
	84	+ forward_backward_func = get_forward_backward_func()
	85	+ losses_reduced = forward_backward_func(
	86	+ forward_step_func=forward_step_func,
	87	+ data_iterator=exp_data,
	88	+ model=model,
	89	+ num_microbatches=get_num_microbatches(),
	90	+ seq_length=seq_len,
	91	+ micro_batch_size=args.micro_batch_size,
	92	+ decoder_seq_length=args.decoder_seq_length,
	93	+ forward_only=False)
	94	+
	95	+ # Empty unused memory.
	96	+ if args.empty_unused_memory_level >= 1:
	97	+ torch.cuda.empty_cache()
	98	+
	99	+ # Update parameters.
	100	+ self.timers('optimizer', log_level=1).start(barrier=args.barrier_with_L1_time)
	101	+ update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
	102	+ self.timers('optimizer').stop()
	103	+
	104	+
	105	+ # Update learning rate.
	106	+ if update_successful:
	107	+ increment = get_num_microbatches() * \
	108	+ args.micro_batch_size * \
	109	+ args.data_parallel_size
	110	+ opt_param_scheduler.step(increment=increment)
	111	+ skipped_iter = 0
	112	+ else:
	113	+ skipped_iter = 1
	114	+
	115	+ # Empty unused memory.
	116	+ if args.empty_unused_memory_level >= 2:
	117	+ torch.cuda.empty_cache()
	118	+
	119	+ if mpu.is_pipeline_last_stage(ignore_virtual=True):
	120	+ # Average loss across microbatches.
	121	+ loss_reduced = {}
	122	+ for key in losses_reduced[0].keys():
	123	+ numerator = 0
	124	+ denominator = 0
	125	+ for x in losses_reduced:
	126	+ val = x[key]
	127	+ # there is one dict per microbatch. in new reporting, we average
	128	+ # over the total number of tokens across the global batch.
	129	+ if isinstance(val, tuple) or isinstance(val, list):
	130	+ numerator += val[0]
	131	+ denominator += val[1]
	132	+ else:
	133	+ # legacy behavior. we average over the number of microbatches,
	134	+ # and so the denominator is 1.
	135	+ numerator += val
	136	+ denominator += 1
	137	+ loss_reduced[key] = numerator / denominator
	138	+
	139	+ return loss_reduced, skipped_iter, grad_norm, num_zeros_in_grad
	140	+ return {}, skipped_iter, grad_norm, num_zeros_in_grad
	141	+
	142	+ def update_log_metrics(self):
	143	+ args = get_args()
	144	+ batch_size = mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
	145	+ args.consumed_train_samples += batch_size
	146	+ self.num_floating_point_operations_so_far += num_floating_point_operations(args, batch_size)
	147	+ self.timers('interval-time').elapsed(barrier=True)
	148	+
	149	+ def step_training_log(self, model, optimizer, loss_dict, skipped_iter, grad_norm, num_zeros_in_grad):
	150	+
	151	+ args = get_args()
	152	+
	153	+ # Logging.
	154	+ loss_scale = optimizer.get_loss_scale().item()
	155	+
	156	+ params_norm = None
	157	+ if args.log_params_norm:
	158	+ params_norm = calc_params_l2_norm(model)
	159	+
	160	+ learning_rate = None
	161	+ decoupled_learning_rate = None
	162	+ for param_group in optimizer.param_groups:
	163	+ if param_group['is_decoupled_lr']:
	164	+ decoupled_learning_rate = param_group['lr']
	165	+ else:
	166	+ learning_rate = param_group['lr']
	167	+
	168	+ return training_log(loss_dict, loss_dict,
	169	+ learning_rate,
	170	+ decoupled_learning_rate,
	171	+ self.iteration, loss_scale,
	172	+ self.report_memory_flag, skipped_iter,
	173	+ grad_norm, params_norm, num_zeros_in_grad)
	174	+
	175	+ def rl_steps(self):
	176	+ """
	177	+ This function is the core process of the entire RL algorithm and will run the followings in the order provided:
	178	+ 1) STEP 1: Actor model generates responses
	179	+ 2) STEP 2: Four model Inference
	180	+ 3) STEP 3: Compute advantages and returns
	181	+ 4) STEP 4: Train actor model and critic
	182	+ """
	183	+
	184	+ args = get_args()
	185	+
	186	+ with torch.no_grad():
	187	+ # STEP 1: Generate data
	188	+ self.trl_ppo_engine.set_model_eval(self.trl_ppo_engine.actor_model.model)
	189	+ self.trl_ppo_engine.set_model_eval(self.trl_ppo_engine.critic_model)
	190	+ rollout_batch = self.rollout()
	191	+
	192	+ # STEP 2: Four model Inference
	193	+ rollout_batch["output_shape"] = [rollout_batch["padded_query_responses"].shape[0],
	194	+ rollout_batch["padded_query_responses"].shape[1] - rollout_batch["context_length"]]
	195	+ rollout_batch["ref_logprobs"] = self.model_forward(self.trl_ppo_engine.ref_model, rollout_batch, "ref_model")
	196	+ rollout_batch["actor_logprobs"] = self.model_forward(self.trl_ppo_engine.actor_model.model, rollout_batch, "actor_model")
	197	+ rollout_batch["scores"] = self.model_forward(self.trl_ppo_engine.reward_model, rollout_batch, "reward_model")
	198	+ rollout_batch["values"] = self.model_forward(self.trl_ppo_engine.critic_model, rollout_batch, "critic_model")
	199	+
	200	+ # broadcast data
	201	+ for key in ["ref_logprobs", "actor_logprobs", "values", "scores"]:
	202	+ if key == "scores":
	203	+ shape = [rollout_batch["padded_query_responses"].shape[0], 1]
	204	+ else:
	205	+ shape = rollout_batch["output_shape"]
	206	+ rollout_batch[key] = broadcast_from_last_pipeline_stage(shape, torch.float32, rollout_batch[key])
	207	+
	208	+ # STEP 3: Compute advantages and returns
	209	+ advantages, returns, padding_mask, padding_mask_p1 = self.trl_ppo_engine.compute_advantages_and_returns(rollout_batch)
	210	+ rollout_batch["advantages"] = advantages
	211	+ rollout_batch["returns"] = returns
	212	+ rollout_batch["padding_mask"] = padding_mask
	213	+ rollout_batch["padding_mask_p1"] = padding_mask_p1
	214	+
	215	+ if args.empty_unused_memory_level >= 1:
	216	+ torch.cuda.empty_cache()
	217	+
	218	+ # STEP 4: Train actor model and critic model
	219	+ self.trl_ppo_engine.set_model_train(self.trl_ppo_engine.actor_model.model)
	220	+ self.trl_ppo_engine.set_model_train(self.trl_ppo_engine.critic_model)
	221	+ self.update_log_metrics()
	222	+
	223	+ actor_loss_dict, actor_skipped_iter, actor_grad_norm, actor_num_zeros_in_grad = \
	224	+ self.train_step(self.trl_ppo_engine.get_actor_forward_output_and_loss_func(),
	225	+ rollout_batch, self.trl_ppo_engine.actor_model.model,
	226	+ self.trl_ppo_engine.actor_optimizer,
	227	+ self.trl_ppo_engine.actor_opt_param_scheduler)
	228	+
	229	+ self.step_training_log(self.trl_ppo_engine.actor_model.model, self.trl_ppo_engine.actor_optimizer, actor_loss_dict,
	230	+ actor_skipped_iter, actor_grad_norm, actor_num_zeros_in_grad)
	231	+
	232	+ critic_loss_dict, critic_skipped_iter, critic_grad_norm, critic_num_zeros_in_grad = \
	233	+ self.train_step(self.trl_ppo_engine.get_critic_forward_output_and_loss_func(),
	234	+ rollout_batch, self.trl_ppo_engine.critic_model,
	235	+ self.trl_ppo_engine.critic_optimizer,
	236	+ self.trl_ppo_engine.critic_opt_param_scheduler)
	237	+ # set report_memory_flag = True only at th first iteration for both actor model and critic model
	238	+ self.report_memory_flag = self.step_training_log(self.trl_ppo_engine.critic_model, self.trl_ppo_engine.critic_optimizer, critic_loss_dict,
	239	+ critic_skipped_iter, critic_grad_norm, critic_num_zeros_in_grad)
	240	+
	241	+ def train(self):
	242	+ self.initialize()
	243	+ args = get_args()
	244	+
	245	+ # Iterations.
	246	+ self.iteration = args.iteration
	247	+ one_logger = get_one_logger()
	248	+ if one_logger and True:
	249	+ iteration_start = self.iteration
	250	+ train_samples_start = args.consumed_train_samples
	251	+ train_samples_target = args.train_samples
	252	+ one_logger.log_metrics({
	253	+ 'train_samples_start': args.consumed_train_samples,
	254	+ 'train_iterations_start': self.iteration,
	255	+ 'train_samples_target': train_samples_target,
	256	+ 'train_iterations_target': args.train_iters,
	257	+ })
	258	+
	259	+ self.num_floating_point_operations_so_far = 0
	260	+ self.timers('interval-time', log_level=0).start(barrier=True)
	261	+ print_datetime('before the start of training step')
	262	+ self.report_memory_flag = True
	263	+
	264	+ if args.manual_gc:
	265	+ # Disable the default garbage collector and perform the collection manually.
	266	+ # This is to align the timing of garbage collection across ranks.
	267	+ assert args.manual_gc_interval >= 0, \
	268	+ 'Manual garbage collection interval should be laerger than or equal to 0.'
	269	+ gc.disable()
	270	+ gc.collect()
	271	+
	272	+ eval_duration = 0.0
	273	+ eval_iterations = 0
	274	+
	275	+ def track_e2e_metrics():
	276	+ # Nested function to track a bunch of E2E APP metrics
	277	+ if one_logger:
	278	+ train_duration = self.timers('interval-time').active_time() # overall_elapsed
	279	+ train_samples = args.consumed_train_samples - train_samples_start
	280	+ train_iterations = self.iteration - iteration_start
	281	+ train_iterations_time_msecs_avg = (train_duration * 1000.0) / train_iterations if train_iterations > 0 else None
	282	+ if eval_iterations > 0:
	283	+ validation_iterations_time_msecs_avg = (eval_duration * 1000.0) / eval_iterations
	284	+ else:
	285	+ validation_iterations_time_msecs_avg = None
	286	+
	287	+ one_logger.log_metrics({
	288	+ 'train_iterations_end': self.iteration,
	289	+ 'train_samples_end': args.consumed_train_samples,
	290	+ 'train_iterations': train_iterations,
	291	+ 'train_samples': train_samples,
	292	+ 'train_iterations_time_msecs_avg': train_iterations_time_msecs_avg,
	293	+ 'validation_iterations_time_msecs_avg': validation_iterations_time_msecs_avg
	294	+ })
	295	+
	296	+ if is_profile_enabled():
	297	+ prof = get_profiler()
	298	+ prof.start()
	299	+
	300	+ self.trl_ppo_engine.set_model_eval(self.trl_ppo_engine.ref_model)
	301	+ self.trl_ppo_engine.set_model_eval(self.trl_ppo_engine.reward_model)
	302	+
	303	+ while self.iteration < args.train_iters:
	304	+ self.iteration += 1
	305	+ self.rl_steps()
	306	+
	307	+ print_datetime('after training is done')
	308	+
	309	+ if self.iteration % args.log_interval == 0:
	310	+ track_e2e_metrics()
	311	+
	312	+ if args.enable_high_availability:
	313	+ args.num_floating_point_operations_so_far = self.num_floating_point_operations_so_far
	314	+ args.iteration = self.iteration
	315	+
	316	+ # Autoresume for actor model and critic model
	317	+ if args.adlr_autoresume and \
	318	+ (self.iteration % args.adlr_autoresume_interval == 0):
	319	+ check_adlr_autoresume_termination(self.iteration, self.trl_ppo_engine.actor_model.model, self.trl_ppo_engine.actor_optimizer,
	320	+ self.trl_ppo_engine.actor_opt_param_scheduler)
	321	+
	322	+ check_adlr_autoresume_termination(self.iteration, self.trl_ppo_engine.critic_model, self.trl_ppo_engine.critic_optimizer,
	323	+ self.trl_ppo_engine.critic_opt_param_scheduler)
	324	+
	325	+ # Checkpointing
	326	+ saved_checkpoint = False
	327	+ if args.exit_signal_handler:
	328	+ signal_handler = get_signal_handler()
	329	+ if any(signal_handler.signals_received()):
	330	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.actor_model.model, self.trl_ppo_engine.actor_optimizer,
	331	+ self.trl_ppo_engine.actor_opt_param_scheduler,
	332	+ self.num_floating_point_operations_so_far,
	333	+ checkpointing_context=None,
	334	+ save_model_type='actor')
	335	+
	336	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.critic_model, self.trl_ppo_engine.critic_optimizer,
	337	+ self.trl_ppo_engine.critic_opt_param_scheduler,
	338	+ self.num_floating_point_operations_so_far,
	339	+ checkpointing_context=None,
	340	+ save_model_type='critic')
	341	+ print_datetime('exiting program after receiving SIGTERM.')
	342	+ break
	343	+
	344	+ if args.save and args.save_interval and \
	345	+ self.iteration % args.save_interval == 0:
	346	+ self.timers('interval-time').stop()
	347	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.actor_model.model, self.trl_ppo_engine.actor_optimizer,
	348	+ self.trl_ppo_engine.actor_opt_param_scheduler,
	349	+ self.num_floating_point_operations_so_far,
	350	+ checkpointing_context=None,
	351	+ save_model_type='actor')
	352	+
	353	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.critic_model, self.trl_ppo_engine.critic_optimizer,
	354	+ self.trl_ppo_engine.critic_opt_param_scheduler,
	355	+ self.num_floating_point_operations_so_far,
	356	+ checkpointing_context=None,
	357	+ save_model_type='critic')
	358	+ saved_checkpoint = True
	359	+ self.timers('interval-time', log_level=0).start(barrier=True)
	360	+
	361	+ # Exiting based on duration
	362	+ if args.exit_duration_in_mins:
	363	+ train_time = (time.time() - _TRAIN_START_TIME) / 60.0
	364	+ done_cuda = torch.cuda.IntTensor(
	365	+ [train_time > args.exit_duration_in_mins])
	366	+ torch.distributed.all_reduce(
	367	+ done_cuda, op=torch.distributed.ReduceOp.MAX)
	368	+ done = done_cuda.item()
	369	+ if done:
	370	+ if not saved_checkpoint:
	371	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.actor_model.model, self.trl_ppo_engine.actor_optimizer,
	372	+ self.trl_ppo_engine.actor_opt_param_scheduler,
	373	+ self.num_floating_point_operations_so_far,
	374	+ checkpointing_context=None,
	375	+ save_model_type='actor')
	376	+
	377	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.critic_model, self.trl_ppo_engine.critic_optimizer,
	378	+ self.trl_ppo_engine.critic_opt_param_scheduler,
	379	+ self.num_floating_point_operations_so_far,
	380	+ checkpointing_context=None,
	381	+ save_model_type='critic')
	382	+
	383	+ print_datetime('exiting program after {} minutes'.format(train_time))
	384	+ break
	385	+
	386	+ # Exiting based on iterations
	387	+ if args.exit_interval and self.iteration % args.exit_interval == 0:
	388	+ if args.save and not saved_checkpoint:
	389	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.actor_model.model, self.trl_ppo_engine.actor_optimizer,
	390	+ self.trl_ppo_engine.actor_opt_param_scheduler,
	391	+ self.num_floating_point_operations_so_far,
	392	+ checkpointing_context=None,
	393	+ save_model_type='actor')
	394	+
	395	+ save_checkpoint_and_time(self.iteration, self.trl_ppo_engine.critic_model, self.trl_ppo_engine.critic_optimizer,
	396	+ self.trl_ppo_engine.critic_opt_param_scheduler,
	397	+ self.num_floating_point_operations_so_far,
	398	+ checkpointing_context=None,
	399	+ save_model_type='critic')
	400	+
	401	+ torch.distributed.barrier()
	402	+ print_datetime('exiting program at iteration {}'.format(self.iteration))
	403	+ break
	404	+ # save actor critic
	405	+
	406	+ if args.empty_unused_memory_level >= 1:
	407	+ torch.cuda.empty_cache()
	408	+
	409	+ if args.manual_gc:
	410	+ if args.manual_gc_interval != 0 and self.iteration % args.manual_gc_interval == 0:
	411	+ gc.collect()
	412	+
	413	+ if is_profile_enabled():
	414	+ prof.step()
	415	+
	416	+ if is_profile_enabled():
	417	+ prof.stop()
	418	+
	419	+ print_datetime('after training is done')
	420	+
	421	+ if args.save and self.iteration != 0 and self.iteration % args.save_interval != 0:
	422	+ save_checkpoint(self.iteration, self.trl_ppo_engine.actor_model.model, self.trl_ppo_engine.actor_optimizer,
	423	+ self.trl_ppo_engine.actor_opt_param_scheduler,
	424	+ self.num_floating_point_operations_so_far,
	425	+ checkpointing_context=None,
	426	+ save_model_type='actor')
	427	+
	428	+ save_checkpoint(self.iteration, self.trl_ppo_engine.critic_model, self.trl_ppo_engine.critic_optimizer,
	429	+ self.trl_ppo_engine.critic_opt_param_scheduler,
	430	+ self.num_floating_point_operations_so_far,
	431	+ checkpointing_context=None,
	432	+ save_model_type='critic')
	433	+
	434	+ def get_batch(self, data_iterator):
	435	+ """Generate a batch."""
	436	+
	437	+ keys = ['input_ids', 'attention_mask']
	438	+ args = get_args()
	439	+ if args.reset_position_ids:
	440	+ keys += ['position_ids']
	441	+ data_type = torch.int64
	442	+
	443	+ if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()):
	444	+ if args.variable_seq_lengths and args.pipeline_model_parallel_size > 2:
	445	+ tokens, attention_mask = get_finetune_data_on_this_tp_rank(data_iterator)
	446	+
	447	+ return tokens, None, None, attention_mask, None
	448	+ else:
	449	+ if args.reset_position_ids:
	450	+ # Broadcast data.
	451	+ data_b = tensor_parallel.broadcast_data(keys, next(data_iterator), data_type)
	452	+ generate_actual_seq_len(data_b)
	453	+
	454	+ return None, None, None, None, None
	455	+
	456	+ # Broadcast data.
	457	+ data_b = tensor_parallel.broadcast_data(keys, next(data_iterator), data_type)
	458	+ tokens = data_b.get('input_ids').long()
	459	+ attention_mask_1d = data_b.get('attention_mask').long()
	460	+
	461	+ if args.reset_position_ids:
	462	+ position_ids = data_b.get('position_ids').long()
	463	+ generate_actual_seq_len(data_b)
	464	+ batch = {
	465	+ 'tokens': tokens,
	466	+ }
	467	+ batch = get_batch_on_this_cp_rank(batch)
	468	+ batch['attention_mask'] = None
	469	+ batch['position_ids'] = position_ids
	470	+ return batch.values()
	471	+
	472	+ attention_mask = get_tune_attention_mask(attention_mask_1d)
	473	+ return tokens, attention_mask, None
	474	+
	475	+ def rollout(self):
	476	+ """
	477	+ Actor model generates responses
	478	+ """
	479	+
	480	+ args = get_args()
	481	+ args.tokenizer_padding_side = "left"
	482	+
	483	+ tokens, attention_mask, position_ids = self.get_batch(self.trl_ppo_engine.train_data_iterator)
	484	+ inputs = self.trl_ppo_engine.tokenizer.batch_decode(tokens, skip_special_tokens=True)
	485	+
	486	+ queries, responses, context_length = (
	487	+ self.trl_ppo_engine.actor_model.generate(input_ids=inputs,
	488	+ do_sample=args.do_sample,
	489	+ top_k=args.top_k,
	490	+ top_p=args.top_p,
	491	+ max_new_tokens=args.max_new_tokens,
	492	+ max_length=args.max_length,
	493	+ stream=False,
	494	+ detokenize=False,
	495	+ return_output_log_probs=False
	496	+ ))
	497	+
	498	+ queries = [query.tolist() for query in queries]
	499	+ responses = [response.tolist() for response in responses]
	500	+ responses_ori_length, responses_pad_length = pad_to_tensor_dict(responses, pad_multi_of=1)
	501	+ prompts_ori_length, prompts_pad_length = pad_to_tensor_dict(queries, "left", pad_multi_of=1)
	502	+ padded_query_responses = torch.tensor([prompt + response for prompt, response in zip(queries, responses)], device=torch.cuda.current_device())
	503	+ sequence_lengths = torch.tensor([(prompts_pad_length.item() + response_length) for response_length in responses_ori_length], device=torch.cuda.current_device())
	504	+ attention_mask = generate_attention_mask(padded_query_responses.tolist(), prompts_ori_length, prompts_pad_length,
	505	+ responses_ori_length, responses_pad_length)
	506	+ position_ids = generate_position_ids_from_attention_mask(padded_query_responses.tolist(), prompts_ori_length, prompts_pad_length)
	507	+
	508	+ rollout_batch = {
	509	+ "context_length": prompts_pad_length.item(),
	510	+ "sequence_lengths": sequence_lengths,
	511	+ "padded_query_responses": padded_query_responses,
	512	+ "attention_mask": torch.tensor(attention_mask, device=torch.cuda.current_device()),
	513	+ "position_ids": torch.tensor(position_ids, device=torch.cuda.current_device()),
	514	+ }
	515	+
	516	+ if args.empty_unused_memory_level >= 1:
	517	+ torch.cuda.empty_cache()
	518	+
	519	+ return rollout_batch
	520	+
	521	+ def model_forward(self, model, input_data, output_type):
	522	+ """
	523	+ Model forward_only step
	524	+ """
	525	+
	526	+ args = get_args()
	527	+ seq_len = input_data["padded_query_responses"].shape[1]
	528	+
	529	+ forward_backward_func = get_forward_backward_func()
	530	+ output_tensor = forward_backward_func(
	531	+ forward_step_func=self.trl_ppo_engine.forward_only_function,
	532	+ data_iterator=input_data,
	533	+ model=model,
	534	+ num_microbatches=get_num_microbatches(),
	535	+ seq_length=seq_len,
	536	+ micro_batch_size=args.micro_batch_size,
	537	+ forward_only=True)
	538	+
	539	+ context_length = input_data["context_length"]
	540	+ tokens = input_data["padded_query_responses"].clone()
	541	+ sequence_lengths = input_data["sequence_lengths"].clone()
	542	+
	543	+ if mpu.is_pipeline_last_stage():
	544	+ output_logits = torch.cat([o['logits'] for o in output_tensor], dim=0)
	545	+ if output_type == 'actor_model' or output_type == 'ref_model':
	546	+ response_logits = output_logits[:, context_length - 1:-1].to(torch.float32)
	547	+ logprob = F.log_softmax(response_logits, dim=-1)
	548	+ tokens_indices = torch.unsqueeze(tokens[:, context_length:], 2)
	549	+ output_tensor = torch.gather(logprob, 2, index=tokens_indices).squeeze(2).contiguous()
	550	+ elif output_type == 'reward_model':
	551	+ output_tensor = output_logits.squeeze(-1).gather(dim=-1, index=(sequence_lengths.unsqueeze(-1) - 1)).to(torch.float32).contiguous()
	552	+ elif output_type == 'critic_model':
	553	+ output_tensor = output_logits.squeeze(-1)[:, context_length - 1:-1].to(torch.float32).contiguous()
	554	+ else:
	555	+ output_tensor = torch.zeros(1)
	556	+
	557	+ if args.empty_unused_memory_level >= 1:
	558	+ torch.cuda.empty_cache()
	559	+ return output_tensor

Amindspeed_llm/tasks/posttrain/trl_ppo/__init__.py+5-0

@@ -0,0 +1,5 @@
	1	+# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.
	2	+
	3	+from .TrlPPOTrainer import TrlPPOTrainer
	4	+
	5	+__all__ = ["TrlPPOTrainer"]

Mmindspeed_llm/tasks/posttrain/trl_ppo/actor_model.py+1-1

	@@ -25,7 +25,7 @@ class ActorModel(MegatronModuleForCausalLM):
25	input = [val[:context_lengths[i]] for i, val in enumerate(output)]	25	input = [val[:context_lengths[i]] for i, val in enumerate(output)]
26		26
27	if not self.include_input:	27	if not self.include_input:
28	- output = [val[context_lengths[i]:min(self.max_length, context_lengths[i]+self.max_new_tokens)] for i, val in enumerate(output)]	28	+ output = [val[context_lengths[i]:] for i, val in enumerate(output)]
29		29
30	# When batch size > 1, you need truncate the tokens after eos_token_id	30	# When batch size > 1, you need truncate the tokens after eos_token_id
31	self._truncate_in_multi_batch(output)	31	self._truncate_in_multi_batch(output)

Mmindspeed_llm/tasks/posttrain/trl_ppo/utils.py+1-1

	@@ -18,7 +18,7 @@ from megatron.training.global_vars import get_timers
18	from megatron.training.training import compute_throughputs_and_append_to_progress_log	18	from megatron.training.training import compute_throughputs_and_append_to_progress_log
19	from megatron.training.utils import unwrap_model, print_rank_0, append_to_progress_log	19	from megatron.training.utils import unwrap_model, print_rank_0, append_to_progress_log
20	from megatron.training.yaml_arguments import core_transformer_config_from_yaml	20	from megatron.training.yaml_arguments import core_transformer_config_from_yaml
21	-from mindspeed_llm.tasks.posttrain.rm.rm_model import GPTRewardModel	21	+from mindspeed_llm.tasks.posttrain.orm.orm_model import GPTRewardModel
22		22
23		23
24	def model_provider(is_reward_model=False, pre_process=True, post_process=True) -> Union[GPTModel]:	24	def model_provider(is_reward_model=False, pre_process=True, post_process=True) -> Union[GPTModel]:

Asources/images/trl_ppo/llama32_trl_ppo_loss.png+0-0

Asources/images/trl_ppo/ppo_in_rlhf.png+0-0