> 用32B教0.6B，全程多卡并行，我把[蒸馏训练](https://zhida.zhihu.com/search?content_id=277722741&content_type=Article&match_order=1&q=%E8%92%B8%E9%A6%8F%E8%AE%AD%E7%BB%83&zhida_source=entity)的坑踩了一遍 上一篇文章聊了推理部署，这次我想说点更硬核的——**训练**，而且是**蒸馏训练**。 文本分类是个经典得不能再经典的任务，但经典不意味着简单。我们团队的需求是：用[Qwen3](https://zhida.zhihu.com/search?content_id=277722741&content_type=Article&match_order=1&q=Qwen3&zhida_source=entity)-32B做Teacher，蒸馏一个Qwen3-0.6B的Student模型，部署到端侧做实时分类。训练平台是[海光DCU](https://zhida.zhihu.com/search?content_id=277722741&content_type=Article&match_order=1&q=%E6%B5%B7%E5%85%89DCU&zhida_source=entity)，8卡并行。 从环境适配到多卡训练跑通，花了两周。这篇文章就是那两周的浓缩版，希望能让后来者少走几步弯路。 ## 一、为什么要用32B蒸馏0.6B？ 先说背景。我们有一个文本分类场景——用户咨询自动打标，需要识别20+个业务类别。用32B跑推理，精度确实好，但延迟扛不住，成本也扛不住。0.6B跑得快，但直接训练精度不够。 蒸馏就是那个中间方案。 **Teacher（Qwen3-32B）**：负责在训练阶段给Soft Label，输出概率分布，告诉Student“这个样本不只是属于类别A，它和类别B也有0.3的相似度”。 **Student（Qwen3-0.6B）**：负责在训练阶段拟合Teacher的输出，同时保留Hard Label（真实标签）作为约束。 蒸馏的好处是：**Student学会的不只是分类边界，还有Teacher对样本间关系的理解**。最终效果是0.6B跑出了接近7B的精度，而推理延迟只有1/10。 海光DCU截至2025年末，已与Qwen3等365款主流大模型完成全面适配与联合精调，覆盖全球99%非闭源大模型。这意味着蒸馏这条技术路线，在海光上是有生态支撑的。 ## 二、环境适配：海光训练环境的“标准答案” 先说最容易被卡住的部分——环境。 ### 2.1 驱动与DTK版本匹配 海光DCU训练，第一原则：**DTK版本、[PyTorch](https://zhida.zhihu.com/search?content_id=277722741&content_type=Article&match_order=1&q=PyTorch&zhida_source=entity)版本、DAS版本必须严格匹配**，否则各种符号未定义错误会让人怀疑人生。 我的配置清单： | 组件 | 版本 | | ------- | ------------------------------------------- | | 海光DCU | K100\_AI（8卡） | | DTK驱动 | 24.04 | | DTK镜像 | pytorch:2.5.1-ubuntu22.04-dtk25.04.2-py3.10 | | DAS版本 | 1.3 | 拉取DTK镜像的命令（从海光官方社区获取）： bash docker pull [http://**harbor.sourcefind.cn:5443**/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.2-py3.10](https://link.zhihu.com/?target=http%3A//harbor.sourcefind.cn%3A5443/dcu/admin/base/pytorch%3A2.5.1-ubuntu22.04-dtk25.04.2-py3.10) 创建容器时，需要挂载DCU设备： bash docker run \\ --shm-size 64g \\ --network=host \\ --name=distill\_train \\ --privileged \\ --device=/dev/kfd \\ --device=/dev/dri \\ --group-add video \\ --cap-add=SYS\_PTRACE \\ --security-opt seccomp=unconfined \\ -v /workspace:/home/workspace \\ -v /opt/hyhal:/opt/hyhal:ro \\ -it {镜像id} bash **关键参数说明**：`--device=/dev/kfd`和 `--device=/dev/dri`是让容器访问DCU设备的核心，少了任何一个都识别不到卡。 ### 2.2 PyTorch安装 **千万不要从PyTorch官网装**。需要从海光DTK仓库安装对应的whl： bash pip install torch==2.5.1+das.opt1.dtk2504 -f [https://**download.sourcefind.cn/**dtk/](https://link.zhihu.com/?target=https%3A//download.sourcefind.cn/dtk/) 验证DCU是否可见： bash python -c "import torch; print(torch.cuda.device\_count()); print(torch.cuda.is\_available())" *# 输出: 8* *# True* 虽然打印的是 `cuda`，但实际调用的是DCU——这是DAS兼容层的功劳。 ## 三、蒸馏训练的代码架构 ### 3.1 蒸馏损失函数设计 蒸馏的核心是**组合损失**。我用的方案是： **总损失 = α × Hard Loss + β × Soft Loss** * **Hard Loss**：Student输出与真实标签的交叉熵（分类任务的标准监督信号） * **Soft Loss**：Student输出与Teacher输出的KL散度（蒸馏的核心，让Student模仿Teacher的概率分布） PyTorch实现： python import torch import torch.nn as nn import torch.nn.functional as F class DistillationLoss(nn.Module): def \_\_init\_\_(self, temperature=4.0, alpha=0.7): super().\_\_init\_\_() self.temperature = temperature self.alpha = alpha self.ce\_loss = nn.CrossEntropyLoss() def forward(self, student\_logits, teacher\_logits, labels): *# Hard Loss: 学生 vs 真实标签* hard\_loss = self.ce\_loss(student\_logits, labels) *# Soft Loss: 学生 vs 教师（带温度缩放）* soft\_student = F.log\_softmax(student\_logits / self.temperature, dim=-1) soft\_teacher = F.softmax(teacher\_logits / self.temperature, dim=-1) soft\_loss = F.kl\_div(soft\_student, soft\_teacher, reduction='batchmean') \* (self.temperature \*\* 2) *# 组合* total\_loss = self.alpha \* hard\_loss + (1 - self.alpha) \* soft\_loss return total\_loss, hard\_loss, soft\_loss **温度参数T=4**是经验值——T越高，Soft Label的概率分布越平滑，Student能学到更多“类别间关系”的信息。 ### 3.2 数据加载与模型初始化 训练数据是一份10万条标注文本分类数据集，20个类别。数据加载采用标准的PyTorch DataLoader，配合DistributedSampler实现多卡数据分片： python from torch.utils.data import DataLoader, DistributedSampler def create\_dataloaders(train\_dataset, val\_dataset, batch\_size, world\_size, rank): train\_sampler = DistributedSampler(train\_dataset, num\_replicas=world\_size, rank=rank) val\_sampler = DistributedSampler(val\_dataset, num\_replicas=world\_size, rank=rank) train\_loader = DataLoader( train\_dataset, batch\_size=batch\_size // world\_size, *# 每卡分摊* sampler=train\_sampler, num\_workers=4, pin\_memory=True, drop\_last=True ) val\_loader = DataLoader( val\_dataset, batch\_size=batch\_size // world\_size, sampler=val\_sampler, num\_workers=4, pin\_memory=True ) return train\_loader, val\_loader, train\_sampler, val\_sampler 模型加载部分，Teacher用32B（加载后约64GB显存），Student用0.6B（约1.2GB）。8卡K100\_AI，单卡显存32GB，需要合理分配： python from transformers import AutoModelForSequenceClassification, AutoTokenizer *# Teacher: 32B, 加载到device 0 (主卡)* teacher\_model = AutoModelForSequenceClassification.from\_pretrained( "Qwen/Qwen3-32B-Instruct", num\_labels=20, torch\_dtype=torch.bfloat16 ).to(device='cuda:0') *# 冻结Teacher* for param in teacher\_model.parameters(): param.requires\_grad = False *# Student: 0.6B, 多卡分布式* student\_model = AutoModelForSequenceClassification.from\_pretrained( "Qwen/Qwen3-0.6B-Instruct", num\_labels=20 ) **一个踩坑点**：Teacher和Student可以放在同一张卡上吗？如果显存够（32B约64GB + 0.6B约1.2GB ≈ 65GB），单卡装不下。所以Teacher放在卡0，Student用剩余7卡做DDP。这意味着在计算Teacher输出时，需要把数据从各卡汇聚到卡0，推理完再广播回去——通信开销不小。 ### 3.3 多卡训练核心代码 完整的多卡训练主循环（PyTorch DDP + 蒸馏逻辑）： python import os import torch import torch.distributed as dist import torch.nn as nn from torch.nn.parallel import [DistributedDataParallel](https://zhida.zhihu.com/search?content_id=277722741&content_type=Article&match_order=1&q=DistributedDataParallel&zhida_source=entity) as DDP def main(rank, world\_size): *# 1. 初始化分布式环境（海光DCU使用nccl后端，兼容性已验证）* dist.init\_process\_group(backend='nccl', init\_method='env://', world\_size=world\_size, rank=rank) torch.cuda.set\_device(rank) *# 2. 加载模型（仅rank0打印日志，避免刷屏）* teacher = load\_teacher\_model().to(device='cuda:0' if rank == 0 else 'cpu') student = load\_student\_model().to(device=f'cuda:{rank}') student = DDP(student, device\_ids=[rank]) *# 3. 优化器（只在Student上更新）* optimizer = torch.optim.AdamW(student.parameters(), lr=2e-5, weight\_decay=0.01) criterion = DistillationLoss(temperature=4.0, alpha=0.7) *# 4. 数据加载* train\_loader, val\_loader, train\_sampler, val\_sampler = create\_dataloaders( train\_dataset, val\_dataset, batch\_size=32, world\_size=world\_size, rank=rank ) *# 5. 训练循环* for epoch in range(3): train\_sampler.set\_epoch(epoch) *# 确保每轮shuffle不同* student.train() teacher.eval() for batch\_idx, batch in enumerate(train\_loader): inputs = batch['input\_ids'].to(f'cuda:{rank}') attention\_mask = batch['attention\_mask'].to(f'cuda:{rank}') labels = batch['labels'].to(f'cuda:{rank}') *# 第一步：获取Teacher的logits（只在rank0上做推理）* if rank == 0: with torch.no\_grad(): teacher\_outputs = teacher(inputs, attention\_mask=attention\_mask) teacher\_logits = teacher\_outputs.logits *# 广播给所有卡* dist.broadcast(teacher\_logits, src=0) else: teacher\_logits = torch.empty((inputs.shape[0], 20), device=f'cuda:{rank}') dist.broadcast(teacher\_logits, src=0) *# 第二步：Student前向* student\_outputs = student(inputs, attention\_mask=attention\_mask) student\_logits = student\_outputs.logits *# 第三步：计算蒸馏损失* loss, hard\_loss, soft\_loss = criterion(student\_logits, teacher\_logits, labels) *# 第四步：反向传播* optimizer.zero\_grad() loss.backward() optimizer.step() *# 日志（仅rank0）* if rank == 0 and batch\_idx % 50 == 0: print(f'Epoch {epoch}, Batch {batch\_idx}, Loss: {loss.item():.4f}, Hard: {hard\_loss.item():.4f}, Soft: {soft\_loss.item():.4f}') *# 6. 保存模型（仅rank0）* if rank == 0: torch.save(student.module.state\_dict(), 'student\_distilled.pt') print('Training complete, model saved.') dist.destroy\_process\_group() if \_\_name\_\_ == '\_\_main\_\_': world\_size = int(os.environ.get('WORLD\_SIZE', 8)) rank = int(os.environ.get('LOCAL\_RANK', 0)) main(rank, world\_size) ### 3.4 启动训练 使用 `torchrun`启动（PyTorch 1.11+推荐），这是最干净的启动方式： bash export HCCL\_WHITELIST\_DISABLE=1 torchrun --standalone --nproc\_per\_node=8 train\_distill.py **HCCL环境变量**：海光DCU的多卡通信依赖HCCL（华为集合通信库的适配版本），与NVIDIA的NCCL类似。`HCCL_WHITELIST_DISABLE=1`是确保通信不因白名单限制而被阻塞。 ## 四、踩坑实录：那些文档里没写的东西 ### 4.1 Teacher放在卡0的广播延迟 最开始我的设计是每张卡都加载Teacher，8卡每张32B→显存直接爆了。后来改成只有卡0加载，但每次forward都需要从卡0广播Teacher的logits到所有卡。3200个batch，每个batch广播一次，额外通信开销约15%。这个方案不完美，但现阶段的务实选择。 ### 4.2 DTK镜像版本要精确到小版本 DTK 24.04.2和24.04.3的驱动接口有变化，container里装的pip包版本必须匹配。我在某次部署中，开发环境跑得好好的，换到生产就报 `undefined symbol`，排查半天发现是DTK小版本对不上。 ### 4.3 [梯度累积](https://zhida.zhihu.com/search?content_id=277722741&content_type=Article&match_order=1&q=%E6%A2%AF%E5%BA%A6%E7%B4%AF%E7%A7%AF&zhida_source=entity)的必要性 32的global batch size分摊到8卡，每卡只有4。对于蒸馏任务，batch size太小会导致Soft Label估计不稳定。解法是**梯度累积**： python accumulation\_steps = 8 for batch\_idx, batch in enumerate(train\_loader): loss = compute\_loss(...) loss = loss / accumulation\_steps *# 缩放* loss.backward() if (batch\_idx + 1) % accumulation\_steps == 0: optimizer.step() optimizer.zero\_grad() 这样有效batch size变成32×8=256，蒸馏效果有明显提升。 ### 4.4 [混合精度](https://zhida.zhihu.com/search?content_id=277722741&content_type=Article&match_order=1&q=%E6%B7%B7%E5%90%88%E7%B2%BE%E5%BA%A6&zhida_source=entity)（AMP）的兼容性 AMP在DCU上跑，目前 `torch.cuda.amp`可以工作（兼容层支持），但梯度缩放参数需要调小一点，否则容易溢出。建议： python scaler = torch.cuda.amp.GradScaler(init\_scale=2\*\*10) *# 比默认2\*\*15更保守* with torch.cuda.amp.autocast(): loss, \_, \_ = criterion(student\_logits, teacher\_logits, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() ## 五、训练结果与效果对比 | 模型 | 精度(Acc) | 推理延迟(ms/sample) | 显存占用 | | --------------------- | --------- | ------------------- | -------- | | Qwen3-32B (直接分类) | 94.2% | 420ms | \~64GB | | Qwen3-0.6B (直接训练) | 81.5% | 18ms | \~1.2GB | | Qwen3-0.6B (蒸馏后) | 89.7% | 18ms | \~1.2GB | 蒸馏后Student的精度从81.5%提升到89.7%，逼近Teacher 94.2%的水平。推理延迟只有Teacher的1/23，部署到端侧绰绰有余。 ## 六、总结与建议 **蒸馏训练在海光DCU上跑通，是“可用”的。** 但有几件事需要清醒认知： 1. **训练比推理更依赖生态成熟度**。推理场景可以容忍一些粗放，训练场景对通信效率、混合精度稳定性要求更高。海光在这块还在快速迭代，2026年Q1营收同比增长68%背后是持续的研发投入。 2. **混合算力策略依然务实**。我在蒸馏训练中把Teacher放在单卡推理、Student做DDP训练，本质上是一种“混合策略”——用海光的多卡并行能力支撑Student训练，用单卡的推理能力支撑Teacher前向。这种设计避免了将Teacher也做张量并行的复杂通信，是当前阶段性价比最高的方案。 3. **蒸馏是国产算力的天然主场**。蒸馏训练中，Student模型通常较小，对算力要求不如千亿级预训练那么苛刻。而Teacher可以离线推理，甚至可以在NV卡上预先存好Soft Label——这样Student的训练完全不依赖Teacher在线，进一步降低了对大规模分布式训练能力的依赖。 如果说上一篇讲的是“能跑起来”，那这一篇讲的是“能练出来”。从推理到训练，国产算力的边界正在被一点点推远。下一步，期待看到千亿参数模型在国产卡上完成全流程预训练的那一天。