tsuyu122/usaf

来源：https://github.com/tsuyu122/usaf

USAF — 极致稀疏自适应微调

许可协议 Python (https://python.org) CUDA (https://developer.nvidia.com/cuda-downloads) 状态

在几乎只能运行推理的硬件上微调 MoE 模型。

Qwen3-30B-A3B 在 fp16 精度下需要 60GB 显存。全参数微调需要 120GB 以上。USAF 在 12GB GPU 上训练了 4.8B 参数中的 26M——这是唯一能在 AMD 显卡上运行的方法，也是唯一同时训练专家权重和路由器的方案。

为什么需要这个项目

我没有 A100、H100，甚至连 RTX 4090 都没有。我只有一块 Radeon RX 6750 XT，12GB 显存，运行 Windows。

现有的微调方法要么无法在这块硬件上加载，要么根本不会触及 MoE 模型真正重要的部分。所以我构建了一个两者都能兼顾的方案。

对比

以 Qwen3-30B-A3B 为例，训练 180 步。LoRA/QLoRA/DoRA 的数值为估算值——目前没有针对该模型、该规模下这些方法的公开基准测试。如果某方法无法运行，我会说明原因。

LoRA 和 QLoRA 在冻结权重上训练适配器矩阵。USAF 则训练实际的专家权重和路由器——它只是挑选出哪些权重是重要的。对于 MoE 模型，门控网络对模型行为的决定性作用超过任何单个专家权重。

为什么在大 GPU 上 USAF 耗时更长

在 A100 上，USAF 每步耗时更长，因为它做了更多工作：

USAF 每步计算 2600 万参数的梯度，而 LoRA 只计算约 10 万——梯度计算量高出 260 倍。但它仅比 LoRA 慢 2-3 倍，这正是稀疏训练的精髓所在。

在消费级硬件上，对比更简单：USAF 能跑，LoRA 跑不了。

结果

在 Qwen3-30B-A3B 上训练 180 步，RX 6750 XT 12GB (AMD)，DirectML。

保留集仓库（Flecs，SFML，EnTT，Box2D）的指标与训练数据一同改善——这是泛化能力提升，而非记忆。

为什么稀疏训练对 MoE 有效

并非所有权重都重要。 MoE 模型将每个 token 路由到少数几个专家。对于给定输入，大多数权重从未被激活。重要性阶段会找出梯度幅度最大的 0.5%。

路由器是杠杆点。 训练门控网络（200 万参数）会改变哪些专家被激活。仅一步就能使损失下降 0.65。适配器方法无法触及路由器。

稀疏性是自适应的。 RigL 重选机制每 50 步替换表现不佳的权重。活跃集不断演变——初始轮换率约 92%，随着模型收敛而降低。

驻留缓存消除了瓶颈。 在 CPU 上进行 4 位反量化速度很慢（每张张量 400ms）。可训练层在 RAM 中保留 fp16 副本——只需反量化一次，永久使用。

快速开始

pip install transformers safetensors psutil

# AMD GPU (DirectML)
python train.py

# NVIDIA GPU (CUDA)
USE_CUDA=1 USE_AMP=1 python train.py

# 多 GPU
USE_CUDA=1 USE_MULTI_GPU=1 MICROBATCH=4 python train.py

无需配置文件。一切通过环境变量配置。

性能

CUDA 数据为估算值，需实际硬件基准测试确认。

支持的模型

自动检测适用于 HuggingFace 上的任何 MoE 模型——只需 config.json 即可。已在 Qwen3-30B-A3B 上测试通过。

我希望能测试的模型

以下这些模型是 USAF 的设计目标。我只是没有对应的 GPU。

所需硬件：每个模型需要 4-8 块 A100 80GB 或同等规格。如果您有硬件并希望看到 USAF 在这些模型上的结果，请通过 GitHub Discussions (https://github.com/tsuyu122/usaf/discussions) 联系我。我来编写训练代码——您提供 GPU。

通用 CLI

python -m usaf.train --model Qwen/Qwen3-30B-A3B --dataset data.jsonl --steps 180
python -m usaf.train --model mistralai/Mixtral-8x7B --dataset data.jsonl

特性

硬件要求

• 12GB 以上显存 GPU 或 32GB 内存（仅 CPU）
• AMD：DirectML（Windows，内置）
• NVIDIA：CUDA 11.8+
• Python 3.10+，PyTorch 2.0+

使用自定义模型

第一步：准备数据集

创建一个 JSONL 文件，包含分词后的序列。每行必须包含 input_ids 和 labels：

{"input_ids": [1, 2, 3, ..., 512], "labels": [1, 2, 3, ..., 512]}

使用模型的分词器对自定义文本进行分词：

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-30B-A3B")

text = "您的训练文本..."
tokens = tokenizer.encode(text)
# 切分成 512 token 的片段
for i in range(0, len(tokens) - 512, 512):
    chunk = tokens[i:i+512]
    sample = {"input_ids": chunk, "labels": chunk[1:] + [tokenizer.eos_token_id]}
    # 将 sample 写入 JSONL

第二步：量化专家权重

USAF 需要将专家权重以 4 位 HQQ 格式存储。目前默认支持 Qwen3-MoE。对于其他模型，您需要生成 experts_q4.pt 文件：

from usaf.quantization import quantize_4bit
import torch

# 加载模型的专家张量（每层的 gate_up_proj 和 down_proj）
q_dict = {}
for layer_idx in range(num_layers):
    for param_name in ["gate_up_proj", "down_proj"]:
        # 加载融合后的专家张量 [num_experts, intermediate, hidden]
        weights = load_expert_weights(model_path, layer_idx, param_name)
        q4_entry = quantize_4bit(weights, group_size=128)
        q_dict[f"model.layers.{layer_idx}.mlp.experts.{param_name}"] = q4_entry

torch.save(q_dict, "my-model-q4/experts_q4.pt")

第三步：配置并运行

# 为您的模型设置以下环境变量
QUANT_PATH="my-model-q4/experts_q4.pt"   # 量化权重路径
TRAIN_FROM=36                            # 第一个可训练层（保留顶部层）
STEPS=360                                # 对于约 190K token 进行 2 个 epoch
FRAC=0.005                               # 0.5% 稀疏度
MICROBATCH=2                             # 批大小（若显存充足可增大）

python train.py

环境变量参考

支持的 GPU 配置

故障排除

“CUDA out of memory”：将 MICROBATCH 减少至 1，或增加 TRAIN_FROM 以冻结更多层。

NVIDIA 上出现 “No module named torch_directml”：这是预期的。代码会自动检测并使用 CUDA。设置 USE_CUDA=1。

损失不下降：确保 FRAC 足够高（>0.001）。尝试使用 EPOCHS=3 进行 2-3 个 epoch。检查数据集质量。

冻结缓存耗时过长：设置 FROZEN_CACHE_N=50 仅缓存前 50 个样本。或通过 USE_FROZEN_CACHE=0 禁用。

未来工作

• 在 A100 级别硬件上与 LoRA/QLoRA/DoRA 进行基准测试
• 完整的 Vulkan 注意力流水线，实现跨厂商加速
• 分布式训练（FSDP）
• 在 DeepSeek-V4 Pro、Kimi K2.5、Mistral Large 3 上进行测试——需要硬件支持

许可协议

Apache 2.0。LICENSE。贡献：CLA。