[CLIP-VIT-L + Qwen] 多模态大模型源码阅读 - trainer篇

参考repo:WatchTower-Liu/VLM-learning; url: VLLM-BASE

前情提要

有关多模态大模型架构中的语言模型部分（MQwen.py）的代码请看（多模态大模型源码阅读 - 1、 多模态大模型源码阅读 - 2， 多模态大模型源码阅读 - 3，多模态大模型源码阅读 - 4）,多模态大模型架构中的视觉模型（visual/CLIP-VIT.py）部分请看多模态大模型源码阅读 - 5
本节主要讲的是项目中的多模态Trainer部分，即项目文件trainer.py，该文件中的代码重构了部分transfomers.trainer的成员方法，以适配多模态场景下的模型训练，包括自定义的损失计算，参数保存，优化器配置，支持分布式训练（多卡场景）。

源码阅读

导包

import torch
from transformers import Trainer
from transformers.trainer import (
    is_sagemaker_mp_enabled,
    get_parameter_names,
    has_length,
    ALL_LAYERNORM_LAYERS,
    logger,
)
import os
from peft import get_peft_model_state_dict

逐行解读

import torch
from transformers import Trainer

torch不必赘述，深度学习的核心出装，构建和训练神经网络的必备库，调包调参侠（我）的福音。
Trainer类主要用于NLP和多模态任务，简化模型训练过程，在后续的代码中作为父类使用。

from transformers.trainer import (
    is_sagemaker_mp_enabled,
    get_parameter_names,
    has_length,
    ALL_LAYERNORM_LAYERS,
    logger,
)

is_sagemaker_mp_enabled检验是否在Amazon SageMaker的模型并行环境中运行。模型并行性允许将模型的不同组件分布到多个GPU设备上，用以加速大规模模型的训练。如果是单卡童鞋就不必在意这个设置~
get_parameter_names用以获取模型中的参数名，在设置优化器参数时，可以区分需要权重衰减的参数和不需要的参数。
has_length检测对象是否有长度信息，用于确定训练过程的迭代次数。在项目代码中没有用到。
ALL_LAYERNORM_LAYERS：包含所有LAYERNORM类型的层，用于在优化器配置中排除这些层的权重衰减。
logger：日志记录，输出训练过程中信息和调试信息。

import os
from peft import get_peft_model_state_dict

os:经常使用的库，主要用来创建文件、文件夹，开关文件。
peft（Parameter-Efficient Fine-Tuning），用于高效微调模型，在微调过程中会冻结预训练模型的大部分参数，仅保留少量的可训练参数，以在尽可能少的资源占用和时间下微调模型适配下游任务， 大名鼎鼎的LoRA、Prefix Tuning、Prompt Tuning 等都在这个库中。get_peft_model_state_dict用于获取微调后的adapter状态字典。例如使用LoRA对模型微调后，可以使用这一方法获取微调后的LoRA adapter状态字典。

compute_loss方法（重构）

class MultiModalTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        return model(
            image=inputs["images"],
            input_ids=inputs["input_ids"],
            labels=inputs["labels"],
        ).loss

整体含义

为多模态场景自定义的损失计算重构方法，以适配多模态形式的输入，如image

逐行解读

class MultiModalTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):

自定义MultiModelTrainer类，继承自transfomers.Trainer，拥有其成员变量和方法。
model:可以同时处理图片和文本类型输入
inputs：包含图片输入，文本索引输入和有监督训练需要的标签数据。
return_outputs：指示是否返回模型输出，考虑到这个项目是科研级代码，所以这个参数没啥用（QWQ）。

        return model(
            image=inputs["images"],
            input_ids=inputs["input_ids"],
            labels=inputs["labels"],
        ).loss

将inputs字典中的对应键下的值传递给model，获取其返回值中的损失值，用于后续的模型优化。

