不到 200 行代码就能微调 Llama-2！

OpenMMLab 官方账号

发布于 2023-08-21 10:51:06

1.1K00

代码可运行

文章被收录于专栏：OpenMMLabOpenMMLab

运行总次数：0

代码可运行

上周 Meta AI 又搞了个大新闻，发布了他们的第二代大语言模型：Llama-2，并在第一时间开源了训练、推理代码，甚至还提供了官方版和 huggingface 版，在良心开源的同时还考虑到了普通用户，真的是太酷啦。在收到消息的第一时间，我就冲进了官方仓库：llama2-recipes 打算体验一下 llama-2 的训练流程。

第一次体验时，可以看出来 code release 还是有一些仓促的，也遇到了一些小意外，没跑几个 iteration 模型就不再收敛了，在细品了一番代码之后，发现了一个小错误。他们把 epoch based 的 scheduler 按照 step based 方式去更新了，导致学习率下降的过快，没过几个迭代，学习率就趋向于 0 了。于是我就赶紧向官方反馈了这个问题：

https://github.com/facebookresearch/llama-recipes/issues/27

官方也回应神速，立马就在第二天修复了这个问题：

https://github.com/facebookresearch/llama-recipes/pull/28

之前遇到类似问题的小伙伴可以更新一下代码，保证药到病除。在解决这个小问题之后，Llama-2 就能正常训练了，再次给 Meta AI 的开源精神点赞（有图有真相）！

经历了这个小插曲后，我也差不多摸清楚了 Llama-2 的训练流程，正如论文里所说，它是用 FSDP 进行训练的，诶？FSDP，MMEngine v0.8.0 不也支持了 FSDP 训练么，于是乎，我就基于 MMEngine 的新特性实现了 Llama-2 的训练流程。

完整的训练示例见：

https://github.com/open-mmlab/mmengine/tree/main/examples/llama2

（文末点击阅读原文可直达）

实现数据类

ctrl c+ ctrl v 大法好，直接参考 llama-recipe 里 alpaca dataset 的实现。

参考代码请见：https://github.com/facebookresearch/llama-recipes/blob/1e0f8a1fb77b9ddccf649970f632dd606a22bd06/ft_datasets/alpaca_dataset.py#L28

构建 FSDPStrategy

FSDPStrategy 的构造函数会初始化分布式环境、随机种子等环境变量，因此需要放在第一步来做。Strategy 是 MMEngine v0.8.0 新引入的特性，旨在解决大模型训练的一些问题。关于 Strategy 的具体解读，大家可以期待一下后续的文章~

strategy = FSDPStrategy(
    model_wrapper=dict(
        auto_wrap_policy=partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls={LlamaDecoderLayer})),
    state_dict_cfg='full',
    env_kwargs=dict(randomness=dict(seed=42)))

构建 dataloader、model

配置完全照搬官方 repo。需要注意的是，官方 repo 默认启用全量参数的 bf16 训练，无需混合精度训练。

# Prepare model
tokenizer = LlamaTokenizer.from_pretrained(args.checkpoint)
tokenizer.add_special_tokens({'pad_token': '<PAD>'})
model = LlamaForCausalLM.from_pretrained(args.checkpoint)
model.to(torch.bfloat16)
model.train()

# Prepare dataset
train_dataset = AlpacaDataset(
    tokenizer=tokenizer, data_path=args.data_root)
train_dataloader = DataLoader(
    train_dataset,
    batch_size=args.batch_size,
    sampler=DefaultSampler(train_dataset, seed=0),
    collate_fn=default_data_collator,
    drop_last=True)

准备 optimizer 和 scheduler

配置与官方 repo 对齐，使用 AdamW 和 StepLR。然后把 model、scheduler、optimizer 传给 strategy，由他来处理 FSDP 相关的逻辑。

optim_cfg = dict(
    optimizer=dict(type=AdamW, lr=1e-4, weight_decay=0.0),
    accumulative_counts=ORI_BATCH_SIZE / args.batch_size)
scheduler_cfgs = [dict(type=StepLR, step_size=1, gamma=0.85)]
model, optimizer, schedulers = strategy.prepare(
    model,
    optim_wrapper=optim_cfg,
    param_scheduler=scheduler_cfgs,
    dispatch_kwargs=dict(max_iters=max_iters, max_epochs=args.max_epoch))

实现 train-loop

用了 strategy，我们就能脱离 Runner，自由自在地实现训练逻辑啦，和原生的 PyTorch 很像有木有？

for epoch in range(args.max_epoch):
    for idx, inputs in enumerate(train_dataloader):
        # Convert inputs to target device.
        inputs = apply_to(inputs, lambda m: isinstance(m, torch.Tensor),
                          lambda m: m.cuda())

        loss = model(**inputs).loss
        optimizer.update_params(loss)

        max_memory = torch.cuda.max_memory_allocated()
        strategy.logger.info(f'Epoch: {epoch+1}/{args.max_epoch}, '
                             f'Iter: {idx+1}/{epoch_length}, '
                             f'Loss: {loss.item():.3f}, '
                             f'Lr: {optimizer.get_lr()["lr"][0]:.6f} '
                             f'Memory: {max_memory/1e9:.3f}G')
        visualizer.add_scalars({'loss': loss.item()})

        torch.cuda.reset_peak_memory_stats()

    for scheduler in schedulers:
        scheduler.step()

    save_dir = f'{args.output_dir}/epoch_{epoch+1}'
    state_dict = model.state_dict()

    if is_main_process():
        model.save_pretrained(save_dir, state_dict=state_dict)
        tokenizer.save_pretrained(save_dir)