可视化神器WandB，大模型训练的必备工具！

前言

最近工作中使用Unsloth 进行大模型微调。与LLamaFacotry 相比，Unsloth 更适合单卡微调大模型，在消耗更小显存的同时具备更高效的微调性能。然而Unsltoh 微调也有瑕疵，它不具备可视化界面，不能像LLamaFactory那样将训练损失loss用折线图表示出来，虽然Unsloth也会打印日志记录loss值，但这样并不能直观体现loss的变化趋势，不利于后续的研究分析。

当然我们完全可以自己编写函数对日志中的loss值等关键节点进行分析，但是我不只想监控Loss，还想监控学习率的变化，grad_norm的变化等，总不能都从日志中扒关键节点分析吧（归根结底就是懒，老想使用现成工具），有没有什么工具可以使用呢？大家还别说，真让笔者给找到了，大模型训练的可视化神器WandB（全名：Weight and Bias）, 帮我们实现大模型训练中的数据可视化需求！

大家对微调、Unsloth 和LLamaFactory 有疑惑的，可以先参考笔者以前的文章。笔者以后也会分享从数据集构建到模型微调到微调结果评估的全流程分享，感兴趣可以关注笔者的掘金号或同名微信公众号：大模型真好玩。

一、WandB基本说明

在大规模模型训练中，我们往往需要监控和分析大量的训练数据，而WandB可以帮助我们实现这一目标。WandB提供了以下重要的功能:

实时可视化: WandB可以实时展示训练过程中关键指标的变化，如损失函数、学习率、训练时间等。通过这些可视化数据，我们能够直观地了解模型的训练进展，快速发现训练中的异常或瓶颈。
自动记录与日志管理: WandB会自动记录每次实验的参数、代码、输出结果，确保实验结果的可追溯性。无论是超参数的设置，还是模型的架构调整，WandB都能够帮助我们完整保留实验记录，方便后期对比与调优。
支持中断与恢复训练： 在长时间的预训练任务中，系统中断或需要暂停是常见的情况。通过WandB 的checkpoint功能，我们可以随时恢复训练，从上次中断的地方继续进行，避免数据和时间的浪费。
多实验对比： 当我们尝试不同的模型配置或超参数时，WandB允许我们在多个实验之间轻松进行对比分析，帮助我们选择最优的模型配置。
团队协作： WandB还支持团队协作，多个成员可以共同查看实验结果，协同调试模型。这对研究和项目开发中团队的合作非常有帮助。

WandB的功能如此强大，那我们该如何注册使用呢？

二、WandB注册流程

首先登录WandB的官方网站，点击右上角 Sign Up注册账号，注册账号需要用到邮箱，我这里使用QQ邮箱进行注册:

注册过程中可以随意填写Full Name 、Company or Institution 字段，但要注意一定要选择Academic选项，否则不免费：

给自己的工作组创建名字，随便起一个即可：

注册成功后WandB会提示你还有一个API Key，一定要记录好这个API-KEY，方便在之后使用该API KEY 验证身份，从而把你在各个服务器上运行的大模型训练中间结果的上传方便进行可视化观察。

在你Unsloth 的虚拟环境下执行pip install wandb安装WandB库。

三、WandB使用指南

WandB 与Unsloth 已经做到高度集成，我们可以很简单的在Unsloth 训练代码中插入WandB 的相关部分，从而让Unsloth的训练流程可视化展示：

在Unsloth微调过程中我们需要创建有监督的微调对象, 然后开启训练，代码如下:

python 复制代码

trainer = SFTTrainer( # 创建有监督微调对象
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        # Use num_train_epochs = 1, warmup_ratio for full training runs!
        warmup_steps=5,
        max_steps=60,
        learning_rate=2e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=10,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
    ),
) 

##----------WandB代码部分-------
import wandb
wandb.login(key="你注册的API-KEY")
##----------WandB代码部分-------

trainer_stats = trainer.train()

只需要加入两行代码，wandb就会自动收集Unsloth的日志并绘制图表，以上代码的执行结果如下，我们可以之后登录View Project显示的网址查看我们绘制的损失曲线等图片，

绘制结果如下图所示，可以看到绘制了包括learning_rate、global_step、grad_norm、epoch和loss在内的五个图片，因为该项目仅仅作为演示没有完成整个训练流程。这里有一个小技巧：大家在训练模型时最好将最后的loss控制在0.4-0.6的范围，这样在保证不过分损失模型原有性能的基础上提升模型专业领域的性能。

四、结语

本篇文章简单分享了可视化神器WandB的使用方法，大家在日常使用中可以尝试用它来解决Unsloth损失可视化的问题。笔者最近在完成大模型微调的相关需求，也总结出很多实践经验，计划分享从数据集构建到模型训练到模型评估的一整套流程，大家感兴趣点个关注吧，也可关注笔者同名微信公众号：大模型真好玩，查阅更多笔者学习工作中免费的大模型经验分享和相关资料~