save_model函数（重构）

    def save_model(self, output_dir=None, _internal_call=False):
        from transformers.trainer import TRAINING_ARGS_NAME
        
        # Ensure output_dir is not None
        if output_dir is None:
            output_dir = self.args.output_dir
        
        # Create the output directory if it doesn't exist
        os.makedirs(output_dir, exist_ok=True)
        
        # Save training arguments
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        
        # Access the original model
        model = self.model.module if hasattr(self.model, 'module') else self.model
        
        # Save LLM parameters
        saved_params_LLM = get_peft_model_state_dict(model.LLM)
        torch.save(saved_params_LLM, os.path.join(output_dir, "adapter_model.bin"))
        
        # Save other parameters
        saved_params_other = model.feature_proj.state_dict()
        torch.save(saved_params_other, os.path.join(output_dir, "other_params.bin"))
        
        # Save configuration
        config = model.LLM.peft_config
        selected_adapters = list(config.keys())
        config[selected_adapters[0]].save_pretrained(output_dir, auto_mapping_dict=None)

整体含义

保存训练过程中的模型及其相关配置到指定的目录，重构后适配了多模态模型模型和配置

逐行解读

    def save_model(self, output_dir=None, _internal_call=False):
        from transformers.trainer import TRAINING_ARGS_NAME

output_dir指定模型和相关配置的保存目录，_internal_call并没有用上，可能与某些内部逻辑有关。
TRAINING_ARGS_NAME用于保存训练模型参数名，是一个常量。

        # Ensure output_dir is not None
        if output_dir is None:
            output_dir = self.args.output_dir
        
        # Create the output directory if it doesn't exist
        os.makedirs(output_dir, exist_ok=True)

如果没有指定模型配置的保存目录，就采用配置参数中的输出路径。同时使用os.makedirs方法在指定输出路径下穿件文件夹，exist_ok保证即使文件夹已经存在也不会报错。

        # Save training arguments
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        
        # Access the original model
        model = self.model.module if hasattr(self.model, 'module') else self.model

保存训练参数到指定文件夹的指定文件中，文件名为TRAINING_ARGS_NAME常量。如果模型被封装在DataParallel 或 DistributedDataParallel 中，通过self.model.module访问模型，否则直接使用self.model。

        # Save LLM parameters
        saved_params_LLM = get_peft_model_state_dict(model.LLM)
        torch.save(saved_params_LLM, os.path.join(output_dir, "adapter_model.bin"))

我们传入的model参数实际上是一个以Qwen为语言模型，SIGLIP/CLIP-VIT为视觉模型的多模态模型参数，所有这里的model.LLM大概率是语言模型，saved_params_llm获取语言模型微调后的adapter状态字典，并将其存储到输出路径下的adapter_model.bin文件中。

        # Save other parameters
        saved_params_other = model.feature_proj.state_dict()
        torch.save(saved_params_other, os.path.join(output_dir, "other_params.bin"))

这段代码用于将多模态模型的中间投影层参数存储到other_params.bin文件中，这里的中间投影层可以参考llava的相关论文，用于将视觉模型的输出映射到语言模型的向量空间，大概如下图所示。

projectionW就是中间投影层，也是整个多模态项目的核心出装。

        # Save configuration
        config = model.LLM.peft_config
        selected_adapters = list(config.keys())
        config[selected_adapters[0]].save_pretrained(output_dir, auto_mapping_dict=None)

peft_config方法获取peftmodel微调需要的参数配置。例如使用LoRA进行微调，config会包含所有必要的参数。config参数包含了adapter的类型、参数和设置。selected_keys获取参数字典中的所有键并将其转换为列表。save_pretrained将选择的适配器参数中的第一个存储到指定目录下，设置自动映射为None。

create_optimizer函数（重构）

    def create_optimizer(self):
        if is_sagemaker_mp_enabled():
            return super().create_optimizer()

        opt_model = self.model

        if self.optimizer is None:
            decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
            decay_parameters = [name for name in decay_parameters if "bias" not in name]
            if self.args.feature_proj_lr is not None:
                projector_parameters = [name for name, _ in opt_model.named_parameters() if "feature_proj" in name]
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.feature_proj_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.feature_proj_lr,
                    },
                ]
            else:
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                ]

            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)

            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

        return self.optimizer

整体含义

创建模型优化器，并且对模型的不同部分采取不同的训练策略，例如权重衰减等，并返回一个自定义的优化器。

