重要

您正在查看 NeMo 2.0 文档。此版本引入了 API 的重大更改和一个新的库 NeMo Run。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

Hugging Face 集成#

io.ConnectorMixin 类可用于使 NeMo 模型与 Hugging Face 兼容。io.ConnectorMixin 使加载 Hugging Face 模型到 NeMo 和以 Hugging Face 格式保存 NeMo 模型成为可能。下面的 GPTModel 类展示了如何实现这一点（我们可以忽略这里的其他 mixin）

class GPTModel(L.LightningModule, io.IOMixin, io.ConnectorMixin, fn.FNMixin):
    ...

通用基类由多个模型扩展，以提供与 Hugging Face 的双向集成。这些模型包括

  • GemmaModel

  • LLamaModel

  • MistralModel

  • MixtralModel

使用 Hugging Face 检查点微调模型#

要微调模型，请使用以下脚本

import nemo_run as run
from nemo.collections import llm
from nemo import lightning as nl

@run.factory
def mistral():
    return llm.MistralModel()

@run.factory
def trainer(devices=2) -> nl.Trainer:
    strategy = nl.MegatronStrategy(tensor_model_parallel_size=devices)

    return nl.Trainer(
        devices=devices,
        max_steps=100,
        accelerator="gpu",
        strategy=strategy,
        plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"),
    )

resume = nl.AutoResume(import_path="hf://mistralai/Mistral-7B-v0.1")
sft = run.Partial(llm.finetune, model=mistral, data=llm.squad, trainer=trainer, resume=resume)

该脚本将尝试在 MistralModel 上加载名为“hf”的 model_importer。然后，它加载 nemo/collections/llm/gpt/model/mistral.py 中的以下类来执行转换

@io.model_importer(MistralModel, "hf")
class HFMistralImporter(io.ModelConnector["MistralForCausalLM", MistralModel]):
    ...

请注意，此转换仅发生一次。之后，转换后的检查点将从 $NEMO_HOME 目录加载。

创建模型导入器#

要实现自定义模型导入器，您可以遵循 HFMistralImporter 类的结构。以下是关于如何创建自定义模型导入器的分步说明。

  1. 定义一个继承自 io.ModelConnector 的新类

    @io.model_importer(YourModel, "source_format")
class CustomImporter(io.ModelConnector["SourceModel", YourModel]):
    # Implementation here

  2. 将 YourModel 替换为您的目标模型类，将 "source_format" 替换为您要导入的格式，并将 "SourceModel" 替换为源模型类型。您可以选择 "source_format" 为任何字符串。在 Mistral 示例中，我们使用 “hf” 字符串来演示我们是从 Hugging Face 导入。

  3. 实现所需的方法

    class CustomImporter(io.ModelConnector["SourceModel", YourModel]):
    def init(self) -> YourModel:
        # Initialize and return your target model
        return YourModel(self.config, tokenizer=self.tokenizer)

    def apply(self, output_path: Path) -> Path:
        # Load source model, convert state, and save target model
        source = SourceModel.from_pretrained(str(self))
        target = self.init()
        trainer = self.nemo_setup(target)
        self.convert_state(source, target)
        self.nemo_save(output_path, trainer)
        # Clean up and return output path
        teardown(trainer, target)
        return output_path

    def convert_state(self, source, target):
        # Define mapping between source and target model states
        mapping = {
            "source_key1": "target_key1",
            "source_key2": "target_key2",
            # ... more mappings ...
        }
        return io.apply_transforms(source, target, mapping=mapping, transforms=[])

    @property
    def tokenizer(self) -> "YourTokenizer":
        # Return the appropriate tokenizer for your model
        return YourTokenizer(str(self))

    @property
    def config(self) -> YourModelConfig:
        # Load source config and convert to target config
        source_config = SourceConfig.from_pretrained(str(self))
        return YourModelConfig(
            # Set appropriate parameters based on source_config
        )

  4. 实现自定义状态转换

    @io.state_transform 装饰器是用于定义源模型状态和目标模型状态之间自定义转换的强大工具。它允许您指定超出简单键重命名的复杂映射。

    @io.state_transform(
    source_key=("source.key1", "source.key2"),
    target_key="target.key"
)
def _custom_transform(ctx: io.TransformCTX, source1, source2):
    # Implement custom transformation logic
    return transformed_data

    以下列表描述了 state_transform 装饰器的关键方面

    1. 源键和目标键

      • source_key：指定源模型状态中的键。可以是单个字符串或字符串元组。

      • target_key：指定目标模型状态中将存储转换后数据的键。

      • 通配符 *：用于跨多个层或组件应用转换。

    2. 转换函数

      • 装饰的函数接收源张量作为参数。

      • 它应该返回目标模型的转换后张量。

    3. 上下文对象

      • 第一个参数 ctx 是一个 TransformCTX 对象。它提供对源模型和目标模型及其配置的访问。

    4. 多个源键

      • 当指定多个源键时，转换函数接收多个张量作为参数。

    5. 灵活的转换

      • 您可以对张量执行任意操作，包括重塑、连接、拆分或应用数学运算。

    以下示例展示了使用通配符的更复杂转换

    @io.state_transform(
    source_key=("model.layers.*.self_attn.q_proj.weight",
                "model.layers.*.self_attn.k_proj.weight",
                "model.layers.*.self_attn.v_proj.weight"),
    target_key="decoder.layers.*.self_attention.qkv.weight"
)
def _combine_qkv_weights(ctx: io.TransformCTX, q, k, v):
    # Combine separate Q, K, V weights into a single QKV tensor
    return torch.cat([q, k, v], dim=0)

    此转换将源模型中单独的 Q、K 和 V 权重矩阵组合成目标模型的单个 QKV 权重矩阵。在键中使用 * 至关重要

    • 在 source_key 中，model.layers.* 匹配源模型中的所有层。

    • 在 target_key 中，decoder.layers.* 对应于目标模型中的所有层。

    通配符确保此转换自动应用于模型的每一层。如果没有它，您需要手动为每一层编写单独的转换。这使代码更简洁且更易于维护，特别是对于具有许多层的模型。

    转换函数本身 (_combine_qkv_weights) 将为每一层调用一次，其中 q、k 和 v 包含该特定层的权重。

  5. 将这些转换添加到 convert_state 方法

    def convert_state(self, source, target):
    mapping = {
        # ... existing mappings ...
    }
    return io.apply_transforms(source, target, mapping=mapping, transforms=[_custom_transform])

通过遵循此结构，您可以创建一个自定义模型导入器，该导入器将模型从源格式转换为目标格式，处理状态映射和任何必要的转换。

本页内容