【HuggingFace LLM】经典NLP微调任务之因果模型

正文

数据集准备

from datasets import load_dataset, DatasetDict

ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train")
ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="validation")

raw_datasets = DatasetDict(
    {
        "train": ds_train,  # .shuffle().select(range(50000)),
        "valid": ds_valid,  # .shuffle().select(range(500))
    }
)

>>>raw_datasets
>>>DatasetDict({
    train: Dataset({
        features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
        num_rows: 606720
    })
    valid: Dataset({
        features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
        num_rows: 3322
    })
})

这是由原始数据集中加载得到的数据格式，内部数据是代码片段，想要做一个代码自动补全功能。

如果直接进行truncation=True的话会丢失很多信息 ，因此使用return_overflow_tokens进行分块

from transformers import AutoTokenizer

context_length = 128
tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")

outputs = tokenizer(
    raw_datasets["train"][:2]["content"],
    truncation=True,
    max_length=context_length,
    return_overflowing_tokens=True,
    return_length=True,
)

print(f"Input IDs length: {len(outputs['input_ids'])}")
print(f"Input chunk lengths: {(outputs['length'])}")
print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")

#Input IDs length: 34
#Input chunk lengths: [128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 117, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 41]
#Chunk mapping: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

可以看到存在一些尾部的不足max_length的batch，可以选择丢弃

def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length:
            input_batch.append(input_ids)
    return {"input_ids": input_batch}


tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)
tokenized_datasets

不过当batch_size过大时，通常会有较大的信息损失，因此可以使用eos_token_id进行连接后，再对其进行分割，这样就只会有一个不足max_length的batch会被丢弃。

max_length = 128
def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=False,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    temp_batch = []

    output_len = len(outputs["input_ids"])

    _ = [temp_batch.extend(input_id + [tokenizer.eos_token_id]) for idx, input_id in enumerate(outputs["input_ids"]) if idx < output_len-1]
    for idx in range(0, len(temp_batch), max_length):
      input_batch.append(temp_batch[idx: min(idx + max_length, len(temp_batch) - 1)])
    return {"input_ids": input_batch}

tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)

初始化新模型

from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

model = GPT2LMHeadModel(config)
model_size = sum(t.numel() for t in model.parameters())
print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")

# GPT-2 size: 124.2M parameters

尝试改变 GPT-2 的默认配置，以适配定制数据集、Tokenizer以及训练需求，覆写配置信息。

1. 初始化模型时，vocab_size需定义为tokenizer的长度，由于模型嵌入层Embedding Layer是一个矩阵，形状是[vocab_size, hidden_size]。因此如果vocab_size还是默认值可能会存在内存浪费 （len(tokenizer < vocab_size)、索引越界 （len(tokenizer > vocab_size)等错误；
2. gpt2的开始、结束标志和tokenizer中的不一致，可能会导致模型持续输出而不停止；

from transformers import DataCollatorForLanguageModeling

tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

1. 使用DataCollatorForLanguageModeling作为填充器，可作为掩码模型mlm以及因果模型clm。
2. gpt2比较类似于上述提到的文本拼接训练，使用一个eos_token_id拼接后，丢弃一部分 不足max_length的bacth，但是这里训练是进行填充 至max_length，因此需要定义pad_token。

使用Trainer API微调模型

from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir="codeparrot-ds",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    eval_strategy="steps",
    eval_steps=5_000,
    logging_steps=5_000,
    gradient_accumulation_steps=8,
    num_train_epochs=1,
    weight_decay=0.1,
    warmup_steps=1_000,
    lr_scheduler_type="cosine",
    learning_rate=5e-4,
    save_steps=5_000,
    fp16=True,
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=data_collator,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["valid"],
)

trainer.train()

Trainer中的所有step参数，都是指优化步数（权重更新次数），因此：

• Mini-batch：GPU一次性读取的批次数per_device_train_batch_size为32；
• gradient_accumulation_steps：攒8*32个数据后再进行梯度回传、参数更新；
• 因此完成这8*32=256条数据处理才算是一个step；

对应的warmup_steps设置为1000，则需要处理完1000*8*32数据后 学习率才会攀升至5e-4。

使用Accelerate进行微调

使用accelerate微调时，可以加入更多定制化的细节。

例如，想要训练结果更多关注常用库，可以对包含常用库的样本添加更多的关注，通过人为放大 这条样本的 Loss，迫使模型在反向传播时，优先优化 这些包含关键词的数据，从而更大概率学会生成这些特定的词。

# 判断样本中是否包含关键词，包含几个
keytoken_ids = []
for keyword in [
    "plt",
    "pd",
    "sk",
    "fit",
    "predict",
    " plt",
    " pd",
    " sk",
    " fit",
    " predict",
    "testtest",
]:
    ids = tokenizer([keyword]).input_ids[0]
    if len(ids) == 1:
        keytoken_ids.append(ids[0])
    else:
        print(f"Keyword has not single token: {keyword}")

自定义损失函数

from torch.nn import CrossEntropyLoss
import torch


def keytoken_weighted_loss(inputs, logits, keytoken_ids, alpha=1.0):
    # Shift so that tokens < n predict n
    shift_labels = inputs[..., 1:].contiguous()
    shift_logits = logits[..., :-1, :].contiguous()
    # Calculate per-token loss
    loss_fct = CrossEntropyLoss(reduce=False)
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    # Resize and average loss per sample
    loss_per_sample = loss.view(shift_logits.size(0), shift_logits.size(1)).mean(axis=1)
    # Calculate and scale weighting
    weights = torch.stack([(inputs == kt).float() for kt in keytoken_ids]).sum(
        axis=[0, 2]
    )
    weights = alpha * (1.0 + weights)
    # Calculate weighted average
    weighted_loss = (loss_per_sample * weights).mean()
    return weighted_loss

