LlaMa-Factory源码解析之预训练LLaMA-Factory/src/llamafactory/train/pt/workflow.py -＞ run_pt()

LLaMA-Factory/src/llmtuner/train/pt/workflow.py at main · hiyouga/LLaMA-Factory · GitHub

截止至2024年7月，该框架workflow.py目录从LLaMA-Factory/src/llmtuner/train/pt/workflow.py

LLaMA-Factory/src/llamafactory/train/pt/workflow.py

# Inspired by: https://github.com/huggingface/transformers/blob/v4.34.1/examples/pytorch/language-modeling/run_clm.py

import math
from typing import TYPE_CHECKING, List, Optional

from transformers import DataCollatorForLanguageModeling

from ...data import get_dataset, split_dataset
from ...extras.ploting import plot_loss
from ...model import load_model, load_tokenizer
from ..utils import create_modelcard_and_push
from .trainer import CustomTrainer


if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback

    from ...hparams import DataArguments, FinetuningArguments, ModelArguments


def run_pt(
    model_args: "ModelArguments",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
    callbacks: Optional[List["TrainerCallback"]] = None,
):
    tokenizer = load_tokenizer(model_args)
    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="pt")
    model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Initialize our Trainer
    trainer = CustomTrainer(
        model=model,
        args=training_args,
        finetuning_args=finetuning_args,
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,
        **split_dataset(dataset, data_args, training_args),
    )

    # Training
    if training_args.do_train:
        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
        trainer.save_model()
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()
        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
            plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])

    # Evaluation
    if training_args.do_eval:
        metrics = trainer.evaluate(metric_key_prefix="eval")
        try:
            perplexity = math.exp(metrics["eval_loss"])
        except OverflowError:
            perplexity = float("inf")

        metrics["perplexity"] = perplexity
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    # Create model card
    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)

对上文中Python代码的逐行中文解释，这段代码主要用于配置和执行语言模型的预训练流程：

导入模块和条件编译

python

复制

from typing import TYPE_CHECKING, List, Optional
from transformers import DataCollatorForLanguageModeling
from ...data import get_dataset, split_dataset
from ...extras.ploting import plot_loss
from ...model import load_model, load_tokenizer
from ..utils import create_modelcard_and_push
from .trainer import CustomTrainer

typing.TYPE_CHECKING: 仅在开发过程中用于类型检查的标志，运行时不需要导入。 List, Optional: 来自typing的类型提示类，用于指定函数参数和返回类型。 DataCollatorForLanguageModeling: 来自transformers库的工具，用于为语言建模创建数据批次。 get_dataset, split_dataset, plot_loss, load_model, load_tokenizer, create_modelcard_and_push: 从同一包的其他模块导入的函数，用于处理数据集、绘图、模型加载和创建模型卡片。 CustomTrainer: 可能是基于transformers库中的Trainer类修改的自定义训练器类，用于特定的训练需求。

条件性导入

python

复制

if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback
    from ...hparams import DataArguments, FinetuningArguments, ModelArguments

这块代码条件性导入了静态类型检查所需的类，这些类在运行时不是必需的，这有助于避免导入循环或运行时不必要的内存使用。

函数定义

python

复制

def run_pt(
    model_args: "ModelArguments",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
    callbacks: Optional[List["TrainerCallback"]] = None,
):

run_pt: 此函数负责语言模型的预训练流程。它接受配置和参数，用于指定模型、数据、训练和微调的参数。 callbacks: 可选的回调函数列表，可以用于训练过程中的自定义操作（例如，日志记录、保存检查点）。

加载Tokenizer和数据集

python

复制

    tokenizer = load_tokenizer(model_args)
    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="pt")

这两行代码根据提供的参数加载tokenizer和数据集。stage="pt"可能表明数据集应该为预训练阶段准备。llamafactory/model/loader.py [源码解析]-CSDN博客 tokenizer

详细说明两行代码的区别

python

复制

    tokenizer_module = load_tokenizer(model_args)
    tokenizer = tokenizer_module["tokenizer"]

tokenizer_module = load_tokenizer(model_args)：

这行代码调用 load_tokenizer 函数，传入 model_args 作为参数，返回一个包含分词器相关信息的字典或对象，并将其赋值给 tokenizer_module 变量。 tokenizer_module 可能包含多个键值对，不仅仅是分词器，还可能包含其他相关信息。

tokenizer = tokenizer_module["tokenizer"]：

