边缘 AI 架构设计：在手机和 IoT 上运行 LLM

# 使用 llama.cpp 的量化工具链
# 1. 转换为 GGUF 格式
# python convert.py models/llama-3-8b --outtype f16 --outfile llama-3-8b-f16.gguf

# 2. 量化为 INT4
# ./quantize llama-3-8b-f16.gguf llama-3-8b-q4_k_m.gguf Q4_K_M

# 在 Python 中使用 llama-cpp-python 加载
from llama_cpp import Llama

llm = Llama(
    model_path="llama-3-8b-q4_k_m.gguf",
    n_ctx=2048,           # 上下文窗口（移动端建议 2048-4096）
    n_threads=4,          # CPU 线程数
    n_gpu_layers=0,       # 移动端通常 0（纯 CPU）
    n_batch=512,          # 批处理大小
    mmap=True,            # 内存映射加载，减少内存占用
    mmap_timeout=30000,
)

response = llm(
    "解释量子计算的基本原理",
    max_tokens=256,
    temperature=0.7,
    top_p=0.9,
)

# 使用 Hugging Face 进行知识蒸馏
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch.nn.functional as F

teacher = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-70B")
student = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-1B")

teacher.eval()
student.train()

optimizer = torch.optim.AdamW(student.parameters(), lr=5e-5)

for batch in dataloader:
    with torch.no_grad():
        teacher_logits = teacher(batch["input_ids"]).logits
        teacher_probs = F.softmax(teacher_logits / 2.0, dim=-1)  # 温度=2

    student_logits = student(batch["input_ids"]).logits
    student_log_probs = F.log_softmax(student_logits / 2.0, dim=-1)

    # KL 散度损失
    loss = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")
    # 加上 hard label loss
    loss += 0.3 * F.cross_entropy(
        student_logits.view(-1, student_logits.size(-1)),
        batch["labels"].view(-1)
    )

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# 结构化剪枝：移除整个注意力头或层
import torch.nn.utils.prune as prune

def structured_prune(model, amount=0.3):
    """移除 30% 最不重要的注意力头"""
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            # 按 L1 范数剪枝
            prune.l1_unstructured(module, name="weight", amount=amount)
            # 固化剪枝
            prune.remove(module, "weight")
    return model

# 评估剪枝后的精度
def evaluate_pruned(model, eval_dataset):
    pruned = structured_prune(model, amount=0.3)
    # 重新微调恢复精度
    trainer = Trainer(
        model=pruned,
        args=TrainingArguments(
            output_dir="./pruned-finetune",
            num_train_epochs=2,
            learning_rate=2e-5,
        ),
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
    )
    trainer.train()
    return pruned

// Android NDK 集成示例
// JNI 接口：从 Java/Kotlin 调用 llama.cpp
#include "llama.h"
#include <jni.h>
#include <android/log.h>

extern "C" JNIEXPORT jlong JNICALL
Java_com_example_llm_LlamaEngine_nativeInit(JNIEnv* env, jobject thiz, jstring model_path) {
    const char* path = env->GetStringUTFChars(model_path, nullptr);

    llama_backend_init(false);

    llama_model_params model_params = llama_model_default_params();
    model_params.n_gpu_layers = 0;  // 移动端纯 CPU
    model_params.use_mmap = true;

    llama_model* model = llama_load_model_from_file(path, model_params);
    env->ReleaseStringUTFChars(model_path, path);
    return (jlong)model;
}

extern "C" JNIEXPORT jstring JNICALL
Java_com_example_llm_LlamaEngine_nativeGenerate(
    JNIEnv* env, jobject thiz, jlong model_ptr, jstring prompt, jint max_tokens
) {
    llama_model* model = (llama_model*)model_ptr;
    const char* prompt_str = env->GetStringUTFChars(prompt, nullptr);

    llama_context_params ctx_params = llama_context_default_params();
    ctx_params.n_ctx = 2048;
    ctx_params.n_threads = 4;
    ctx_params.n_batch = 512;

    llama_context* ctx = llama_new_context_with_model(model, ctx_params);

    // Tokenize
    std::string text(prompt_str);
    std::vector<llama_token> tokens = llama_tokenize(ctx, text, true);

    // 生成循环
    std::string result;
    for (int i = 0; i < max_tokens; i++) {
        if (llama_decode(ctx, llama_batch_get_one(&tokens.back(), 1, i, 0))) break;

        llama_token new_token = llama_sample_token(ctx, nullptr, 0.8f, 0.9f, 32);
        if (new_token == llama_token_eos(ctx)) break;

        tokens.push_back(new_token);
        result += llama_token_to_piece(ctx, new_token);
    }

    llama_free(ctx);
    env->ReleaseStringUTFChars(prompt, prompt_str);
    return env->NewStringUTF(result.c_str());
}

