大模型微调原理与代码实战案例(五):IA3

随着ChatGPT的快速崛起,大型模型的时代正在发生革命性变化。但对于很多人而言,进行大型模型的预训练或全面微调似乎是遥不可及的。

不过随着多种高效参数微调技术的涌现,科研人员和普通开发者都有机会尝试微调这些庞大的模型了。

本文我将分享了大模型微调技术的原理及代码案例**,完整版代码,可在文末获取**。

IA3 简述

IA3(论文:Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning),通过学习向量来对激活层加权进行缩放,从而获得更强的性能,同时仅引入相对少量的新参数,如下图左边所示,它的诞生背景是为了改进 LoRA。

为了使微调更有效,IA3(通过抑制和放大内部激活注入适配器)使用学习向量重新调整内部激活。这些学习到的向量被注入到典型的基于transformer的架构中的attention和feedforward模块中。原始权重保持冻结,这些学习到的向量是微调期间唯一可训练的参数。与学习 LoRA 更新低秩权重矩阵不同,处理学习向量可以使可训练参数的数量少得多。

与 LoRA 类似,IA3 具有许多相同的优点:

  • IA3 通过大幅减少可训练参数的数量,使微调更加高效。对于 T0 模型,使用 IA3 只有大约 0.01% 的可训练参数,而使用 LoRA 有 > 0.1% 的可训练参数。

  • 原始的预训练权重保持冻结状态,这意味着您可以拥有多个轻量级、便携式 IA3 模型,用于在其之上构建的各种下游任务。

  • 使用 IA3 微调的模型的性能与完全微调的模型的性能相当。

  • IA3 不会增加任何推理延迟,因为适配器(adapter)权重可以与基础模型合并。

原则上,IA3 可以应用于神经网络中权重矩阵的任何子集,以减少可训练参数的数量。根据作者的实现,IA3 权重被添加到 Transformer 模型的 key, value 和 feedforward 层。给定注入 IA3 参数的目标层,可训练参数的数量可以根据权重矩阵的大小确定。

IA3 微调实战

与 PEFT 支持的其他方法一样,要使用 IA3 微调模型,您只需要以下几步:

  • 实例化基本模型。

  • 创建一个配置 (IA3Config),在其中定义 IA3 特定的参数。

  • 使用 get_peft_model() 包装基础模型以获得可训练的 PeftModel。

  • 像平常训练基础模型一样训练 PeftModel。

为了不影响阅读体验,完整版代码在公众号:机器学习社区,回复:ia3,即可获取,这里仅列出关键步骤。

第一步,引进必要的库,如:IA3 配置类 IA3Config

from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, IA3Config, TaskType

第二步,创建 IA3 微调方法对应的配置。

peft_config = IA3Config(task_type=TaskType.CAUSAL_LM,
                        target_modules=["query_key_value", "mlp.dense_4h_to_h"],
                        inference_mode=False, 
                        feedforward_modules=["mlp.dense_4h_to_h"])

参数说明:

  • task_type:指定任务类型。如:条件生成任务(SEQ_2_SEQ_LM),因果语言建模(CAUSAL_LM)等。

  • inference_mode:是否在推理模式下使用Peft模型。

  • target_modules:要替换为 IA3 的模块名称列表或模块名称的正则表达式,例如,注意力块。在 PEFT 中支持的模型中默认的模块名如下所示:

TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING = {
    "t5": ["k", "v", "wo"],
    "mt5": ["k", "v", "wi_1"],
    "gpt2": ["c_attn", "mlp.c_proj"],
    "bloom": ["query_key_value", "mlp.dense_4h_to_h"],
    "roberta": ["key", "value", "output.dense"],
    "opt": ["q_proj", "k_proj", "fc2"],
    "gptj": ["q_proj", "v_proj", "fc_out"],
    "gpt_neox": ["query_key_value", "dense_4h_to_h"],
    "gpt_neo": ["q_proj", "v_proj", "c_proj"],
    "bart": ["q_proj", "v_proj", "fc2"],
    "gpt_bigcode": ["c_attn", "mlp.c_proj"],
    "llama": ["k_proj", "v_proj", "down_proj"],
    "bert": ["key", "value", "output.dense"],
    "deberta-v2": ["key_proj", "value_proj", "output.dense"],
    "deberta": ["in_proj", "output.dense"],
}
  • feedforward_modules:target_modules 中被视为前馈(feedforward)层的模块名称列表或模块名称的正则表达式。虽然学习向量与注意力块的输出激活相乘,但向量与经典前馈层的输入相乘。在 PEFT 中支持的模型中默认的前馈层模块名如下所示:
TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING = {
    "t5": ["wo"],
    "mt5": [],
    "gpt2": ["mlp.c_proj"],
    "bloom": ["mlp.dense_4h_to_h"],
    "roberta": ["output.dense"],
    "opt": ["fc2"],
    "gptj": ["fc_out"],
    "gpt_neox": ["dense_4h_to_h"],
    "gpt_neo": ["c_proj"],
    "bart": ["fc2"],
    "gpt_bigcode": ["mlp.c_proj"],
    "llama": ["down_proj"],
    "bert": ["output.dense"],
    "deberta-v2": ["output.dense"],
    "deberta": ["output.dense"],
}
  • module_to_save:除了 IA3 层之外要设置为可训练并保存在最终检查点中的模块列表。这些通常包括模型的自定义头(head),该头是为微调任务随机初始化的。例如,在序列分类或Token分类任务中,最后一层classifier/score是随机初始化的,因此需要可训练和保存。

第三步,通过调用 get_peft_model 方法包装基础的 Transformer 模型。

model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

通过 print_trainable_parameters 方法可以查看到 IA3 可训练参数的数量(仅为172,032)以及占比(仅为0.0307%)。

trainable params: 172,032 || all params: 559,386,624 || trainable%: 0.0307536849504646

IA3 模型类结构如下所示:

PeftModelForCausalLM(
  (base_model): IA3Model(
    (model): BloomForCausalLM(
      (transformer): BloomModel(
        (word_embeddings): Embedding(250880, 1024)
        (word_embeddings_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (h): ModuleList(
          (0): BloomBlock(
            (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (self_attention): BloomAttention(
              (query_key_value): Linear(
                in_features=1024, out_features=3072, bias=True
                (ia3_l): ParameterDict(  (default): Parameter containing: [torch.FloatTensor of size 3072x1])
              )
              (dense): Linear(in_features=1024, out_features=1024, bias=True)
              (attention_dropout): Dropout(p=0.0, inplace=False)
            )
            (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (mlp): BloomMLP(
              (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)
              (gelu_impl): BloomGelu()
              (dense_4h_to_h): Linear(
                in_features=4096, out_features=1024, bias=True
                (ia3_l): ParameterDict(  (default): Parameter containing: [torch.FloatTensor of size 1x4096])
              )
            )
          )
          ...
          (23): BloomBlock(
            ...
          )
        )
        (ln_f): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
      )
      (lm_head): Linear(in_features=1024, out_features=250880, bias=False)
    )
  )
)

第四步,模型训练的其余部分均无需更改,当模型训练完成之后,保存高效微调部分的模型权重以供模型推理即可。

peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}"
model.save_pretrained(peft_model_id)

输出的模型权重文件如下所示:

/data/nfs/llm/model/bloomz-560m_IA3_CAUSAL_LM
├── [ 398]  adapter_config.json
├── [689K]  adapter_model.bin
└── [ 129]  README.md

0 directories, 3 files

注意:这里只会保存经过训练的增量 PEFT 权重。其中,adapter_config.json 为 IA3 配置文件;adapter_model.bin 为 IA3 权重文件。

第五步,加载微调后的权重文件进行推理。

from peft import PeftModel, PeftConfig

peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}"
config = PeftConfig.from_pretrained(peft_model_id)
# 加载基础模型
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
# 加载PEFT模型
model = PeftModel.from_pretrained(model, peft_model_id)

# 编码
inputs = tokenizer(f'{text_column} : {dataset["test"][i]["Tweet text"]} Label : ', return_tensors="pt")

# 模型推理
outputs = model.generate(
        input_ids=inputs["input_ids"], 
        attention_mask=inputs["attention_mask"], 
        max_new_tokens=10, 
        eos_token_id=3
    )

# 解码
print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))

至此,我们完成了 IA3 的训练及推理。

结语

本文对 IA3 基本原理进行了简述;同时,讲解了 IA3 进行模型训练及推理。

如果觉得我的文章能够能够给您带来帮助,期待您的点赞收藏加关注~~

技术交流群

完整版代码在公众号:机器学习社区,回复:ia3,即可获取。

建了实战技术交流群!想要进交流群的同学,可以直接加微信号:mlc2060。加的时候备注一下:研究方向 +学校/公司+知乎,即可。然后就可以拉你进群了。

前沿技术资讯、算法交流、求职内推、算法竞赛、面试交流(校招、社招、实习)等、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企开发者互动交流~

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值