这行代码从 tokenizer_module 中提取键为 "tokenizer" 的项，并将其赋值给 tokenizer 变量。 tokenizer 变量现在存储了实际的分词器对象，可以用于后续的处理步骤。

总结来说，第一行代码是加载整个分词器模块（可能包含多个信息），第二行代码是从这个模块中提取出具体的分词器对象。

模型加载和数据整合器设置

python

复制

    model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

load_model: 此函数用tokenizer和相应的参数配置初始化模型。 DataCollatorForLanguageModeling: 准备训练数据批次。mlm=False指示未使用掩码语言建模，这表明可能是因果（自回归）语言建模设置。

训练器初始化

python

复制

    trainer = CustomTrainer(
        model=model,
        args=training_args,
        finetuning_args=finetuning_args,
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,
        **split_dataset(dataset, data_args, training_args),
    )

使用加载的模型、训练参数、tokenizer、数据整合器和可选回调初始化CustomTrainer。split_dataset函数可能返回训练/验证分割作为关键字参数。

训练块

python

复制

    if training_args.do_train:
        ...

如果training_args.do_train启用，此块处理训练过程。包括保存模型、记录指标，并可选地绘制损失图。

评估块

python

复制

    if training_args.do_eval:
        ...

这一部分处理模型的评估，基于评估损失计算困惑度，并记录评估指标。

创建模型卡片

    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)

最后，这一行调用函数创建并发布模型卡片，提供关于模型、配置及其性能的全面文档和元数据，这对于机器学习工作流的可重现性和透明度至关重要。

tokenizer ->

llamafactory/model/loader.py [源码解析]-CSDN博客

总结

**文章总结**：
本文介绍了`workflow.py`脚本，该脚本是`LLaMA-Factory`项目中用于配置和执行语言模型预训练流程的关键部分。它通过一系列步骤加载模型参数、数据集、tokenizer，设置训练环境，执行训练与评估，并最终创建模型卡片以文档和发布模型结果。以下是详细的中文概述：
### 主要模块和类导入
- **Type Hints（类型提示）**：从`typing`模块导入`TYPE_CHECKING`, `List`, 和 `Optional`，用于函数参数和返回值的静态类型检查。
- **Transformers库**：导入`DataCollatorForLanguageModeling`，用于将数据准备成可训练的批次。
- **内部模块**：从相对路径导入自定义函数`get_dataset`, `split_dataset`, `plot_loss`, `load_model`, `load_tokenizer`, 和 `create_modelcard_and_push`，以及自定义训练器类`CustomTrainer`。
### 条件性导入
- 仅在类型检查时（使用`TYPE_CHECKING`），从`transformers`库和内部`hparams`模块导入更多配置和参数类（`Seq2SeqTrainingArguments`, `TrainerCallback`, 等），用于更详细的类型注解。
### 函数定义
`run_pt`函数是整个预训练流程的主控函数，它接收多个参数（如`model_args`, `data_args`, `training_args`, `finetuning_args`, 和可选的回调函数列表）来控制训练细节。
### 步骤详解
1. **初始化Tokenizer和数据集**：使用`load_tokenizer`加载分词器，然后基于该分词器和其他参数调用`get_dataset`来加载数据集。
- **Tokenizer扩展说明**：`load_tokenizer`函数返回一个包含分词器相关信息的对象，可能还包含其他信息，从中提取实际的分词器对象使用。
2. **加载模型和设置数据整合器**：通过`load_model`函数加载预训练或微调过的模型，并使用`DataCollatorForLanguageModeling`设置训练数据的整合器。
3. **初始化CustomTrainer**：将加载的模型、参数、tokenizer、数据整合器以及可选的回调函数传递给`CustomTrainer`实例，并使用`split_dataset`返回的训练/验证数据分割作为关键字参数初始化训练器。
4. **训练过程**：如果启用了训练（`training_args.do_train`为真），则调用训练器的`train`方法开始训练，并保存模型、指标和状态。如果需要，还会绘制训练损失图。
5. **评估过程**：如果启用了评估（`training_args.do_eval`为真），则通过调用训练器的`evaluate`方法对模型进行评估，并基于评估损失计算困惑度，最后记录所有评估指标。
6. **创建模型卡片**：最后，调用`create_modelcard_and_push`函数创建包含模型详情、配置和性能的模型卡片，并将其发布到适当的位置以确保工作的可复现性和透明度。
### 总结
该脚本通过将复杂的训练流程分解为清晰的步骤和函数调用，为语言模型的预训练和微调提供了一个结构化的框架。这不仅简化了模型开发和部署过程，还提高了代码的可读性和可维护性。