8卡微调Grok-1实战教程

        本文是根据魔搭社区推出的轻量级训练推理工具SWIFT微调实战教程。SWIFT（Scalable lightWeight Infrastructure for Fine-Tuning）是一套基于PyTorch的轻量级、开箱即用的模型微调、推理框架，让AI爱好者能够轻松地在消费级显卡上运行大模型和AIGC。

        Grok-1是一个具有314B参数的基础模型，由于采用了Rust+JAX框架构建，与transformers生态不兼容，导致使用其进行微调训练的成本较高。为了解决这个问题，Colossal-AI及时推出了解决方案，提供了一个更加方便易用的Python+PyTorch+HuggingFace版本的Grok-1——grok-1-pytorch。目前，该模型已在HuggingFace和ModelScope上开源。

魔搭SWIFT模型训练框架已经率先支持对grok-1-pytorch在8卡环境下的微调，并且已经将训练脚本完全开源。这使得用户能够更方便地进行模型的微调和训练。

01 环境准备

git clone https://github.com/modelscope/swift.git
cd swift
pip install -e .[llm]

02 微调

02 微调

实验环境

• GPU:8*A100 80G

• 镜像:ModelScope官方镜像1.13.1版本

• peft:0.10.0

数据集准备

Grok是base模型，因此我们使用了问题生成数据集DuReader作为训练集。该数据集约15000条，max-length设置为512，训练数据约10000条（平均长度305±92 tokens）。

模型准备

Grok模型我们使用了ColossalAI提供的版本，其中我们额外准备了符合transformers标准的tokenizer。

模型链接：

• https://www.modelscope.cn/models/colossalai/grok-1-pytorch/summary

• https://www.modelscope.cn/models/AI-ModelScope/grok-1-tokenizer/summary

训练

由于Grok模型过大，device_map和deepspeed zero3非offload均无法运行训练，因此本次实验我们使用了LoRA+deepspeed zero3 offload模式运行训练。训练完整脚本如下：

# cd examples/pytorch/llm first
nproc_per_node=8

PYTHONPATH=../../.. \
torchrun \
    --nproc_per_node=$nproc_per_node \
    --master_port 29500 \
    llm_sft.py \
    --model_type grok-1 \
    --sft_type lora \
    --tuner_backend swift \
    --dtype bf16 \
    --output_dir output \
    --ddp_backend nccl \
    --dataset dureader-robust-zh \
    --train_dataset_sample -1 \
    --num_train_epochs 1 \
    --max_length 512 \
    --check_dataset_strategy warning \
    --lora_rank 8 \
    --lora_alpha 32 \
    --lora_dropout_p 0.05 \
    --lora_dtype bf16 \
    --lora_target_modules DEFAULT \
    --gradient_checkpointing true \
    --batch_size 2 \
    --weight_decay 0.1 \
    --learning_rate 1e-4 \
    --gradient_accumulation_steps $(expr 16 / $nproc_per_node) \
    --max_grad_norm 0.5 \
    --warmup_ratio 0.03 \
    --eval_steps 100 \
    --save_steps 100 \
    --save_total_limit 2 \
    --logging_steps 10 \
    --deepspeed_config_path scripts/grok-1/lora_ddp_ds/zero3.json \
    --save_only_model true \

改脚本需要一个zero3.json文件，完整的训练文件可以在这里找到。

下面是训练过程的一些benchmark

显存占用没到24G，按道理4090也可以，训练时长约4小时。

03 推理   

    SWIFT框架目前还不支持使用DeepSpeed进行推理，因此我们继续依靠Transformers库中的device_map功能来实现分布式推理。鉴于模型的规模庞大，某些层（layers）会被卸载到CPU上运行，这会影响到LoRA（Low-Rank Adaptation）的加载并导致推理过程中出现错误。为了应对这一问题，我们对PEFT（Performance-Efficient Tuning）的实现进行了特定的修改：当原始的Linear层在元设备（meta device）上时，我们不会迁移LoRA参数，而是在运行时动态地迁移权重（weights）。这样的调整有助于确保模型即使在部分层被卸载到CPU的情况下也能正确执行推理。

 

# cd examples/pytorch/llm first
PYTHONPATH=../../.. \
python llm_infer.py \
    --ckpt_dir output/grok-1/vx-xxx-xxx/checkpoint-xxx \
    --dtype bf16 \
    --load_dataset_config true \
    --max_new_tokens 64 \
    --do_sample true \
    --dtype bf16 \
    --eval_human false \
    --merge_lora false \

结果

[PROMPT]Task: Question Generation
Context: 我个人感觉是吕颂贤版,剧情和原著差别不大,虽然TVB演员颜值和风光没有大陆的好。但是香港特区人口和地域的限制,只能注重在演员的演技方面发挥很出色,楼主看过大陆排《笑傲江湖》吧!在台词上表现的很生硬没有香港的注重神色配台词,比如杜燕歌把吕颂贤表情和性格几乎和原著差别不大。武打几乎沿用徐克和程小东动作的风格很注重实际技巧,没有大陆版的在武打场面依靠电脑特效表现的太夸张了。李亚鹏版的武打动作和导演还是香港的元彬,大陆毕竟还是在武侠剧起步的比较晚,主要是还是靠明星大腕压阵而香港却是恰恰相反。
Question:[OUTPUT]笑傲江湖哪个版本好看</s>
Answer: 吕颂贤版