TinyML: Развёртывание LLM на edge-устройствах
llmtinymledgedeploymentembedded
Введение
TinyML — это область машинного обучения, ориентированная на запуск моделей на микроконтроллерах и других устройствах с крайне ограниченными ресурсами. Для LLM это означает возможность обработки естественного языка прямо на устройстве без подключения к облаку.
Resource Constraints
| Платформа | RAM | Flash | CPU | Power |
|---|---|---|---|---|
| ESP32 | 520 KB | 4 MB | 240 MHz | 500 mA |
| STM32 | 2 MB | 2 MB | 480 MHz | 300 mA |
| Raspberry Pi Pico 2 | 264 KB | 2 MB | 133 MHz | 200 mA |
| Arduino Nano 33 | 256 KB | 256 KB | 48 MHz | 50 mA |
| Raspberry Pi 5 | 8 GB | SD | 2.4 GHz | 3 A |
Quantization Techniques
INT8 Quantization
import torch
import torch.quantization as quantization
def quantize_model(model: torch.nn.Module, target_path: str):
"""INT8 квантование модели для edge-деплоя"""
model.eval()
# Квантование весов
quantized = quantization.quantize_dynamic(
model,
{torch.nn.Linear},
dtype=torch.qint8
)
# Проверка точности
accuracy = evaluate(quantized, validation_data)
print(f"Accuracy: {accuracy:.2%}")
torch.jit.save(torch.jit.script(quantized), target_path)
PTQ vs QAT
PTQ (Post-Training Quantization):
┌──────────┐ ┌───────────┐ ┌──────────┐
│ FP32 │ → │ Quantize │ → │ INT8 │
│ Model │ │ (no train)│ │ Model │
└──────────┘ └───────────┘ └──────────┘
Accuracy: ~95%
QAT (Quantization-Aware Training):
┌──────────┐ ┌───────────┐ ┌──────────┐
│ FP32 │ → │ SimQ Train│ → │ INT8 │
│ Model │ │ (fake quant)│ │ Model │
└──────────┘ └───────────┘ └──────────┘
Accuracy: ~99%
Binary и Ternary Networks
class BinaryLinear(torch.nn.Module):
"""Бинарный линейный слой: веса ∈ {-1, 0, +1}"""
def forward(self, x):
weights = torch.sign(self.weight)
return torch.nn.functional.linear(x, weights)
class TernaryLinear(torch.nn.Module):
"""Тернарный линейный слой: веса ∈ {-1, 0, +1}"""
def __init__(self, in_features, out_features, sparsity=0.5):
super().__init__()
self.weight = torch.nn.Parameter(torch.randn(out_features, in_features))
self.sparsity = sparsity
def forward(self, x):
w = self.weight
# Threshold для sparsity
threshold = torch.quantile(w.abs(), self.sparsity)
w = torch.sign(w) * torch.relu(w.abs() - threshold)
return torch.nn.functional.linear(x, w)
Model Compression
Pruning
import torch.nn.utils.prune as prune
def prune_model(model: torch.nn.Module, amount: float = 0.5):
"""Структурный pruning модели"""
for name, module in model.named_modules():
if isinstance(module, torch.nn.Linear):
# Unstructured pruning
prune.l1_unstructured(module, name='weight', amount=amount)
return model
# Structured pruning (по каналам)
def structured_pruning(model, sparsity_per_layer: dict):
for name, module in model.named_modules():
if isinstance(module, torch.nn.Linear) and name in sparsity_per_layer:
prune.ln_unstructured(
module, name='weight',
n=2, # channel-wise
amount=sparsity_per_layer[name]
)
Knowledge Distillation for Edge
def distill_teacher_to_student(
teacher: torch.nn.Module,
student: torch.nn.Module,
train_loader,
temperature: float = 4.0,
epochs: int = 50
):
"""Дистилляция знаний: большая модель → маленькая"""
optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
criterion = torch.nn.KLDivLoss(reduction="batchmean")
for epoch in range(epochs):
for inputs, _ in train_loader:
# Teacher inference (no grad)
with torch.no_grad():
teacher_logits = teacher(inputs) / temperature
# Student inference
student_logits = student(inputs) / temperature
# Loss: KL divergence between distributions
loss = criterion(student_logits, torch.softmax(teacher_logits, dim=1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
Deployment Frameworks
TensorFlow Lite
import tensorflow as tf
# Convert PyTorch → ONNX → TFLite
def convert_to_tflite(model_path: str, output_path: str):
converter = tf.lite.TFLiteConverter.from_concrete_files(model_path)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# INT8 full quantization
def representative_dataset():
for _ in range(100):
yield [np.random.rand(1, 1, 224, 224).astype(np.float32)]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_model = converter.convert()
with open(output_path, 'wb') as f:
f.write(tflite_model)
MicroTVM
import tvm
from tvm import relay
def compile_for_microcontroller(model, target_device: str):
"""Компиляция модели для микроконтроллера через MicroTVM"""
target = tvm.target.Target(target_device)
mod, params = relay.frontend.from_pytorch(model, input_shape)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
# Deploy on microcontroller
from tvm.contrib import utils, download
temp = utils.tempdir()
lib_path = temp.relpath("model.so")
lib.export_library(lib_path)
return lib
ONNX Runtime Micro
import onnxruntime as ort
# Edge inference с ONNX Runtime
def create_edge_session(model_path: str):
providers = [
('TensorrtExecutionProvider', {'cudnn_conv_algo_search': 'HEURISTIC'}),
('CUDAExecutionProvider'),
'CPUExecutionProvider'
]
session = ort.InferenceSession(
model_path,
providers=providers,
sess_options=ort.SessionOptions()
)
return session
Optimisation Strategies
KV Cache Quantization
class QuantizedKVCache:
"""INT8 KV Cache для экономии памяти"""
def __init__(self, num_heads: int, head_dim: int, max_len: int):
self.num_heads = num_heads
self.head_dim = head_dim
self.max_len = max_len
self.cache_size = num_heads * head_dim * max_len
# INT8 storage (4x compression vs FP16)
self.k_cache = torch.int8.zeros(self.cache_size)
self.v_cache = torch.int8.zeros(self.cache_size)
self.k_scale = torch.float32(1.0)
self.v_scale = torch.float32(1.0)
def update(self, step: int, k: torch.Tensor, v: torch.Tensor):
k_quant = torch.quantize_per_tensor(k, self.k_scale, 0, torch.int8)
v_quant = torch.quantize_per_tensor(v, self.v_scale, 0, torch.int8)
self.k_cache[step] = k_quant.int_repr()
self.v_cache[step] = v_quant.int_repr()
def get(self) -> tuple:
k = torch.dequantize(self.k_cache)
v = torch.dequantize(self.v_cache)
return k, v