tsuyu122/usaf

Source: https://github.com/tsuyu122/usaf

USAF — Ultra Sparse Adaptive Fine-Tuning

Fine-tune MoE models on hardware that can barely run inference.

Qwen3-30B-A3B needs 60GB in fp16. Full fine-tuning needs 120GB+. USAF trains 26M out of 4.8B parameters on a 12GB GPU — the only method that works on AMD and the only one that trains expert weights and the router.

Why This Exists

I don’t have an A100, an H100, or even an RTX 4090. I have a Radeon RX 6750 XT with 12GB. On Windows.

Every existing fine-tuning method either won’t load on this hardware or won’t touch the parts of MoE models that actually matter. So I built something that does both.

Comparison

Qwen3-30B-A3B, 180 steps. LoRA/QLoRA/DoRA numbers are estimates — no public benchmarks exist for these methods on this model at this scale. Where a method can’t run, I explain why.

LoRA and QLoRA train adapter matrices on frozen weights. USAF trains the actual expert weights and router — it just picks which ones matter. For MoE models, the gate determines model behavior more than any single expert weight.

Why USAF Takes Longer on Big GPUs

On an A100, USAF is slower per-step because it does more work:

USAF computes gradients for 26M parameters per step vs ~100K for LoRA — 260× more gradient work. That it’s only 2-3× slower is the entire point of sparse training.

On consumer hardware, the comparison is simpler: USAF runs. LoRA doesn’t.

Results

180 steps on Qwen3-30B-A3B, RX 6750 XT 12GB (AMD), DirectML.

Held-out repositories (Flecs, SFML, EnTT, Box2D) improved alongside training data — generalization, not memorization.

Why Sparse Training Works for MoE

Not all weights matter. MoE models route each token to a handful of experts. Most weights never activate for a given input. The importance phase finds the 0.5% with highest gradient magnitude.

The router is leverage. Training the gating network (2M parameters) changes which experts fire. A single step drops loss by 0.65. Adapter methods can’t touch the router.

Sparsity adapts. RigL reselection replaces underperforming weights every 50 steps. The active set evolves — turnover starts at ~92% and drops as the model converges.

Resident caching kills the bottleneck. 4-bit dequantization is slow on CPU (400ms per tensor). Trainable layers keep fp16 copies in RAM — dequant once, use forever.

Quick Start

pip install transformers safetensors psutil

# AMD GPU (DirectML)
python train.py

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

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

No config files. Everything via environment variables.

Performance

CUDA numbers are estimates pending real hardware benchmarks.

Supported Models

Auto-detection works for any MoE model from HuggingFace — config.json is all it needs. Tested on Qwen3-30B-A3B.

Models I Want to Test

These are the models USAF was designed for. I just don’t have the GPUs.

Hardware needed: 4-8× A100 80GB or equivalent per model. If you have access and want to see USAF results on these, reach out via GitHub Discussions. I’ll write the training code — you bring the GPUs.

Universal 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

Features

Hardware

• GPU with 12GB+ VRAM or 32GB RAM (CPU-only)
• AMD: DirectML (Windows, built-in)
• NVIDIA: CUDA 11.8+
• Python 3.10+, PyTorch 2.0+

Using Your Own Model

Step 1: Prepare the dataset

Create a JSONL file with tokenized sequences. Each line must have input_ids and labels:

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

To tokenize your own text with the model’s tokenizer:

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

text = "Your training text here..."
tokens = tokenizer.encode(text)
# Chunk into 512-token segments
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]}
    # Write sample to JSONL

Step 2: Quantize the expert weights

USAF needs the expert weights in 4-bit HQQ format. Currently supports Qwen3-MoE out of the box. For other models, you need to generate the experts_q4.pt file:

from usaf.quantization import quantize_4bit
import torch

# Load your model's expert tensors (gate_up_proj and down_proj for each layer)
q_dict = {}
for layer_idx in range(num_layers):
    for param_name in ["gate_up_proj", "down_proj"]:
        # Load the fused expert tensor [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")

Step 3: Configure and run

# Set these environment variables for your model
QUANT_PATH="my-model-q4/experts_q4.pt"   # Path to quantized weights
TRAIN_FROM=36                            # First trainable layer (keep top layers)
STEPS=360                                # 2 epochs for ~190K tokens
FRAC=0.005                               # 0.5% sparsity
MICROBATCH=2                             # Batch size (increase if VRAM allows)

python train.py

Environment Variables Reference

Supported GPU Configurations

Troubleshooting

“CUDA out of memory”: Reduce MICROBATCH to 1 or increase TRAIN_FROM to freeze more layers.

“No module named torch_directml” on NVIDIA: Expected. The code auto-detects and uses CUDA. Set USE_CUDA=1.

Loss not decreasing: Ensure FRAC is high enough (>0.001). Try 2-3 epochs with EPOCHS=3. Check dataset quality.

Frozen cache takes too long: Set FROZEN_CACHE_N=50 to only cache the first 50 samples. Or disable with USE_FROZEN_CACHE=0.

Future Work

• Benchmarks against LoRA/QLoRA/DoRA on A100-class hardware
• Full Vulkan attention pipeline for cross-vendor acceleration
• Distributed training (FSDP)
• Tests on DeepSeek-V4 Pro, Kimi K2.5, Mistral Large 3 — need hardware

License

Apache 2.0. LICENSE. Contributions: CLA.