逐行解读

   def create_optimizer(self):
        if is_sagemaker_mp_enabled():
            return super().create_optimizer()

        opt_model = self.model

如果启用了模型并行，则调用父类的创建优化器方法，否则将成员变量self.model赋值给opt_model，后续将根据opt模型的参数属性创建自定义的优化器。

        if self.optimizer is None:
            decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
            decay_parameters = [name for name in decay_parameters if "bias" not in name]

如果成员变量optimizer为None，代表我们尚未创建一个优化器，进入代码内部。用get_parameter_names获取opt_model中所有LAYERNORM类型层的参数名，并去除掉名字中带有‘bias’（偏置）的参数，这是因为我们不对偏置项进行权重衰减。其余的参数在后续都将应用权重衰减。

            if self.args.feature_proj_lr is not None:
                projector_parameters = [name for name, _ in opt_model.named_parameters() if "feature_proj" in name]

如果设置了投影层的学习率，我们获取opt_model中所有名字里带有‘feature_proj’的参数，这些参数都是投影层参数，代表我们的模型是多模态模型，具有投影层。

                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.feature_proj_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in projector_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.feature_proj_lr,
                    },
                ]

初始化优化器参数组，列表中总共有四个字典，我们逐一来看（看着唬人，其实简单QAQ）。
第一个字典：‘params’键对应于不是投影层参数的所有权重衰减参数；‘weight_dacay’键对应配置参数中的权重衰减值，代表对这些参数应用权重衰减。
第二个字典：‘params’键对应于不是投影层参数的所有权重衰减参数；‘weight_dacay’键的值为0，代表不应用权重衰减。
第三个字典：‘params’键对应于投影层参数的所有权重衰减参数；‘weight_dacay’键对应配置参数中的权重衰减值，“lr”(学习率)为配置参数中的学习率值，代表应用权重衰减，并且初始化学习率。
第四个字典：‘params’键对应于投影层参数的所有权重衰减参数；‘weight_dacay’键的值为0，“lr”(学习率)为配置参数中的学习率值，代表不应用权重衰减，并且初始化学习率。
总而言之，对于非投影层的权重衰减参数，一组应用权重衰减，一组不应用权重衰减。这里的权重衰减参数是如何选出的参考上一段代码。对于投影层的权重衰减参数，一组应用权重衰减，一组不应用权重衰减，并且都有初始的学习率。

            else:
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                ]

与上一段代码相反，这里是不应用投影层的情况。参考上一段代码中的非投影层权重衰减参数的配置即可，同样是一组运用权重衰减，一组不应用权重衰减。

            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)

            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

        return self.optimizer

使用父类的get_optimizer_cls_and_kwargs方法获取优化器类和优化器参数，传递入配置参数。用之前代码定义的优化器参数组和优化器参数初始化优化器实例，并返回。

create_optimizer_and_scheduler函数（重构）

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        super().create_optimizer_and_scheduler(num_training_steps)
        if self.args.local_rank != -1:
            self.model = torch.nn.parallel.DistributedDataParallel(
                self.model,
                device_ids=[self.args.local_rank],
                output_device=self.args.local_rank,
                find_unused_parameters=True
            )

整体含义

这段代码主要用于分布式训练，让模型能够在多个GPU上并行计算。

逐行解读

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        super().create_optimizer_and_scheduler(num_training_steps)

根据传入的训练迭代次数调用父类的create_optimizer_and_scheduler函数。子类在父类的功能上进行拓展。

        if self.args.local_rank != -1:
            self.model = torch.nn.parallel.DistributedDataParallel(
                self.model,
                device_ids=[self.args.local_rank],
                output_device=self.args.local_rank,
                find_unused_parameters=True
            )

如果local_rank为-1，代表不处于分布式训练环境中，反之local_rank指定了GPU的索引，调用torch.nn.parallel.DistributedDataParallel方法，创建DDP模型。DDP可以让模型进行分布式数据并行。其中self.model为模型实例，device_ids指定模型训练用的GPU编码，output_device指定模型输出的GPU编码，find_unused_parameters检查模型在前向传播后是否有未使用的参数。
至此，项目的Trainer.py源码讲解完毕。