同上，需要自定义缩放Loss：

1. 输入和标签是一个移位关系 ，首先都是copy自例如['我','爱','学习','AI']，对应的输入应该为['我','爱','学习']，而标签应该是['爱','学习','AI']，所以上述代码中一个是取切片[1:]，而另一个是取切片[:-1]；
2. 选择CrossEntropyLoss(reduce=False)，以保留每一个token的损失而非平均值；
3. mini-batch输入形状为[32,128]，展平计算一个32*128的loss，再恢复原始形状并在axis=1，即计算得到一个[32, 1]的loss_per_sample，用于计算每个样本对应的损失；
4. 计算每个样本中的关键词个数，并计入最终损失缩放系数；

weight_decay = 0.1

def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for n, p in model.named_parameters():
        if any(nd in n for nd in no_decay):
            params_without_wd.append(p)
        else:
            params_with_wd.append(p)
    return [
        {"params": params_with_wd, "weight_decay": weight_decay},
        {"params": params_without_wd, "weight_decay": 0.0},
    ]

在权重衰减weight_decay过程中：

1. 由于bias主要起到平移作用，不需要对其进行衰减；
2. LayerNorm.weight是层归一化权重，用于恢复数据本身分布，不需要衰减；
3. LayerNorm.bias是层归一化偏置，作用同上，也不需要衰减。代码逻辑中任何包含bias的都不需要衰减。

from torch.utils.data.dataloader import DataLoader

tokenized_datasets.set_format("torch")
train_dataloader = DataLoader(tokenized_datasets["train"], batch_size=32, shuffle=True)
eval_dataloader = DataLoader(tokenized_datasets["valid"], batch_size=32)

# 设置评估函数，返回困惑度
def evaluate():
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch["input_ids"], labels=batch["input_ids"])

        losses.append(accelerator.gather(outputs.loss))
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss)
    except OverflowError:
        perplexity = float("inf")
    return loss.item(), perplexity.item()

# 初始化模型配置
model = GPT2LMHeadModel(config)

# 设置优化器
from torch.optim import AdamW
optimizer = AdamW(get_grouped_params(model), lr=5e-4)

# 配置accelerate
from accelerate import Accelerator
accelerator = Accelerator(fp16=True)
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

# 设置优化器衰减速度
from transformers import get_scheduler
num_train_epochs = 1
num_update_steps_per_epoch = len(train_dataloader)
num_training_steps = num_train_epochs * num_update_steps_per_epoch

lr_scheduler = get_scheduler(
    name="linear",
    optimizer=optimizer,
    num_warmup_steps=1_000,
    num_training_steps=num_training_steps,
)

# 配置上传仓库信息
from huggingface_hub import Repository, get_full_repo_name

model_name = "codeparrot-ds-accelerate"
repo_name = get_full_repo_name(model_name)
output_dir = "codeparrot-ds-accelerate"
repo = Repository(output_dir, clone_from=repo_name)

# 主训练流程
from tqdm.notebook import tqdm

gradient_accumulation_steps = 8
eval_steps = 5_000

model.train()
completed_steps = 0
for epoch in range(num_train_epochs):
    for step, batch in tqdm(
        enumerate(train_dataloader, start=1), total=num_training_steps
    ):
        logits = model(batch["input_ids"]).logits
        loss = keytoken_weighted_loss(batch["input_ids"], logits, keytoken_ids)
        if step % 100 == 0:
            accelerator.print(
                {
                    "samples": step * samples_per_step,
                    "steps": completed_steps,
                    "loss/train": loss.item() * gradient_accumulation_steps,
                }
            )
        loss = loss / gradient_accumulation_steps
        accelerator.backward(loss)
        if step % gradient_accumulation_steps == 0:
            accelerator.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            completed_steps += 1
        if (step % (eval_steps * gradient_accumulation_steps)) == 0:
            eval_loss, perplexity = evaluate()
            accelerator.print({"loss/eval": eval_loss, "perplexity": perplexity})
            model.train()
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
            if accelerator.is_main_process:
                tokenizer.save_pretrained(output_dir)
                repo.push_to_hub(
                    commit_message=f"Training in progress step {step}", blocking=False
                )

#重要

当step运行到gradien_accumulation_step倍数的时候：

1. accelerator.clip_grad_norm_(model.parameters(), 1.0)进行梯度裁剪，超过1的梯度都会被强制缩放为1.0，防止梯度过大出现爆炸情况；
2. 梯度累积8个steps后，才进行optimizer.step()梯度更新；
3. completed_step是模型权重更新的真实步长，而step是Mini-batch次数；

当执行到指定步长时，开始评估并保存：

1. 这里指定的是eval_steps*gradient_accumulation_steps，也可以换算成completed_steps，目的都是需要在完成梯度更新的步长整数倍数条件下进行保存；
2. model.train()是由于评估时通常会开启model.eval()，需再次开启训练模型才能学习；

!NOTE 梯度累积的优势

核心好处：用显存很小的显卡，训练出大显存显卡的效果。（以时间换空间）

1. 突破显存限制：假设显卡只有 16GB 显存。设置 batch_size = 256，显存爆炸程序崩溃；设置 batch_size = 32 才能跑起来。通过 gradient_accumulation_steps = 8，在逻辑上等效于用了 batch_size = 256；
2. 提高训练稳定性：小 Batch Size (比如 32)：样本太少，噪音很大 。可能这 32 条数据比较特殊，导致计算出的梯度方向是偏的。模型权重会"东跳西窜"，难以收敛到最优解 ；大 Batch Size (比如 256)：样本多了噪音被平均掉，梯度的方向更准确，指向最优解的路径更平滑。

问题

总结