# MLC-LLM：使用 TVM 编译器优化，支持 Android/iOS
# 编译模型到移动端
"""
# 1. 编译模型
python -m mlc_llm compile \
    --model llama-3-8b \
    --quantization q4f16_1 \
    --target android \
    --output llama-3-8b-q4f16_1-android.tar

# 2. 在 Android 上加载
"""

# Swift (iOS) 示例
"""
import MLCEngine

let engine = MLCEngine()
try await engine.load(
    model: "Llama-3-8B-q4f16_1",
    modelPath: Bundle.main.path(forResource: "llama-3-8b", ofType: "mlc")!
)

let response = try await engine.generate(
    prompt: "写一首关于秋天的诗",
    maxTokens: 256,
    temperature: 0.7
)
"""

# Qualcomm AI Hub：在 Snapdragon NPU 上运行 LLM
"""
import qai_hub as hub

# 1. 编译模型为 Qualcomm NPU 格式
model = hub.upload_model("llama-3-8b-q4.gguf")

# 2. 在 NPU 上推理
job = model.run(
    device=hub.Device("Samsung Galaxy S24"),
    inputs={"prompt": "Hello, how are you?"}
)
result = job.get_result()
"""

class HybridLLMRouter:
    """根据请求复杂度路由到端侧或云端"""

    def __init__(self, local_engine, cloud_client):
        self.local = local_engine  # 端侧 3B 模型
        self.cloud = cloud_client  # 云端 70B 模型
        self.complexity_threshold = 0.5

    async def generate(self, prompt: str, context: dict) -> str:
        complexity = self._estimate_complexity(prompt, context)

        if complexity < self.complexity_threshold and self._local_capable(prompt):
            # 简单任务：端侧处理
            return await self.local.generate(prompt)

        if complexity > 0.8:
            # 复杂任务：云端处理
            return await self.cloud.generate(prompt)

        # 中等任务：端侧生成草稿，云端优化
        draft = await self.local.generate(prompt)
        if self._quality_check(draft):
            return draft
        return await self.cloud.refine(prompt, draft)

    def _estimate_complexity(self, prompt, context):
        score = 0.0
        # 长度因素
        score += min(len(prompt) / 2000, 0.3)
        # 推理深度因素（关键词检测）
        reasoning_keywords = ["分析", "推理", "比较", "为什么", "证明"]
        score += sum(0.15 for kw in reasoning_keywords if kw in prompt)
        # 上下文长度
        score += min(context.get("context_length", 0) / 50000, 0.2)
        return min(score, 1.0)

# 路由策略配置
routing_rules:
  - name: "short_qa"           # 简单问答
    condition: "tokens < 200 and no_reasoning"
    target: "local"             # 端侧 3B
    fallback: "cloud"

  - name: "long_context"        # 长上下文
    condition: "context_tokens > 4096"
    target: "cloud"             # 云端 70B
    fallback: null

  - name: "code_generation"     # 代码生成
    condition: "task_type == 'code'"
    target: "cloud"
    fallback: "local"           # 离线时用端侧

  - name: "streaming_chat"      # 流式聊天
    condition: "streaming == true"
    target: "local"             # 端侧低延迟
    fallback: "cloud"

# 微型 LLM 部署方案（<100M 参数）
# 适合 ESP32-S3 / 树莓派 Zero 等微型设备

# 方案 1：TinyLlama 1.1B 量化到 INT2
# 大小：~300MB，可在 512MB 内存的树莓派上运行

# 方案 2：使用专门设计的小模型
from transformers import AutoModelForCausalLM

# Qwen2.5-0.5B：仅 500M 参数
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-0.5B",
    torch_dtype=torch.float16,
    device_map="cpu"
)

# 量化到 INT4
from bitsandbytes import quantize
model = quantize(model, quant_type="nf4")
# 内存占用：~150MB