Implementación de la clase ChatGPT con RLHF desde cero: análisis de Microsoft DeepSpeed Chat línea por línea

escribe al frente

Este artículo se escribió por primera vez en abril de 2023 en este artículo "Implementación de ChatGPT con RLHF desde cero: de TRL / ChatLLaMA / ColossalChat a DeepSpeed ​​​​Chat". Más tarde, tuve que hablar sobre DSC en el "Campamento fuera de línea de desarrollo de proyectos de modelos grandes". Para evitar que el texto original sea demasiado largo, la parte DSC al final del artículo se extrae de este artículo.

Prefacio

Como se menciona en este artículo, el DeepSpeed ​​​​Chat (DSC para abreviar) de código abierto de Microsoft está bien implementado. Su método de capacitación de tres etapas similar a ChatGPT le brindará un proceso de implementación de código "algoritmo PPO / RLHF" completo y transparente (OK Información puede ayudarle a obtener el resultado con la mitad del esfuerzo)

Después de desenterrar su código clave, encontrará que corresponde uno a uno con la parte principal escrita en otro artículo de este blog ( si aún no ha leído el principio, se recomienda leer este artículo primero: Análisis del principio técnico de ChatGPT , solo comprendiendo los principios se puede comprender mejor la implementación o la implementación real, especialmente la tercera parte del artículo  ), y la correspondencia uno a uno de artículos, principios/algoritmos, fórmulas y códigos puede hacer un cambio cualitativo en su comprensión.

DSC aprovecha al máximo el plan de optimización del proyecto DeepSpeed, teniendo en cuenta que los actores tienen que cambiar repetidamente entre los modos tren (actualización de parámetros) y evaluación (recopilación de experiencias) durante la etapa RLHF. Sin optimización, la velocidad general será muy lenta. Pero la aceleración del tren original y la aceleración de evaluación de DeepSpeed ​​son dos soluciones disociadas. DSC ha diseñado un motor llamado DeepSpeedHybridEngine, que permite a los actores disfrutar de la optimización de la aceleración del tren y de la evaluación durante la etapa RLHF, mejorando la velocidad general de RLHF.

 Para resumirlo en una frase: DeepSpeed ​​​​llegó a acelerar RLHF y se convirtió en chat de velocidad profunda.

Nota: Un estudiante de la clase ChatGPT en línea de julio "Primavera de la clase de fanfarronería" escribió este modelo con gran detalle (más de 5 meses desde el comienzo del año, además de la serie ChatGPT en este blog, este análisis de chat de velocidad profunda en La primavera es mi Los únicos artículos que he visto que son lo suficientemente profundos y detallados, principalmente porque son realmente profundos y detallados. Por un lado, la tecnología es demasiado nueva y, por otro lado, hay demasiados detalles involucrados), por lo que la mayor parte del análisis en este artículo se basa en su blog .

En general, los métodos de entrenamiento de tres fases de DeepSpeed ​​​​Chat e instructGPT son similares: las tres fases están representadas por fase1, fase2 y fase3 respectivamente.

A continuación se presenta una breve descripción de las tres etapas de la formación:

Parte 1 DSC fase-1: Ajuste supervisado

1.1 Proceso de formación en OFV

El código central de la fase 1 se puede encontrar en: apps/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py. El proceso de capacitación se muestra en la siguiente figura (haga clic con el botón derecho en la imagen: abra la imagen en una nueva pestaña para ver  la imagen grande de alta resolución  )

  1. Cargar tokenizador(1-2)
  2. Cargue el modelo base (actualmente solo se admiten algunos modelos CausalLM) (3-4)
  3. Determine si la tecnología LoRA está habilitada en función de si lora_dim (la dimensión de rango bajo de LoRA) está configurada. Si está habilitado, la estructura del modelo base se transformará en LoRA (ver detalles más adelante) y se devolverá el modelo transformado (5-6 )
  4. Determine si "Actualizar solo parámetros LoRA" está habilitado. Si está habilitado, congele los parámetros estructurales restantes y devuelva el modelo congelado (7-8)
  5. Obtener conjunto de datos (9-10)
  6. Crear una instancia de DataLoader(11)
  7. Utilice la tecnología de optimización DeepSpeedEngine de DeepSpeed ​​para envolver objetos como modelos (12)
  8. Antes de iniciar la capacitación formal, primero realizar una evaluación de indicadores, el indicador seleccionado es perplejidad (13-14).
  9. Iniciar entrenamiento, ciclo de época:

1.2 Explicación sobre LoRA y Perplejidad

Hay dos detalles en el proceso anterior que vale la pena mencionar:

  1. Para obtener una explicación detallada de LoRA , consulte la sección 2.2.3 de este artículo "Alpaca-LoRA: ajuste fino de "Alpaca basado en LLaMA" en GPU de consumo a través de la biblioteca PEFT".
  2. DeepSpeed-Chat eligió la perplejidad como índice de evaluación durante el entrenamiento de la fase 1.
    La perplejidad es un índice que mide el rendimiento de un modelo de lenguaje. Mide qué tan bien el modelo entrenado se ajusta a los datos de prueba. Para cada token de la oración de salida, puede obtener el valor de probabilidad de confianza de su salida. Al multiplicar estos valores y tomar el recíproco de su media geométrica, se puede calcular la perplejidad. Es más conciso usar la fórmula para expresar: Entre ellos, la oración de salida tiene un total de tokens
    \text { perplejidad }=\left(\prod_{t=1}^{T} p_{t}\right)^{-\frac{1}{T}}
    y  el valor de probabilidad de confianza del token es y el proceso de tentrenamiento del modelo CausalLM generalmente se optimiza utilizando la pérdida de probabilidad logarítmica. La fórmula de pérdida de salida es la siguiente: entre ellos, la oración de salida tiene un total de tokens y la confianza Por lo tanto , el valor de probabilidad del token t ttth es En realidad, existe la siguiente relación entre la perplejidad y la pérdida de CausalLM: El cálculo de la perplejidad del código fuente relevante también se basa en la fórmula anterior: primero ingrese los datos de verificación en el modelo para obtener el modele la producción de pérdida y luego calcule la perplejidad a través de la relación exponencial entre perplejidad y pérdida.tp_t


    \text {pérdida}=-\frac{1}{T} \sum_{t=1}^{T} \log p_{t}
    tp_t


    \text { perplejidad }=\exp (\text { pérdida })
        def evaluation(model, eval_dataloader):
            """
            以困惑度perplexity为评估指标进行验证
            """
            model.eval()
            losses = 0
            for step, batch in enumerate(eval_dataloader):
                """
                batch: 由input_ids、attention_mask、labels共3个部分组成的dict。
                其中每个部分的shape均为(bs, max_seq_len)
                """
                batch = to_device(batch, device)
                with torch.no_grad():
                    outputs = model(**batch)
    
                """Causal LM 的损失函数为交叉熵损失"""
                loss = outputs.loss
                losses += loss.float()
            losses = losses / (step + 1)
    
            try:
                """困惑度perplexity通常可以通过exp(CELoss)计算得到"""
                perplexity = torch.exp(losses)
            except OverflowError:
                perplexity = float("inf")
    
            try:
            	"""
            	- get_all_reduce_mean中调用了torch.distributed.all_reduce(perplexity, op=torch.distributed.ReduceOp.SUM)
            	- 对所有进程、或者说GPU(因为通常情况下就是单个进程控制单个GPU)中的perplexity进行求和
            	- 然后再除以全局进程数torch.distributed.get_world_size()得到平均的perplexity结果
            	"""
                perplexity = get_all_reduce_mean(perplexity).item()
            except:
                pass
            return perplexity

Parte 2 DSC fase-2: Ajuste del modelo de recompensa

2.1 Datos de entrenamiento

El modelo se entrenará con pares de datos en la forma de los siguientes ejemplos basados ​​en la pérdida de clasificación y, finalmente, se obtendrá un RM (modelo de recompensa) con capacidades de puntuación similares a las humanas.

Nombre del formato de datos    ilustrar     Muestra
frase_elegida El diálogo completo de preferencia humana, obtenido mediante la rápida respuesta de preferencia de conexión elegida, aplicable a la fase 1 y la fase 2. “Humano: ¿Cuéntame sobre Microsoft en unas pocas frases? Asistente: Microsoft es una empresa de software que desarrolla, otorga licencias y brinda soporte para productos de software, incluidos Windows, Office y Windows Phone. Es la empresa de software más grande del mundo por ingresos y la segunda empresa de software más grande del mundo por capitalización de mercado. Microsoft también es un importante proveedor de servicios de computación en la nube, incluida la plataforma de computación en la nube Microsoft Azure y el conjunto de productos Microsoft Office 365”.
rechazar_oración El diálogo completo del rechazo humano se obtiene conectando el mensaje a la respuesta de rechazo rechazado, y es adecuado para la fase 2.  “Humano: ¿Cuéntame sobre Microsoft en unas pocas frases? Asistente: No estoy seguro de lo que quieres decir”.

2.2 Proceso de formación

El proceso de entrenamiento general de la fase 2 se muestra en el diagrama de tiempo UML ( haga clic derecho en la imagen: abra la imagen en una nueva pestaña para ver la imagen grande de alta definición ):

  1. Cargar tokenizador(1-2)
  2. Cargar el modelo (rm_model), lo que implica ciertos cambios estructurales (3-8)
  3. Determine si la tecnología LoRA está habilitada en función de si lora_dim (la dimensión de rango bajo de LoRA) está configurada. Si está habilitado, la estructura del modelo base se transformará en LoRA (ver detalles más adelante) y se devolverá el modelo transformado (9-10 )
  4. Determine si está habilitado "Actualizar solo parámetros LoRA". Si está habilitado, congele los parámetros estructurales restantes y devuelva el modelo congelado (11-12).
  5. Obtener conjunto de datos (13-14)
  6. Crear una instancia de DataCollator para organizar aún más los datos cargados (15-16)
  7. Instanciar DataLoader(17)
  8. Utilice la tecnología de optimización DeepSpeedEngine de DeepSpeed ​​para envolver objetos como rm_model (18)
  9. Antes de comenzar la capacitación formal, primero se realiza la evaluación del indicador, el indicador seleccionado es la precisión de los resultados de la clasificación (19-20)
  10. Iniciar entrenamiento, ciclo de época:

2.3 Explicación detallada de los códigos clave

2.3.1 Estructura específica de RM

Primero, use la clase de transformador AutoModel para leer la red troncal del modelo especificado (sin definir directamente la estructura de la red con un cabezal de salida) y luego introduzca una capa lineal que pueda lograr una reducción de dimensionalidad de tamaño_oculto a 1. Esta capa lineal Se utilizará como salida de la red troncal. Encabezado, genera 1 puntuación para cada posición en la secuencia de entrada.

# applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/main.py
"""
rm_model调用了create_critic_model进行载入
默认情况下rm_model是不启用dropout的
"""
rm_model = create_critic_model(···)

# applications/DeepSpeed-Chat/training/utils/model/model_utils.py
def create_critic_model(···):
    """此处的模型读取方法用的是“AutoModel”,因此此处critic_model只有主干部分"""
    critic_model = create_hf_model(AutoModel, ···)

    """
    critic_model传入RewardModel,将额外得到线性层输出头,
    因此此处的critic_model结构为“v_head + 主干部分”
	"""
    critic_model = RewardModel(critic_model, ···)
    ...
    return critic_model

# applications/DeepSpeed-Chat/training/utils/model/reward_model.py
class RewardModel(nn.Module):
    """
    将读取得到的model的结构修改为适用于RewardModel的形式,
    总的来说即是使用载入的主干网络进行特征提取,
    其所提取的特征(最后层的各位置输出特征hidden_states)将被传入线性层,输出得到1个数值,
    该数值即为分值,因此max_seq_len维度的每个位置均会得到1个分值
    """
    def __init__(self, base_model, ...):
        super().__init__()
		···
        if hasattr(self.config, "word_embed_proj_dim"):
        	"""
			OPT系列模型的word_embed_proj_dim为embedding层的输出维度,
			通常在transformer模型中也就等于 hidden_size,
			v_head将基于主干网络的输出特征 hidden_state 进行分值预测,共输出max_seq_len个分值
			"""
            self.v_head = nn.Linear(self.config.word_embed_proj_dim,
                                    1,
                                    bias=False)
        ···
        """base_model即为主干网络,因此RM最终由1个主干网络和1个线性层构成"""
        self.rwtranrsformer = base_model

La estructura del modelo de RM es básicamente la siguiente (el modelo base aquí es "facebook/opt-125m"), que consta de la red troncal rwtransformer y el cabezal de salida v_head:

RewardModel(
  (v_head): Linear(in_features=768, out_features=1, bias=False)
  (rwtranrsformer): OPTModel(
    (decoder): OPTDecoder(
      (embed_tokens): Embedding(50272, 768, padding_idx=1)
      (embed_positions): OPTLearnedPositionalEmbedding(2050, 768)
      (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (layers): ModuleList(
        (0-11): 12 x OPTDecoderLayer(
          (self_attn): OPTAttention(
            (k_proj): Linear(in_features=768, out_features=768, bias=True)
            (v_proj): Linear(in_features=768, out_features=768, bias=True)
            (q_proj): Linear(in_features=768, out_features=768, bias=True)
            (out_proj): Linear(in_features=768, out_features=768, bias=True)
          )
          (activation_fn): ReLU()
          (self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          (fc1): Linear(in_features=768, out_features=3072, bias=True)
          (fc2): Linear(in_features=3072, out_features=768, bias=True)
          (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        )
      )
    )
  )
)

2.3.2 Formulario de entrada requerido por DataCollator y RM

El recopilador de datos data_collator utilizado en la fase 2 es DataCollatorReward(). El único ejemplo de muestra tomado en esta fase es en realidad un par de datos elegido-rechazado (consulte el bloque de código a continuación).

Es decir, un lote de tamaño de lote elimina pares de datos de tamaño de lote. El recopilador de datos dividirá los pares de datos en oración_elegida y oración_rechazada (el ejemplo se divide en dos). Por lo tanto, de hecho, la cantidad real de datos ingresados ​​en el modelo para un lote debería ser "tamaño_lote" * 2"

# applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/main.py
"""phase2使用的data_collator为DataCollatorReward()"""
data_collator = DataCollatorReward()

# applications/DeepSpeed-Chat/training/utils/data/data_utils.py
class DataCollatorReward:
    def __call__(self, data):
        """
        对dataloader取到的数据 data 进一步整理,将数据整理成batch输入形式
        入参 data 的具体样式可见下个代码块
        """
        batch = {}

        """f为data中的1个tuple,tuple的第0个元素和第2个元素
        分别为chosen_sentence和reject_sentence的input_ids"""
        batch["input_ids"] = torch.cat([f[0] for f in data] + 
        							   [f[2] for f in data],
                                       dim=0)

        """f为data中的1个tuple,tuple的第1个元素和第3个元素
        分别为chosen_sentence和reject_sentence的attention_mask"""
        batch["attention_mask"] = torch.cat([f[1] for f in data] +
                                            [f[3] for f in data],
                                            dim=0)

        """batch的具体样式可见下个代码块"""
        return batch

Y los datos de entrada son una lista de datos por lotes, cada elemento del cual es un par de datos elegidos-rechazados:

    (
	 chosen_sentence_input_ids, 
	 chosen_sentence_attention_mask,
	 reject_sentence_input_ids,
	 reject_sentence_attention_mask
	)

Los elementos 0.º y 2.º de cada conjunto de datos son input_ids, y el 1.º y 3.º elementos son atencion_mask.

El lote de salida es un diccionario: {"input_ids": tensor([...]), "attention_mask": tensor([...])}
y en el valor del diccionario, el elegido está en la primera mitad y el rechazado está en la segunda mitad:

    {
    "input_ids": [
                  chosen_sentence_1_input_ids,
                  chosen_sentence_2_input_ids,
                  ...,
                  reject_sentence_1_input_ids,
                  reject_sentence_2_input_ids,
                  ...
                 ]
    "attention_mask": [
                       chosen_sentence_1_attention_mask,
                       chosen_sentence_2_attention_mask,
                       ...,
                       reject_sentence_1_attention_mask,
                       reject_sentence_2_attention_mask,
                       ...
                      ]
        
    }

Después de la entrada posterior en el modelo, los datos se dividen directamente en la primera mitad y la segunda mitad y se yuxtaponen para obtener los correspondientes pares de datos elegidos-rechazados.

2.3.3 Diseño de recompensas y pérdida de clasificación por pares de toda la conversación

El proceso de propagación hacia adelante de RM no es complicado, en general es:

  1. Los datos pasan a través de la red troncal para obtener la característica de salida de la última capa estados_ocultos con forma (bs*2, max_seq_len, tamaño_oculto);
  2. Luego envíe las características de salida a la capa lineal v_head para obtener las recompensas de puntuación con forma (bs*2, max_seq_len)

La parte más complicada es en realidad "cálculo de la pérdida de clasificación por pares" y "diseño de agregación de puntuaciones".

2.3.3.1 Pérdida de clasificación por pares

\operatorname{pérdida}(\theta)=\mathrm{E}_{\left(\mathrm{x}, \mathrm{y}_{\mathrm{c}}, \mathrm{y}_{\mathrm{ r}}\right) \sim \mathrm{D}}\left[-\log \left(\sigma\left(\mathrm{r}_{\theta}\left(\mathrm{x}, \mathrm{ y}_{\mathrm{c}}\right)-\mathrm{r}_{\theta}\left(\mathrm{x}, \mathrm{y}_{\mathrm{r}}\right)\ derecha)\derecha)\derecha]

Entre ellos, r_\thetase encuentra RM, Xse solicita, y_cse elige, año_rse rechaza (x, y_c)y (x, año_r)se elige_sentencia y rechaza_sentencia respectivamente.
El propósito de esta función de pérdida es maximizar la diferencia entre "clasificación elegida/buena/superior" y "rechazada/mala/clasificación inferior", promoviendo así el aprendizaje r_\thetadel patrón de clasificación correspondiente.

Cuando DeepSpeed-Chat implementó esta parte, r_\theta(x,y_c)eligió r_\theta(x,y_r)la parte de alineación de la respuesta de frase_elegida y frase_rechazada respectivamente. La descripción del texto es ligeramente abstracta. Ver el bloque de código a continuación le ayudará a comprender este concepto:

max_seq_len为10,pad_token_id为0,
有同属同个prompt的chosen_sentence和reject_sentence:
prompt: [11, 22, 33]
chosen_sentence: [11, 22, 33, 44, 55, 66, 0, 0, 0, 0]
reject_sentence: [11, 22, 33, 40, 50, 0, 0, 0, 0, 0]

“两者answer的对齐部分”即为“非prompt部分也非padding部分、但长度要对齐”:
chosen_truncated: [44, 55, 66]
reject_truncated: [40, 50, 0]

所以当上面的chosen_sentence的answer比较长时,reject_sentence在取相应部分时要取至与chosen部分等长为止;
类似的,如果reject_sentence的answer较长时,同理

Para obtener la "parte de alineación" mencionada anteriormente, el código realiza una operación de índice relativamente oscura y abstracta, pero siempre que comprenda que su propósito final es obtener la recompensa de la parte alineada de la frase_elegida y la frase_rechazada, puede calcular la pérdida.

2.3.3.2 Diseño de recompensa por diálogo

        Aunque la recompensa de la "parte alineada" se utiliza para calcular la pérdida de clasificación por pares, la puntuación de predicción del modelo RM para una conversación en realidad toma la recompensa del último token válido del texto de la conversación (generalmente el "token final"). El siguiente bloque proporciona un ejemplo sencillo que ilustra esta situación.

pad_token_id = 0
conversation = [11, 22, 33, 44, 55, 66, 0, 0, 0, 0]
conversation_rewards = [2.01, 0.23, 2.89, 0.66, 0.33, 2.25, 0.36, 0.99, 1.32, 1.62]
token_id为66的token作为该对话的最后1个有效token,
其对应的reward“2.25”将被用于表示整个对话的reward

El código general es el siguiente.

# applications/DeepSpeed-Chat/training/utils/model/reward_model.py
class RewardModel(nn.Module):
	def __init__(self, ···):
		···
	···
	def forward(self, input_ids=None, ···):
		"""获得主干网络的输出的特征"""
		transformer_outputs = self.rwtranrsformer(···)

		"""
		取最后一层的输出特征
		hidden_states.shape: (bs*2, max_seq_len, hidden_size)
		"""
		hidden_states = transformer_outputs[0]

		"""
		将特征送入全连接层得到分数回归值
		rewards.shape: (bs*2, max_seq_len)
		"""
		rewards = self.v_head(hidden_states).squeeze(-1)

		"""先前提及过,实际的bs应该是输入bs的一半"""
		bs = input_ids.shape[0] // 2

		"""区分出chosen和reject"""
		chosen_ids = input_ids[:bs]
		rejected_ids = input_ids[bs:]
		chosen_rewards = rewards[:bs]
		rejected_rewards = rewards[bs:]
		
		loss = 0
        for i in range(bs):
            """
            取出同组chosen和rejected的token_id和分值reward
            chosen_id.shape: (max_seq_len, )
            """
            chosen_id = chosen_ids[i]
            rejected_id = rejected_ids[i]
            chosen_reward = chosen_rewards[i]
            rejected_reward = rejected_rewards[i]

			"""
			下方本应有各种取index相关的操作,
			基于源码解读的可读性考量,且这些部分只是逻辑形式上的弯弯绕绕,与相关原理并不存在直接关系,所以选择暂且将它们忽略
			"""
			
			"""
			c_ind为chosen_sentence的answer后的第一个pad_token的index
			例如pad_token_id=0,sentence[11,22,33,44,55,66,0,0,0,0],c_ind即为第一个pad_token的index=6 """
            c_ind = ···

            """
            r_ind同理,为reject_sentence的answer后的第一个pad_token的index"""
            r_ind = ···

            """end_ind则为两者的较大者"""
            end_ind = max(c_ind, r_ind)

            # 取chosen和rejected第一个不同的地方的index,可以理解为“response中两个回答自由发挥的第1个token的index”
            """divergence_ind为chosen_sentence和reject_sentence两者answer的第1个token的index"""
            divergence_ind = ···

            """
            以chosen_sentence和reject_sentence最先不同的地方为起始、生成结束的地方为终止,取两者在这个片段的对应分值
            这部分其实就是上个代码块提及的“对齐部分”
            """
            c_truncated_reward = chosen_reward[divergence_ind:end_ind]
            r_truncated_reward = rejected_reward[divergence_ind:end_ind]

            """
            (c_truncated_reward - r_truncated_reward).shape: (truncated_seq_len,)
            计算损失时使用了rank loss的形式,并且是对chosen和rejected“对齐片段”进行计算的
            """
            loss += -torch.log(
                torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()
		
        loss = loss / bs
        
        """取代表结束的pad token所在位置的前一个位置(可以理解为的最后一个有效token的位置)的分值作为参考分值"""
            chosen_mean_scores.append(
                chosen_reward[c_ind - 1])  #use the end score for reference
            rejected_mean_scores.append(rejected_reward[r_ind - 1])
        chosen_mean_scores = torch.stack(chosen_mean_scores)
        rejected_mean_scores = torch.stack(rejected_mean_scores)
        
        """返回损失和参考分值"""
        return {
            "loss": loss,
            "chosen_mean_scores": chosen_mean_scores,
            "rejected_mean_scores": rejected_mean_scores,
        }
   ···

2.3.4 Evaluación de indicadores de la fase 2

La métrica de evaluación utilizada por DeepSpeed-Chat en la fase 2 es la precisión de clasificación correcta. El proceso principal es:

  1. Ingrese varios pares de datos elegidos-rechazados (divididos en oración_elegida y oración_rechazada mediante data_collator en el proceso) en RM para razonar y obtener las puntuaciones de cada oración;
  2. Compare la puntuación de la frase_elegida y la puntuación de la frase_rechazada que pertenecen al mismo mensaje. Cuando la puntuación de la frase_elegida es mayor que la puntuación de la frase_rechazada, es una "predicción correcta", de lo contrario es una "predicción incorrecta";
  3. Se cuentan los resultados de las predicciones correctas y la precisión se calcula como índice de evaluación.
  4. Además, las "puntuaciones" promedio de la frase elegida se calcularán durante el proceso de evaluación como referencia.
def evaluation_reward(model, eval_dataloader):
    model.eval()
    """统计预测(赋分)正确的结果
    即 chosen_reward > rejected_reward 的结果数"""
    correct_predictions = 0

    """统计预测总数"""
    total_predictions = 0
    scores = 0
    for step, batch in enumerate(eval_dataloader):
        batch = to_device(batch, device)
        with torch.no_grad():
            """outputs: {'loss':tensor(), 
            			'chosen_mean_scores':tensor(bs,), 
            			'rejected_mean_scores':tensor(bs,)}"""
            outputs = model(**batch)

        """chosen.shape: (bs,)"""
        chosen = outputs["chosen_mean_scores"]

        """rejected.shape: (bs,)"""
        rejected = outputs["rejected_mean_scores"]

        """"赋分正确"即为chosen分值大于rejected分值"""
        correct_predictions += (chosen > rejected).sum()
        total_predictions += chosen.shape[0]

        """累加每个step的平均chosen分值"""
        scores += outputs["chosen_mean_scores"].mean().float()

        if step == 99:  # For faster evaluation and debugging
            break
    """计算acc指标"""
    acc = correct_predictions / total_predictions

    """计算当前step的平均chosen分值"""
    scores = scores / (step + 1)
    try:
        """多进程结果求和求平均"""
        acc = get_all_reduce_mean(acc).item()
        scores = get_all_reduce_mean(scores).item()
    except:
        pass
    return scores, acc

Respecto a RM, lo último que vale la pena mencionar es que en la implementación de DeepSpeed-Chat, la puntuación de predicción de una conversación mediante el modelo RM en realidad se basa en la recompensa del último token del texto de la conversación. Por supuesto, esta no es la "La única forma de usar esto. Este es un diseño de estrategia abierta para calificar conversaciones de esta manera, pero el equipo de DeepSpeed-Chat ha adoptado esta implementación. Por supuesto, los usuarios también pueden desarrollar sus propias estrategias de procesamiento de calificación, como la recompensa promedio en la respuesta. parte, recompensa en secuencia y luego completa. La capa de conexión obtiene la recompensa de agregación, etc.

En nuestra implementación, utilizamos el token final de la secuencia o el primer token de relleno como puntuación agregada y los comparamos. Otros también pueden utilizar la puntuación promedio de la respuesta completa como alternativa.

Parte 3 DSC fase 3: Ajuste fino de RLHF

Esta sección es una adaptación de la tercera parte del análisis del chat de velocidad profunda realizado por estudiantes en el curso en línea ChatGPT en julio de primavera.

3.1 Datos de entrenamiento de RLHF

Nombre del formato de datos    ilustrar     Muestra
inmediato 
 
La descripción de la situación actual proporciona información de entrada de instrucciones para la generación del modelo, puede entenderse como una "pregunta" en el sentido popular y es adecuada para la fase 3.

"Humano: ¿Cuéntame sobre Microsoft en unas pocas frases?

Asistente: "(El ejemplo de texto se proporciona para facilitar la comprensión; de hecho, es input_ids)

secuencia La secuencia de diálogo completa generada por el actor en función de la entrada del mensaje.  

"Humano: ¿Cuéntame sobre Microsoft en unas pocas frases?

Asistente: Microsoft es una empresa de renombre mundial". El ejemplo de texto se proporciona para facilitar la comprensión; de hecho, es input_ids)

problemas de registro  Logits/logaritmo de estrategia del actor basado en la salida de secuencia forma: debería ser (seq_bs, max_seq_len, vocab_size). Después del procesamiento de recopilación, solo se toma el valor log_logit del token de etiqueta real, que es (seq_bs, max_seq_len, 1).
ref_logprobs referencia/logits SFT/logaritmo de estrategia basado en la salida de secuencia forma: debería ser (seq_bs, max_seq_len, vocab_size). Después del procesamiento de recopilación, solo se toma el valor log_logit del token de etiqueta real, que es (seq_bs, max_seq_len, 1).
valor El crítico evalúa el valor de cada posición en la secuencia en función de la salida de la secuencia.  forma: (seq_bs, max_seq_len)
premio eward/RM se basa en la recompensa (ambiental) por todo el resultado de la conversación por secuencia, y se agregará un término de penalización β cuando se implemente el código real. forma: (seq_bs,)
   
máscara_de_atención  Se utiliza para filtrar elementos no válidos. forma: (seq_bs, max_seq_len)

Hay dos puntos que vale la pena mencionar:

  1. Las definiciones de datos empíricos en cada marco no son exactamente las mismas. Por ejemplo, los datos empíricos definidos por ColossalChat tienen más términos "adv" y "recompensa" que aquí (esta recompensa no es la otra recompensa, y la recompensa de ColossalChat se refiere a " corregido por divergencia KL"). "KL_Reward" después de "), pero son esencialmente iguales, excepto que el alcance del marco es diferente, porque adv (función de ventaja Adventure) y KL_Reward se pueden calcular a partir de los elementos existentes logprobs, ref_logprobs , recompensa y valor
  2. Desde la perspectiva de la eficiencia del código, la definición de datos de experiencia de ColossalChat es relativamente más rigurosa, porque las recompensas de ventaja y penalización de KL se pueden calcular a partir de datos de experiencia básicos y se pueden calcular en un solo paso durante la etapa de generación de experiencia
    . está organizado para calcularse en la fase de capacitación y calcularse en cada iteración de PPO (después de todo, las ventajas y las recompensas de penalización de KL se calculan en función de los datos de experiencia básicos, y los datos de experiencia básicos se han determinado en la fase de generación de experiencia, por lo que incluso en diferentes iteraciones de PPO, las ventajas y las recompensas de penalización de KL tampoco cambian, por lo que DeepSpeed-Chat calcula repetidamente las recompensas de penalización de adv y KL. Se estima que el orden de cálculo de este enlace será ajustado por los equipos relevantes en el futuro)

3.2 Todo el proceso de formación de RLHF

Todo el proceso de capacitación de RLHF se muestra en la siguiente figura ( haga clic con el botón derecho en la imagen: abra la imagen en una nueva pestaña para ver la imagen grande en alta definición )

  1. Cargar tokenizador(1-2);
  2. Obtenga el conjunto de datos y cree una instancia de DataCollator (3-9): obtenga el conjunto de datos (4-5) del mensaje para recopilar experiencia, si el entrenamiento no supervisado está habilitado, luego obtenga el conjunto de datos de los datos no supervisados ​​(6-7) y la instancia DataCollator se utiliza para organizar aún más los datos cargados.

    data_collator se crea una instancia de DataCollatorRLHF. Esta clase implementa principalmente "relleno hasta max_prompt_len (el valor predeterminado es la mitad de max_seq_len) y luego voltea". ¿Por qué necesitamos realizar específicamente la operación de volteo (voltear)
    en ficha de aviso? ?
    La razón es que el propósito de usar el mensaje en la fase 3 es ingresar el mensaje en el modelo de actor, y el actor generará contenido posterior de manera autorregresiva basado en el mensaje, para recopilar experiencia

    . 125 m como ejemplo, este modelo La longitud máxima de secuencia (max_seq_len) que se puede admitir es 512, y la fase 3 también se preestablecerá con una longitud máxima de mensaje (max_prompt_len), que generalmente es la mitad de max_seq_len, es decir, 256. El resto se utilizará la mitad de la longitud para generar la
    entrada. Cuando el mensaje no cumple con la longitud máxima del mensaje max_prompt_len, será necesario rellenar el mensaje (reflejado en el código data_collator de la fase 3), y la operación de relleno generalmente agrega el token de relleno directamente a la parte posterior de la secuencia, y la entrada después del relleno se convertirá en En forma de [ indicador, relleno ], la tarea de generación autorregresiva se generará después de pad_token; esto no es razonable
    Por lo tanto, primero debe voltear la entrada del mensaje, realizar una operación de relleno después de voltear y luego voltearla hacia atrás. La entrada rellenada adquiere la forma de [relleno, mensaje]. Para tareas autorregresivas, el contenido del mensaje continúa durante generación.es razonable.

    Debería poder comprender mejor el propósito de esta operación a través de los siguientes ejemplos.
    max_prompt_len = 5
    pad_token_id = 0
    
    prompt_token_ids = [233, 11, 22]
    # padding位于后侧 ×
    prompt_token_ids.padding() = [233, 11, 22, 0, 0]
    
    prompt_token_ids.flip(0) = [22, 11, 233]
    prompt_token_ids.flip(0).padding() = [22, 11, 233, 0, 0]
    # padding位于前侧 √
    prompt_token_ids.flip(0).padding().flip(0) = [0, 0, 233, 11, 22]
  3. InstanciarDataLoader(10);
  4. Utilice DeepSpeedRLHFEngine() para cargar varios modelos (actor, árbitro/SFT, crítico, recompensa/RM) necesarios para el entrenamiento de PPO y encapsúlelos para obtener rlhf_engine(11-12);
  5. Instanciar al formador en gestión de formación del PPO(13-14);
  6. Cree una instancia de MiniDataset para el entrenamiento de PPO (diferente del conjunto de datos mencionado anteriormente, que se utiliza para obtener datos para toda la ronda grande. MiniDataset administra además los datos proporcionados por el conjunto de datos para su asignación a rondas de PPO, es decir, rondas pequeñas para entrenamiento). (15-16);
  7. Iniciar entrenamiento, época redonda grande (prompt_epoch)

3.3 Explicación detallada del código clave de la etapa tres: step3_rlhf_finetuning

3.3.1 Inicialización de cada modelo en la etapa tres: main.py, rlhf_engine.py en step3_rlhf_finetuning

En cuanto a la inicialización del modelo, la clase DeepSpeedRLHFEngine se utiliza en el código fuente para inicializar el actor, ref/SFT, crítico, recompensa/RM, actor_ema y otros modelos, esta clase implementa principalmente:

  1. Al leer el modelo, aunque también admite extraer el modelo correspondiente directamente desde huggingface hub, generalmente los modelos entrenados de fase1 y fase2 se leen desde la ruta local:
    \flecha correcta  actor, ref/SFT y actor_ema ( EMA es ExponentialMovingAverage, que se denomina media móvil exponencial). En chino, es una técnica de entrenamiento de modelos. Los parámetros obtenidos por el modelo durante la k-ésima actualización no utilizan directamente los k-ésimos nuevos parámetros, sino que se obtienen mediante el promedio ponderado de los parámetros históricos en k-1 y el k-ésimos nuevos parámetros. Principalmente para mejorar la estabilidad del proceso de entrenamiento  ) generalmente inicializan el modelo obtenido del entrenamiento de la fase 1;
    \flecha correcta  la crítica y la recompensa generalmente inicializan el modelo obtenido del entrenamiento de la fase 2
  2. Establezca diferentes configuraciones de DeepSpeed ​​(ds_config) para cada modelo relacionado y use DeepSpeedEngine para la encapsulación. Los actores usarán DeepSpeedHybridEngine para la encapsulación de forma predeterminada. A continuación se puede ver una breve introducción a DeepSpeedHybridEngine.
  3. Finalmente, se obtiene un objeto rlhf_engine que contiene todos los modelos relacionados.

El código correspondiente es el siguiente.

  1. En main.py en step3_rlhf_finetuning , puedes ver que se llama a DeepSpeedRLHFEngine

    # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py
    """
    使用DeepSpeedRLHFEngine类直接初始化模型
    当然其内部仍旧调用了“create_hf_model”方法来读取模型,
    但其中实现了更为精细的DeepSpeed控制
    """
    rlhf_engine = DeepSpeedRLHFEngine(···)
  2. La implementación de DeepSpeedRLHFEngine se encuentra en step3_rlhf_finetuning/rlhf_engine.py , que implica la inicialización de cuatro modelos, como actor, árbitro, crítico y recompensa.

    # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py
    class DeepSpeedRLHFEngine():
        def __init__(···):
            """
            加载模型并进行DS封装
            1. actor与ref(以及actor_ema)通常都初始化自phase1训练所得的模型;
            2. critic与reward通常都初始化自phase2训练所得的模型
            根据它们的入参就能知道
            """
            ···
    
            """此处的actor是模型经过DeepSpeed封装后得到的DeepSpeedHybridEngine对象"""
            self.actor = self._init_actor(actor_model_name_or_path)
    
            """此处的reference是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
            self.ref = self._init_ref(actor_model_name_or_path)
    
            self.actor_ema = None
            """如果开启了ema,则初始化并封装ema"""
            if self.args.enable_ema:
                """此处的ema是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
                self.actor_ema = self._init_ema(actor_model_name_or_path)
    
            """此处的critic是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
            self.critic = self._init_critic(critic_model_name_or_path)
    
            """此处的reward是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
            self.reward = self._init_reward(critic_model_name_or_path)
  3. Entonces, los detalles de inicialización del actor son los siguientes

    # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py
    def _init_actor(self, actor_model_name_or_path):
        """
        初始化actor并使用DeepSpeedHybridEngine封装
        :param actor_model_name_or_path: phase1训练好的actor模型路径
        :return: 经DeepSpeedHybridEngine封装的actor
        """
    	···
    
        """
        DS Config
        根据传参构建ds config,
        与其他相关模型不同的地方在于,如果传参指定启用了enable_hybrid_engine,
        那么HybridEngine将作用于actor,对actor进行封装,
        因为HybridEngine可以使得模型可以在训练与推理两种模式中进行自动切换,
        同时享有训练与推理的优化,
        这对于既需要进行推理生成、又需要进行训练的actor来说是有增益作用的。
        """
        ds_config = get_train_ds_config(···,
            enable_hybrid_engine=self.args.enable_hybrid_engine,
            ···)
    	···
    
        # Model
        """使用CausalLM结构载入模型及权重,实例化actor"""
        actor_model = create_hf_model(
            model_class=AutoModelForCausalLM,
            model_name_or_path=actor_model_name_or_path,
            ds_config=ds_config,
            ···)
    
        # LoRA
        """如果开启LoRA训练则添加LoRA旁路"""
        if self.args.actor_lora_dim > 0:
            actor_model = convert_linear_layer_to_lora(···)
            if self.args.only_optimize_lora:
                actor_model = only_optimize_lora_parameters(actor_model)
    
        # Optimizer
        """实例化优化器:分组权重衰减等"""
        AdamOptimizer = DeepSpeedCPUAdam if self.args.offload else FusedAdam
        optim_params = get_optimizer_grouped_parameters(
            actor_model, self.args.actor_weight_decay)
        optim = AdamOptimizer(optim_params,
                              lr=self.args.actor_learning_rate,
                              betas=(0.9, 0.95))
    
        # LR Scheduler
        """实例化学习率调度器"""
        lr_scheduler = get_scheduler(
            name=self.args.lr_scheduler_type,
            optimizer=optim,
            num_warmup_steps=self.args.num_warmup_steps,
            num_training_steps=self.num_total_iters,
        )
    	
    	"""
        DeepSpeedEngine封装
        若ds_config中定义了启用HybridEngine,
        则返回的actor_engine不仅是个DeepSpeedEngine实例,
        确切地说还是个DeepSpeedHybridEngine实例,集成有HybridEngine的优化
        """
        actor_engine, *_ = deepspeed.initialize(model=actor_model,
                                                optimizer=optim,
                                                lr_scheduler=lr_scheduler,
                                                config=ds_config)
        ···
        return actor_engine

    La inicialización del resto de ref, actor_ema, critica y recompensa es casi la misma, excepto que las configuraciones de ds_config son diferentes, pero eventualmente devolverán objetos encapsulados por DeepSpeedEngine.

3.3.2  La diferencia entre recompensa_score y valores y la adquisición de datos empíricos

3.3.2.0 Adquisición de datos empíricos: step3_rlhf_finetuning/ppo_trainer.py

Similar a la figura siguiente, el proceso de obtención de datos de experiencia en esta etapa de DeepSpeed-Chat es el siguiente:

  1. Preparar datos de aviso (prompt_input_ids, Prompt_attention_mask);
  2. Utilice el actor actual para responder la solicitud y obtener la secuencia de diálogo completa (es decir, la secuencia en la imagen de arriba);
  3. Ingrese seq al actor actual, genere los logits de política actuales (antiguos) (es decir, action_logits en la figura anterior) y tome el logaritmo de logprobs;
  4. Ingrese seq a ref/SFT, genere la estrategia de referencia ref_logits (es decir, sft_logits en la figura anterior) y tome el logaritmo ref_logprobs;
  5. Ingrese seq para recompensa/RM, y la salida es la puntuación de recompensa ambiental (es decir, r (x, y) en la figura anterior);
  6. Ingrese la secuencia al crítico actual y genere los valores estimados del valor actual (antiguo) (es decir, el valor en la imagen de arriba);
  7. En este punto, se han obtenido todos los datos empíricos básicos utilizados para el entrenamiento de PPO. En cuanto a la ventaja, la recompensa y otros datos que se muestran en la figura, en DeepSpeed-Chat, el cálculo comienza durante el proceso de entrenamiento específico (por cierto, la recompensa que se muestra en la
    figura se refiere a la "recompensa KL" mencionada en InstructGPT: para evitar el exceso de confianza en la recompensa aprendida en la fase 2, la divergencia KL de SFT y logits se introduce como recompensa de penalización, que se explicará en detalle a continuación)

La implementación del código relevante se puede ver en el bloque de código a continuación.

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def generate_experience(self, prompts, mask):
    """
    生成经验
    :param prompts: prompt input ids,(bs, max_prompt_len)
    :param mask: prompt attention mask, (bs, max_prompt_len)
    :return:
    """
    
    """将actor、reference、critic、reward转换为eval模式"""
    self.eval()
    
    """
    seq.shape: (seq_bs, max_seq_len)
    seq_bs指:排除较短answer后的batch_size。
    所谓“较短answer”在默认设定中是“序列长度小于1的answer”,
    短answer的seq都被滤掉了,
    所以可能batch_size会比之前小,
    但这个可能性极低,DS-C认为只有在使用未经phase1训练的模型来生成才会出现该情况。
    
    _generate_sequence()更具体的细节可见后续详解
    """
    seq = self._generate_sequence(prompts, mask)
    
    """将actor、critic转换为train模式,因为后续两者仍需要进行训练"""
    self.train()
    
	···
	
    with torch.no_grad():
    	"""
        经验采集:这部分其实就是在获取计算phase3损失函数所需的内容
        1. actor:(旧)策略-output.logits
        2. reference:SFT策略-output_ref.logits
        3. reward:奖励-reward_score,InsructGPT中的r_\theta
        4. critic:(旧)价值估计-values
        """
        output = self.actor_model(seq, attention_mask)
        output_ref = self.ref_model(seq, attention_mask)

        # (seq_bs, max_seq_len, vocab_size)
   		logits = output.logits

    	# (seq_bs, max_seq_len, vocab_size)
    	logits_ref = output_ref.logits
		
		"""价值函数的forward_value()更具体的细节下文马上讲解 """
		"""reward_score取的是answer最后一个token的value"""
        # reward_score.shape: (seq_bs,)
        reward_score = self.reward_model.forward_value(
            seq, attention_mask,prompt_length=self.prompt_length)['chosen_end_scores'].detach()

        """critic_model.forward_value(return_value_only=True)
将返回shape为(seq_bs, max_seq_len)的序列各token的value"""
        # 相当于就输出了旧价值values序列
        values = self.critic_model.forward_value(
        	seq, attention_mask, return_value_only=True).detach()[:, :-1]

    # 返回的dict是“进行PPO所需要使用的一组数据”
    # prompts.shape: (bs, max_prompt_len)
    # logits[:, :-1, :].shape: (seq_bs, max_seq_len - 1)
    # seq[:, 1:].shape: (seq_bs, max_seq_len - 1)
    # gather_log_probs()相当于输入logits和labels,对logits进行log_softmax后取出对应label位置的logit值
    # 因此logprobs.shape: (seq_bs, max_seq_len - 1),ref_logprobs.shape: (seq_bs, max_seq_len - 1)
    # values.shape: (seq_bs, max_seq_len - 1)
    # rewards.shape: (seq_bs,),reward_score在InstructGPT中就是r_\theta
    # input_ids.shape: (seq_bs, max_seq_len)
    # attention_mask.shape: (seq_bs, max_seq_len)
    """gather_log_probs()更具体的细节可见后续详解。"""
    return {
        'prompts': prompts,
        'logprobs': gather_log_probs(logits[:, :-1, :], seq[:, 1:]),
        'ref_logprobs': gather_log_probs(logits_ref[:, :-1, :], seq[:,
                                                                    1:]),
        'value': values,
        'rewards': reward_score,
        'input_ids': seq,
        "attention_mask": attention_mask

A continuación, hay tres puntos que es necesario enfatizar:

3.3.2.1 Generación de secuencias: step3_rlhf_finetuning/ppo_trainer.py

Para el mensaje de este lote, se ingresará al actor actual (para el actor que se obtendrá iterativamente en función de datos empíricos, el "actor actual" en este momento puede considerarse como la "red de estrategia anterior") para generar la respuesta (como se muestra en la figura siguiente) y luego unir el mensaje y la respuesta para obtener la secuencia

El código de muestra es el siguiente.

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def _generate_sequence(self, prompts, mask):
    """
	生成seq
    """
	···
	"""
    获取prompt拼接上answer后的最大长度,实际上相当于max_seq_len,用于对生成长度做限制
    """
    max_min_length = self.max_answer_seq_len + prompts.shape[1]

    with torch.no_grad():
        """调用actor,输入input_ids和attention_mask进行生成"""
        seq = self.actor_model.module.generate(prompts,
                                               attention_mask=mask,
                                               max_length=max_min_length,
                                               min_length=max_min_length)

    """下方操作是为了过滤掉只有极短answer(有效长度小于1)的seq"""
    batch_size = seq.shape[0]

    """prompt长度:实际上就是max_prompt_len"""
    prompt_length = prompts.shape[1]

    """取出answer部分,此时还含有pad token"""
    ans = seq[:, prompt_length:]

    """统计answer的有效长度(去掉pad token后的长度)"""
    valid_ans_len = (ans != self.tokenizer.pad_token_id).sum(dim=-1)

    """排除较短(此处设置为有效长度小于1)的answer,余下的answer将被存入out_seq作为最终返回"""
    out_seq = []
    for i in range(batch_size):
        if valid_ans_len[
                i] <= 1:  # if the answer is shorter than 1 token, drop it
            continue
        else:
            out_seq.append(seq[i:i + 1])
    out_seq = torch.cat(out_seq, dim=0)  # concate output in the batch dim

    # out_seq.shape: (valid_batch_size, max_seq_len)
    return out_seq
3.3.2.2 La diferencia entre recompensa recompensa_score y valores de estimación de valor: utils/model/ reward_model.py

"Recompensa/recompensa ambiental/puntuación_recompensa" proporciona principalmente un valor/puntuación de recompensa para toda la secuencia de diálogo.
"Valor estimado/valores" proporciona una predicción de valor para cada posición en la secuencia de diálogo , que está relacionada con el paso/estado de tiempo. relacionado

Por ejemplo, existe la secuencia de diálogo seq=[11, 22, 33, 44, 55, 66 , 0, 0, 0, 0], y su conversación_rewards = [2.01, 0.23, 2.89, 0.66, 0.33, 2.25 , 0,36, 0,99, 1,32, 1,62]

  • La recompensa recompensa_score solo será un escalar, específicamente recompensa_score_seq = 2,25 correspondiente al último token válido;
  • Sus valores estimados son una matriz unidimensional, como value_seq=[0.21, 1.26, 2.52, 0.03, 0.59, 1.55, 1.75, 2.12, 2.22, 1.32]

De la siguiente manera, la clase de modelo RewardModel del modelo de recompensa implementa el método para obtener recompensas ambientales y estimaciones de valor, a saber, forward_value(),有两点需要重点强调下

  1. Si se llama a este forward_value en la Sección 3.3.2 de esta fase de generación de experiencia, los valores obtenidos son antiguos.
            # 相当于就输出了旧价值values序列
            values = self.critic_model.forward_value(
            	seq, attention_mask, return_value_only=True).detach()[:, :-1]
    Si se llama a este forward_value durante el cálculo de pérdida en la sección "3.3.4.4 Cálculo final de pérdida de valor" a continuación, los valores obtenidos serán nuevos
        # 且此时因为是计算价值损失,所以这里计算的是新价值估计
        value = self.critic_model.forward_value(**batch,
                                                return_value_only=True,
                                                use_cache=False)[:, :-1]
  2. Es diferente de otro método forward() usado en este tipo de entrenamiento de RewardModel. El otro método forward() implementa principalmente la adquisición de recompensas ambientales y el cálculo de la pérdida de clasificación. En resumen, la clase RewardModel implementa el forward()
    usado en entrenamiento.método , y también implementa el método forward_value() utilizado en la inferencia

Finalmente, forward_value se implementa en la clase RewardModel de la siguiente manera:

# applications/DeepSpeed-Chat/training/utils/model/reward_model.py
class RewardModel(nn.Module):

    def __init__(self, base_model, tokenizer, num_padding_at_beginning=0):
        ···
    ···
    def forward(···):
    	"""forward()在之前“2.3.3 整个对话的reward设计和成对排序损失”中已经进行过详解,且与此处所述内容无关,此处不再赘述"""
        ···

    def forward_value(···, return_value_only=False, ···):
        """
        和forward的差别在于:forward需要针对输入的chosen-rejected对计算排序损失并返回
        而forward_value只需要考虑一个输入,然后返回分值
        说白了,forward的输入是数据对,因为要计算数据对的排序损失,而forward value的输入是单个数据,直接推理出其分值
        至于参数return_value_only: 如果设置为True,则在计算出values(在序列上每个位置的分值预测)后直接返回
        """
        
        """经过主干网络正向传播得到输出"""
        transformer_outputs = self.rwtranrsformer(···)

        # hidden_states.shape: (bs, max_seq_len, hidden_size)
        hidden_states = transformer_outputs[0]

        """将隐状态特征传入线性层v_head输出得到分值"""
        # values.shape: (bs, max_seq_len)
        values = self.v_head(hidden_states).squeeze(-1)
        
        if return_value_only:
        	"""
			如果传参中预设了“return_value_only=True”,
			那么将直接返回 values: (bs, max_seq_len)
			"""
            return values
        else:
        	"""否则还将进一步取得reward_score"""
            # 相当于为true  返回values序列,为false 返回values序列和reward标量值 
            bs = values.size(0)
            seq_len = input_ids.shape[1]
            chosen_end_scores = []
            for i in range(bs):
            	···
                # value.shape: (max_seq_len,)
                value = values[i]

                """c_ind即为prompt之后的序列片段中,第一个pad_token的index"""
                c_ind = ···

                """取c_ind的前一个index(实际上就是answer的最终位置)作为reward_score"""
                ···
                chosen_end_scores.append(value[c_ind - 1])
            
            """返回values和reward_score"""
            return {
                "values": values,
                "chosen_end_scores": torch.stack(chosen_end_scores),
            }
3.3.2.3 Procesamiento adicional de logits del modelo de políticas

La forma de los logits generados por el modelo de política (actor, ref/SFT) es (bs, max_seq_len, vocab_size), pero al calcular la penalización por divergencia de KL y el peso de importancia, no es necesario calcular los logits de todos los vocabularios. Solo se requiere el elemento de verdad fundamental. Simplemente calcule los logits de (elementos correspondientes a cada token en la secuencia)

batch_size = 1
max_seq_len = 4
vocab_size  = 3

logits = [
          [[1.23, 2.11, -0.56], 
           [-1.52, -1.11, 1.66], 
           [0.32, 0.13, 1.55], 
           [-0.55, -0.23, -1.62]]
         ]

seq = [
       [2, 2, 0, 1]
      ]

Para CausalLM, el valor de confianza de los logits en el paso de tiempo t es para predecir el token de secuencia en el paso t+1, por lo que los logits[, :-1, :] y seq[:, 1:] son ​​la "predicción y etiqueta " "Relación:

logits[, :-1, :] = [
                      [[1.23, 2.11, -0.56], 
                       [-1.52, -1.11, 1.66], 
                        [0.32, 0.13, 1.55]]
                    ]
seq[:, 1:] = [
              [2, 0, 1]
             ]

Solo necesita extraer los logits de acuerdo con las etiquetas correspondientes de la predicción. Tomando el ejemplo anterior como ejemplo, el resultado final probs es

probs = [
             [-0.56, -1.52, 0.13]
            ]

Por lo tanto, DeepSpeed-Chat define la función together_log_probs() para posprocesar los logits de salida para obtener el resultado logarítmico log_probs

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def gather_log_probs(logits, labels):
    """
    相当于输入logits和labels,对logits进行log_softmax后取出对应label位置耳朵logit值
    :param logits: (bs, seq_len, vocab_size)
    :param labels: (bs, seq_len)
    :return: log_probs_labels.squeeze(-1): (bs, seq_len)
    """

    # log_probs.shape: (bs, seq_len, vocab_size)
    log_probs = F.log_softmax(logits, dim=-1)

    """
    此处gather()可以根据labels(index)来从log_probs中获取对应index的值
    总的来说就是取出logits中对应labels数值位置的值
    log_probs_labels.shape: (bs, seq_len, 1)
    """
    log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
    return log_probs_labels.squeeze(-1)

3.3.3 Gestión de datos de capacitación de PPO-MiniDataset: utils/data/data_utils.py

El conjunto de datos se cargó una vez al principio, pero el conjunto de datos cargado al principio era para la gestión de todos los datos de entrenamiento, mientras que el MiniDataset utilizado en este momento era principalmente para la gestión de los datos utilizados en las iteraciones de entrenamiento de PPO. El proceso de gestión de datos previo a la formación PPO puede entenderse como:

  1. Primero, el cargador de datos extrae del conjunto de datos: 1 dato no supervisado de Prompt_batch y 1 dato de solicitud de Prompt_batch
    ( Nota: Los datos no supervisados ​​aquí son para implementar el elemento ptx. En cuanto a por qué existe este elemento ptx, el motivo es: no El entrenamiento supervisado permite que el modelo tenga la capacidad básica de generar oraciones fluidas, y la introducción de ptx en la etapa RLHF permite que el modelo persiga las preferencias humanas sin olvidar la capacidad básica de generación). Para este último, si los datos rápidos de 1 Prompt_batch se utiliza para la recopilación de experiencias
    , obtendrá 1 dato de experiencia de Prompt_batch.
  2. Después de eso, los datos no supervisados ​​de 1 Prompt_batch y los datos empíricos de 1 Prompt_batch se enviarán a sus respectivas instancias de MiniDataset para su administración: 1 Prompt_batch se dividirá en varios ppo_batch para varias iteraciones de entrenamiento de PPO, como se muestra en el siguiente código ( de step3_rlhf_finetuning/main.py )
    # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py
    """经验数据以及无监督数据都将被MiniDataset所管理"""
    exp_mini_dataset = MiniDataset(···)
    unsup_mini_dataset = MiniDataset(···)  
    
    # out为经验数据     
    out = trainer.generate_experience(···)
    exp_dataset = exp_mini_dataset.add(out)
    unsup_dataset = unsup_mini_dataset.add(batch_unsupervised)

El paso 2 anterior es lo que tiene que hacer MiniDataset, y la clase MiniDataset se define en utils/data/data_utils.py y realiza las siguientes tres operaciones:

  1. separate (): subdividido en datos ppo_batch, su código de implementación específico es
    # applications/DeepSpeed-Chat/training/utils/data/data_utils.py
    class MiniDataset:
        def __init__(self, max_size, small_batch_size):
            """
            :param max_size: batch数。通常此处指“用于给actor做生成的prompt的batch数(注意是batch数不是batch_size)”
            :param small_batch_size: batch size。通常此处指“PPO训练的batch_size”。
            """
            self.dataset = []
            self.max_size = max_size
            self.small_batch_size = small_batch_size
    
        def seperate(self):
        	"""维护1个small_dataset"""
            small_dataset = []
    
            # 从self.dataset中逐个取batch
            for large_batch in self.dataset:
                """判断batch的数据类型(列表/元组/字典),
                根据数据类型取其batch_size,赋值给large_size"""
                if type(large_batch) == list or type(large_batch) == tuple:
                    large_size = len(large_batch[0])
                elif type(large_batch) == dict:
                    large_size = len(large_batch[list(large_batch.keys())[0]])
                else:
                    large_size = len(large_batch)
                """
    
                以下部分代码略微抽象,需要举例说明
                - 比如prompt的batch_size设置为3,PPO训练用的batch_size设置为4,则最后能取来用、存入small_dataset的也就只有3条数据
                - (因为生成用的dataloader只采样出了3条,最多也就只有3条)
    
                - 比如prompt的batch_size设置为5,PPO训练用的batch_size设置为4,则最后能取来用、存入small_dataset的就是2组数据
                - (第1组为idx0,idx1,idx2,idx3共4条数据、第2组为idx4共1条数据)
    
                - 比如prompt的batch_size设置为9,PPO训练用的batch_size设置为4,则最后能取来用、存入small_dataset的就是3组数据
                - ([0,1,2,3],[4,5,6,7],[8])
                """
                for i in range(0, large_size, self.small_batch_size):
                    if type(large_batch) == list or type(large_batch) == tuple:
                        small_dataset.append(
                            [x[i:i + self.small_batch_size] for x in large_batch])
                    elif type(large_batch) == dict:
                        small_dataset.append({
                            k: v[i:i + self.small_batch_size]
                            for k, v in large_batch.items()
                        })
                    else:
                        small_dataset.append(large_batch[i:i + self.small_batch_size])
            """清空self.dataset"""
            self.free()
            
            """返回最终取用的数据,该ppo_batch数据将用于ppo训练迭代"""
            return small_dataset
  2. add(): obtiene datos por lotes (prompt_batch);

        def add(self, data):
            """
    		在最开始的时候可以传参预设“生成X个batch再进行PPO训练”,
    		此处的max_size就是其中的X,
    		如果少于max_size则将batch数据加入至MiniDataset中,
    		直至达到max_size个batch
    		"""
            if len(self.dataset) < self.max_size:
                self.dataset.append(data)
                if len(self.dataset) == self.max_size:
                    """
                    seperate()主要实现了
                    1. 在batch的基础上,再细分ppo_batch并返回
                    2. 清空MiniDataset中的数据
                    """
                    return self.seperate()
                else:
                    return None
            else:
                raise ValueError(
                    "The dataset is full but we did not stop it. There is a bug in the code."
                )
  3. free(): borra los datos del lote obtenidos y devuelve los datos ppo_batch
        def free(self):
            """清空self.dataset中的数据"""
            self.dataset = []

3.3.4 Entrenamiento de PPO bajo la arquitectura AC: bajo el RM con penalización β y truncamiento, la estrategia se itera continuamente y el valor se estima a través de datos empíricos

Para un lote de datos empíricos recopilados, utilice MiniDataset para procesarlo en varios lotes de datos ppo_batch para múltiples iteraciones de entrenamiento de modelos relacionados.

Los ppo_epochs establecidos en DeepSpeed-Chat, desde la perspectiva del aprendizaje por refuerzo, en realidad representan el número de reutilizaciones de un lote de datos empíricos:

  • Si ppo_epochs se establece en 1, durante el entrenamiento, el lote de datos de experiencia introducido se descartará directamente después de un recorrido completo y luego se realizará la siguiente ronda de Prompt_epoch, momento en el cual se recopilará nuevamente un nuevo lote de datos de experiencia.
  • Si ppo_epochs se establece en n, durante el entrenamiento, el lote de datos de experiencia introducido se recorrerá n veces antes de descartarse, lo que equivale a que este lote de datos de experiencia se reutilice n veces para el entrenamiento fuera de la política.
# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py,以下是其中的第470-490行
for ppo_ep in range(args.ppo_epochs):
    """ppo_epoch循环"""
    for i, (exp_data, unsup_data) in enumerate(zip(exp_dataset, unsup_dataset)):
        """
        ppo_step循环:
        从MiniDataset返回的数据中,
        取1个ppo_batch的经验数据和无监督数据来训练
        """

        """经验数据训练,返回actor_loss和critic_loss"""
        actor_loss, critic_loss = trainer.train_rlhf(exp_data)

        """累加本ppo_step的指标,后续将除以内层迭代次数计算均值"""
        actor_loss_sum += actor_loss.item()
        critic_loss_sum += critic_loss.item()
        average_reward += exp_data["rewards"].mean()

        """无监督数据训练"""
        if unsupervised_training_enabled:
            """返回无监督损失"""
            unsup_loss = trainer.train_unsupervised(unsup_data, 
            											args.unsup_coef)
            """累加本ppo_step的无监督损失,后续将除以内层迭代次数计算均值"""
            unsup_loss_sum += unsup_loss.item()

        """PPO训练迭代次数(ppo_step)+1"""
        inner_iter += 1

        """是否启用指数移动平均技术"""
        if args.enable_ema:
            moving_average(rlhf_engine.actor,
                           rlhf_engine.actor_ema,
                           zero_stage=args.actor_zero_stage)

	"""打乱数据供off-policy复用"""
    random.shuffle(exp_dataset)
    random.shuffle(unsup_dataset)

1 La capacitación de PPO se administra mediante el método train_rlhf (), que implementa principalmente " Nota, si no comprende el siguiente contenido, puede combinarlo con " Análisis de principios técnicos de ChatGPT: desde el algoritmo PPO de RL, RLHF a GPT4, instructGPT " Sección 3.2 del artículo "Mejorar la comprensión ":

  1. En el cálculo de la recompensa de penalización de divergencia de KL old_rewards, para evitar el exceso de confianza en las recompensas ambientales aprendidas en la fase 2 r(x,y), se agrega el término de penalización de divergencia de KL:

    r_{KL} = r(x,y) - \beta \log \frac{\pi_{old}^{RL}(y|x)}{\pi^{SFT}(y|x)}

  2. Cálculo de ventajas y rendimientos.
    La implementación de ventajas de la mayoría de los marcos, incluido este marco, no utiliza exclusivamente el error TD, sino que combina el método MC basado en el error TD, es decir, GAE (estimación de ventaja generalizada); para
    todos ttrayectoria tcon una \lambda=1longitud \lambda=0de En otras palabras,  el retorno cuando alcanza un cierto paso de tiempo es
    \begin{array}{c} \hat{A}_{t}=\delta_{t}+(\gamma \lambda) \delta_{t+1}+(\gamma \lambda)^{2} \delta_ {t+2}+\cdots+(\gamma \lambda)^{T-t+1} \delta_{T-1} \\ \text { donde } \delta_{t}=r_{KL, t}+\ gamma \cdot V_{\text {antiguo }}\left(s_{t+1}\right)-V_{\text {antiguo }}\left(s_{t}\right) \end{array}
    tt

    R_t = \hat{A}_t + V_t

  3. En 1 ppo_batch, la fórmula de cálculo de pérdidas del actor es:
    p g_{-} pérdida=E_{\tau \sim \pi_{\text {antiguo }}^{RL}} E_{\left(s_{t}, a_{t}\right) \sim \tau}\ izquierda[\max \left(-\hat{A}_{t} \cdot \frac{p_{\text {nuevo }}^{RL}\left(a_{t} \mid s_{t}\right) }{p_{\text {antiguo }}^{RL}\left(a_{t} \mid s_{t}\right)},-\hat{A}_{t} \cdot \operatorname{clip}\ left(\frac{p_{nuevo}^{RL}\left(a_{t} \mid s_{t}\right)}{p_{\text {antiguo }}^{RL}\left(a_{t} \mid s_{t}\right)}, 1-\epsilon, 1+\epsilon\right)\right)\right]
    Entre ellas, \ podersolo se refiere al contenido de la parte de "respuesta" y no incluye la parte de aviso.
  4. En 1 ppo_batch, la fórmula de cálculo de pérdida del crítico es:
    adaptar el nuevo valor estimado V_ {nuevo}para que no se desvíe demasiado del valor estimado anterior cuando se recopila la experiencia, de modo que la reproducción de la experiencia aún pueda ser efectiva:

    V_{clip} = clip(V_{nuevo}, V_{antiguo}-\phi, V_{old}+\phi)

    El crítico se ajustará al retorno R:

    vf\_loss = \frac{1}{2} \cdot E_{\tau \sim \pi_{old}^{RL}} E_{s_t \sim {\tau}} [\max((V_{nuevo}( s_t)-R_t)^2, (V_{clip}(s_t)-R_t)^2)]

    Entre ellos, \ podersolo se refiere al contenido de la parte de "respuesta" y no incluye la parte de aviso, lo que equivale a enfatizar que "esta fórmula de pérdida solo calcula la parte de respuesta y no incluye la pérdida de la parte de aviso". en esta fórmula."

  A continuación, veamos la implementación del código. Para garantizar la fluidez de la lectura, los estudiantes de la clase ChatGPT en línea de julio ajustaron parte del código en la primavera, de modo que el código de función correspondiente esté conectado detrás de su llamada, para facilitar la comparación específica de sus parámetros pasados, de modo que para distinguir las estrategias nuevas y antiguas pasadas, las estrategias antiguas y nuevas Para que la estimación del valor, etc. sea
        más clara, dividí el código en varias secciones y agregué una serie de fórmulas, diagramas, explicaciones y explicaciones. Finalmente, combiné "código y diagramas" para hacer un análisis más intuitivo y brindarte una explicación única y transparente.

3.3.4.1 La primera es una serie de definiciones y una penalización de KL por las recompensas en la etapa dos.

La fórmula correspondiente a agregar una penalización de KL a la recompensa en la etapa 2 se amplía a (de 3.1.3 Etapa 3 de entrenamiento de InstructGPT en otro artículo de este blog: Análisis de los principios técnicos de ChatGPT : la estrategia de cómo optimizar aún más el modelo a través de la algoritmo PPO)

\begin{alineado} objetivo(\phi ) &= E_{(x,y)\sim D_{\pi _{\phi }^{RL}}} [r_\theta (x,y) - \beta log( \pi _{\phi }^{RL}(y|x) / \pi ^{SFT}(y|x) )] + \gamma E_{x\sim D_{pretrain}} [log(\pi _{ \phi }^{RL})] \\&= E_{(x,y)\sim D_{\pi _{ }^{RL'}}} \left [ \frac{\pi _{\phi }^ {RL}(y|x)}{\pi ^{RL'}(y|x)}r_{\theta'}(x,y) - \beta log(\pi^{RL'}(y|x ) / \pi ^{SFT}(y|x) ) \right ] + \gamma E_{x\sim D_{pretrain}} [log(\pi _{\phi }^{RL})] \\&= E_{(x,y)\sim D_{\pi _{ }^{RL'}}} \left [ \min \left(\frac{\pi_{\phi }^{RL}(y|x)} {\pi ^{RL'}(y|x)} r_{\theta'}(x,y),{clip}\left(\frac{\pi_{\phi }^{RL}(y|x) }{\pi ^{RL'}(y|x)}, 1-\varepsilon, 1+\varepsilon\right) r_{\theta'}(x,y)\right) - \beta log(\pi^ {RL'}(y|x) / \pi ^{SFT}(y|x) ) \right ]+ \gamma E_{x\sim D_{pretrain}} [log(\pi _{\phi }^{ RL})]\\&= E_{(x,y)\sim D_{\pi _{ }^{RL'}}} \left [ \min \left(\frac{\pi_{\phi }^{RL}(y|x)}{\pi ^{ RL'}(y|x)} A^{\theta^{RL'}}\left(x,y\right),{clip}\left(\frac{\pi_{\phi }^{RL}( y|x)}{\pi ^{RL'}(y|x)}, 1-\varepsilon, 1+\varepsilon\right) A^{\theta^{RL'}}\left(x,y\ right)\right) \right ]+ \gamma E_{x\sim D_{pretrain}} [log(\pi _{\phi }^{RL})] \end{aligned}

El diagrama correspondiente es
"Y hay dos puntos que merecen especial atención:

  1. Cuando se implementa el código real, cuando se aplica la penalización de KL con β a RM, el numerador toma la estrategia anterior en los datos empíricos ( como se muestra en la fórmula anterior, el correspondiente π(RL')). Por supuesto, incluso si el numerador es la estrategia anterior en los datos empíricos, la relación de penalización correspondiente a β sigue siendo la relación entre lo antiguo y lo nuevo: π(RL')/π(SFT) , porque aunque π(RL') se inicializa inicialmente con π( SFT), después de uno o más pasos, π(RL') Es decir, se actualiza. Si se actualiza en un paso o en varios pasos depende de que las ppo_epochs mencionadas anteriormente sean iguales a 1 o n. En cuanto al índice de penalización en el implementación del código, corresponde a: estrategia antigua/estrategia SFT = log_probs/ref_log_probs), no Los action_logits ambiguos en la Figura 1 a continuación
  2. Además, cuando se implementó el código real, se agregó la penalización de KL a RM y se truncó el clip de recompensa por algunas razones de seguridad, pero el clip de recompensa no se reflejó en la Figura 1 a continuación.

El código correspondiente es

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def train_rlhf(self, inputs):
    """
    使用1个ppo_batch的经验数据,执行1次rlhf训练迭代
    """
    # prompt input ids
    prompts = inputs['prompts']

    # (旧)策略
    log_probs = inputs['logprobs']

    # SFT策略
    ref_log_probs = inputs['ref_logprobs']

    # RM奖励
    reward_score = inputs['rewards']

    # (旧)价值估计
    values = inputs['value']
    attention_mask = inputs['attention_mask']

    # seq input ids
    seq = inputs['input_ids']
	
	"""
    获取prompts的最后1个位置作为start
    比如prompt_len为256,start则为 256-1=255
    这个start主要是用于取出经验数据中的“非prompt”部分(也即“answer+padding”部分)
    """
    start = prompts.size()[-1] - 1
	
	"""
	action_mask相当于取 attention_mask除了第0个序列位置外的部分,
	需要注意的是:
	1. 多数情况下,包括此处在内的transformers风格代码中,
	attention_mask指的实际上是“padding_mask”而非“sequence_mask”;
	2. 之所以要进行[:, 1:]切片,是为了去除第0个位置从而与seq对齐,
	因此 action_mask.shape: (bs, max_seq_len - 1)
	3. 后续将被用于过滤掉pad token位置的信息
	4. 但实际上在后续的使用中,
	基本都会结合上方定义的start,从action_mask中再切片出“非prompt”部分,
	例如 action_mask[start:],实际上就相当于取“非prompt”部分,
	action_mask[start:].shape: (bs, max_answer_len)
	"""
	action_mask = attention_mask[:, 1:]    
    ···

    """经验数据中的价值估计为“旧”价值估计"""
    old_values = values
    with torch.no_grad():
    	
    	###计算KL惩罚修正的奖励################################################
        """
        通过KL散度惩罚,以及r_\theta(来自phase-2的reward model)计算得到修正的奖励,
        注意此处的入参:
        1. log_probs为经验数据中的旧策略『特别注意这里 是针对旧策略log_probs(如上公式所示),非图示中的action_logits』
        2. ref_log_probs为经验数据中的SFT策略
        3. reward_score为经验数据中的RM赋分
        """
        old_rewards = self.compute_rewards(prompts, log_probs,
                                           ref_log_probs, reward_score,
                                           action_mask)
        def compute_rewards(self, prompts, log_probs, ref_log_probs, reward_score,
                    action_mask):
		    """
		    计算实际rewards,涉及(旧)策略与SFT的KL散度惩罚、RM的reward
		    """
		    """计算经验采样时actor与SFT的KL散度惩罚"""
		    kl_divergence_estimate = -self.kl_ctl * (log_probs - ref_log_probs)
		    rewards = kl_divergence_estimate

		    """
		    找到answer的起始start:即prompt的最后1个token位置
		    比如prompts长度为256,answer的起始则为256-1=255
		    """
		    start = prompts.shape[1] - 1

			"""
			ends为batch中各个数据的最后1个有效token的index,
			每个数据的最末有效token位置很大可能是不一样的,
			因此ends是个数组
			"""
		    ends = ···

		    """
		    将RM得到的奖励值限定在一定范围,默认为(-5,5)
            相当于既对RM加了修正,同时又对RM做了截断reward_clip
		    """
		    reward_clip = torch.clamp(reward_score, -self.clip_reward_value,
		                              self.clip_reward_value)
			···

		    """
		    因为batch中每个数据的最末有效token位置很可能不一样,
		    所以无法通过矩阵来并行,需要使用for循环逐个数据处理
		    """
		    for j in range(batch_size):
		    	"""
		        KL_reward = KL + reward
		        加和只在最末有效token上进行
		        """
		        rewards[j, start:ends[j]][-1] += reward_clip[j]	
		    """返回KL rewards"""
		    return rewards
3.3.4.2 A continuación, calcule las ventajas y retornos: get_advantages_and_returns en step3_rlhf_finetuning/ppo_trainer.py

(Una cosa más, como se mencionó al comienzo de la Sección 3.1, de hecho, debido a que la recompensa de ventaja y penalización de KL se puede calcular a partir de los datos básicos de la experiencia, la recompensa de ventaja y penalización de KL se puede calcular en un solo paso durante la generación de experiencia. escenario)


        ###计算优势与回报################################################
        """
        计算优势advantages和回报returns
		注意此处的入参:
        4. old_value为经验数据中的(旧)价值估计
        5. old_rewards为刚才计算得到的KL_reward
		"""
        advantages, returns = self.get_advantages_and_returns(
            old_values, old_rewards, start)
		def get_advantages_and_returns(self, values, rewards, start):
		    """
		    计算优势与回报
		    实现基本与上述公式相同
		    """
		    lastgaelam = 0
		    advantages_reversed = []
		    length = rewards.size()[-1]
		    
		    """反向遍历计算各个时间步的优势advantage"""
		    for t in reversed(range(start, length)):
		    	"""获取下个时间步的价值估计V_{old}(s_{t+1})"""
		        nextvalues = values[:, t + 1] if t < length - 1 else 0.0

		        """计算单步TD-error"""
                # gamma代表折扣因子γ,δ 表示TD误差
		        delta = rewards[:, t] + self.gamma * nextvalues - values[:, t]

		        """累计优势"""
		        lastgaelam = delta + self.gamma * self.lam * lastgaelam

		        """存储各个时间步的优势"""
		        advantages_reversed.append(lastgaelam)

		    """对逆序的优势列表进行正序处理,得到正常时间步排列的优势"""
		    advantages = torch.stack(advantages_reversed[::-1], dim=1)

		    """
		    return_t = adv_t + v_(s_t)
		    由优势计算得到回报
		    """
		    returns = advantages + values[:, start:]

		    """返回优势与回报"""
		    return advantages.detach(), returns

————————————————
Además, cité una ventaja informática completa en la Sección 3.2 del artículo  " Análisis de los principios técnicos de ChatGPT: del algoritmo RL PPO, RLHF a GPT4, instructGPT " Los ejemplos de valores y valores de retorno pueden aclarar todo el proceso de cálculo a través de ejemplos.
"Para ayudarlo a aclarar aún más cada detalle de los conceptos relacionados, echemos un vistazo más de cerca a cómo DSC calcula los retornos.

En DSC, la función de ventaja se define como la suma descontada de los errores TD (error de diferencia temporal, δ aquí) en todos los pasos de tiempo futuros. Este método de calcular la función de ventaja se llama estimación de ventaja generalizada (GAE).

Primero, necesitamos calcular la secuencia de error TD. Aquí usamos la secuencia de valores V_old y la secuencia de recompensa proporcionadas para calcular:
valores = [-0.2761, -2.3945, 0.1729, -0.0919, -0.0867, -0.0818, -0.0758]
KL_rewards = [-4.6873e-04, -3.1257e-04, 5.8591e-05, -5.5084e-03, -4.0741e-03, -5.5275e-03, -8.5999e-02] γ=0.9
...
"
y el código de cálculo correspondiente a este ejemplo es

# -*- coding: utf-8 -*-
"""
Created on Wed Jan 19 20:42:44 2022
@author: Remixa
"""
import torch

def get_advantages_and_returns(values, rewards, start, length, gamma, lam):
    lastgaelam = 0
    advantages_reversed = []
    length = rewards.size(-1)
    print("#" * 10 + "init:" + "#" * 10)
    print(f"start: {start}, length: {length},\ngamma: {gamma}, lam: {lam}\nvalues: {values}\nrewards: {rewards}")
    for t in reversed(range(start, length)):
        nextvalues = values[:, t + 1] if t < length - 1 else torch.zeros_like(values[:, t])
        delta = rewards[:, t] + gamma * nextvalues - values[:, t]
        lastgaelam = delta + gamma * lam * lastgaelam
        advantages_reversed.append(lastgaelam)
        print("#" * 10 + f"t = {t}" + "#" * 10)
        print(f"t = {t}:\nnextvalues: {nextvalues}\ndelta: {delta}\nlastgaelam: {lastgaelam}\nadvantages_reversed: {advantages_reversed}")
    advantages = torch.stack(advantages_reversed[::-1], dim=1)
    returns = advantages + values[:, start:]
    print("#" * 10 + "result:" + "#" * 10)
    print(f"advantages: {advantages}\nreturns: {returns}")
    return advantages.detach(), returns

if __name__ == '__main__':
    old_values = torch.tensor([[-0.2761, -2.3945, 0.1729, -0.0919, -0.0867, -0.0818, -0.0758]])
    old_rewards = torch.tensor([[-4.6873e-04, -3.1257e-04, 5.8591e-05, -5.5084e-03, -4.0741e-03, -5.5275e-03, -8.5999e-02]])
    start = 3
    length = 7
    gamma = 0.9
    lam = 0.95
    advantages, returns = get_advantages_and_returns(old_values, old_rewards, start, length, gamma, lam)
3.3.4.3 Luego calcule la pérdida de la estrategia: actor_loss_fn de step3_rlhf_finetuning/ppo_trainer.py

El diagrama correspondiente es

Pero antes de mirar el código, hay dos puntos que conviene mencionar:

  1. Se refleja la relación de importancia de las estrategias antigua y nueva: logprobs / old_logprobs  , pero el truncamiento de la nueva estrategia logprobs no se refleja en la figura anterior. Después de todo, ya sabemos que esta relación debe limitarse mediante truncamiento o divergencia KL para evitar que se compare la nueva estrategia. Las actualizaciones son demasiado grandes que la estrategia anterior.
  2. log_probs  es la estrategia (antigua) en los datos empíricos previamente definidos. En consecuencia, la nueva estrategia del actor se define como  actor_log_probs . La relación correspondiente entre estas cuatro definiciones en la implementación del código se muestra en la siguiente figura (la razón de esto puede ser como spring Como dije, una persona escribió la función de pérdida de estrategia y otra persona escribió los datos empíricos (en resumen, preste atención para evitar confusiones)

El código correspondiente es

    ###计算actor损失并更新################################################
    batch = {'input_ids': seq, "attention_mask": attention_mask}

    """将seq经验数据输入至actor,进行自回归预测"""
    actor_prob = self.actor_model(**batch, use_cache=False).logits

    """取出probs,此处为新策略"""
    actor_log_prob = gather_log_probs(actor_prob[:, :-1, :], seq[:, 1:])

    """
    计算actor损失
    注意此处的入参:
    1. actor_log_probs为方才刚输出的新策略
    2. log_probs为之前定义的经验数据中的(旧)策略
    3. advantages为之前计算出的优势
    """
    actor_loss = self.actor_loss_fn(actor_log_prob[:, start:],
                                    log_probs[:, start:], advantages,
                                    action_mask[:, start:])
	def actor_loss_fn(self, logprobs, old_logprobs, advantages, mask):
	    """计算actor的损失"""
	    
	    """
	    重要性采样权重计算:ratio = exp(log(new)-log(old)) 
	    """
	    log_ratio = (logprobs - old_logprobs) * mask
	    ratio = torch.exp(log_ratio)

		"""计算策略梯度损失的2个情况:加权优势 与 裁剪加权优势"""
	    pg_loss1 = -advantages * ratio
	    pg_loss2 = -advantages * torch.clamp(ratio, 1.0 - self.cliprange,
	                                         1.0 + self.cliprange)

	    """
		从策2种情况中选择损失较大者作为真正的损失,
		并且基于ppo_batch内所有数据的所有有效时间步计算平均损失值
		"""
	    pg_loss = torch.sum(torch.max(pg_loss1, pg_loss2) * mask) / mask.sum()
	    return pg_loss

	"""actor反向传播、更新参数"""
    self.actor_model.backward(actor_loss)
    self.actor_model.step()
3.3.4.4 Cálculo final de pérdida de valor: critic_loss_fn de step3_rlhf_finetuning/ppo_trainer.py

    ###计算critic损失并更新################################################
    """将seq经验数据输入至critic,预测得到新价值估计"""
    # 调用的forward_value即是上文「9.3.2.2 奖励reward_score和价值估计values的区别」中分析的那个
    # 且此时因为是计算价值损失,所以这里计算的是新价值估计
    value = self.critic_model.forward_value(**batch,
                                            return_value_only=True,
                                            use_cache=False)[:, :-1]

    """
    计算critic损失
    注意此处的入参:
    1. values为方才刚输出的新价值估计
    2. old_values为经验数据中的(旧)价值估计
    3. returns为之前计算出的回报
    """
	critic_loss = self.critic_loss_fn(value[:, start:], old_values[:,start:],
                                      returns, action_mask[:, start:])
	def critic_loss_fn(self, values, old_values, returns, mask):
	    """计算价值损失"""

	    """裁剪当前新values,使得其不至于太偏离经验采样阶段的旧values"""
	    values_clipped = torch.clamp(
	        values,
	        old_values - self.cliprange_value,
	        old_values + self.cliprange_value,)
	    
	    """计算当前values与回报的L2 Loss"""
	    vf_loss1 = (values - returns)**2

	    """计算裁剪后的当前values与回报的L2 Loss"""
	    vf_loss2 = (values_clipped - returns)**2

	    """
	    选择损失较大者作为真正的损失,
		并且基于ppo_batch内所有数据的所有有效时间步计算平均损失值,
		此外critic损失项的系数为0.5。
	    """
	    vf_loss = 0.5 * torch.sum(
	        torch.max(vf_loss1, vf_loss2) * mask) / mask.sum()
	    return vf_loss
	    
    """critic反向传播、更新参数"""
    self.critic_model.backward(critic_loss)
    self.critic_model.step()
	
	
	"""本次ppo_step将返回actor_loss和critic_loss供指标统计"""
    return actor_loss, critic_loss

Por cierto, durante el entrenamiento RLHF de la fase 3, para permitir que el modelo mantenga el rendimiento del modelo previamente entrenado en la resolución de tareas mientras aprende las preferencias humanas, se introdujo el modelado de lenguaje autorregresivo tradicional para el entrenamiento conjunto.

\mathrm{p}(\mathrm{x})=\prod_{\mathrm{t}=1}^{\mathrm{T}} \mathrm{p}\left(\mathrm{x}_{\mathrm{ t}} \mid \maths{x}_{<\maths{t}}\right)

El código de muestra correspondiente es

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py
unsup_loss = trainer.train_unsupervised(unsup_data, args.unsup_coef)
def train_unsupervised(self, inputs, unsup_coef):
    """
    1个ppo_batch的无监督训练
    :param inputs: dict:input_ids, attention_mask, labels
    :param unsup_coef: 无监督损失系数
    """
    """确保actor处于训练模式,否则将返回报错"""
    self._validate_training_mode()

    """actor进行常规的CausalLM训练"""
    outputs = self.actor_model(**inputs, use_cache=False)
    loss = outputs.loss
    """反向传播、更新参数"""
    self.actor_model.backward(unsup_coef * loss)
    self.actor_model.step()

    return loss

Finalmente, me gustaría citar nuevamente algunos puntos resumidos hechos por los estudiantes en la primavera:

  1. "El entrenamiento de RLHF implica aprendizaje por refuerzo. El proceso de entrenamiento es extremadamente sensible a la configuración de hiperparámetros. Después de probar una variedad de configuraciones de parámetros, el equipo de DeepSpeed-Chat finalmente configuró per_device_train_batch_size (es decir, Prompt_batch_size) = per_device_mini_batch_size (es decir, ppo_batch_size) de forma predeterminada. y la capacitación generada comienza inmediatamente después de Prompt_batch; de esta manera, lo que realmente se hace es un aprendizaje reforzado según la política, recopilando una vez y aprendiendo una vez, y la tasa de utilización de datos no es alta.
  2. Además, el equipo de DeepSpeed-Chat también descubrió que es muy difícil establecer el coeficiente (unsup_coef) para la pérdida de entrenamiento no supervisado, y el proceso de entrenamiento se volverá más oscilante, pero el equipo no gastó mucha energía en ajustar este coeficiente. parámetro.

Por supuesto, estas no son las mejores configuraciones de hiperparámetros. El equipo de DeepSpeed-Chat aún anima a los usuarios a probar más y compartir sus propias experiencias de ajuste de parámetros.

Supongo que te gusta

Origin blog.csdn.net/v_JULY_v/article/details/132939877
Recomendado
Clasificación