RockMan256/needle-onnx-lfm

Needle — ONNX export (finetuned, LFM dataset)

Вывод файнтюна + конвертации Flax→PyTorch→ONNX.

Файлы

Файл	Описание	Примерный размер
encoder.onnx	Encoder. Input: input_ids(B,T), Output: encoder_out(B,T,512)	~55 MB
decoder_step.onnx	Decoder step (1 токен). KV cache in/out.	~85 MB
needle_torch.pt	PyTorch state_dict (веса)	~100 MB
needle_torch.config.json	Конфиг модели	tiny
needle.model	SentencePiece BPE (vocab=8192)	125 KB
tokenizer-specials.json	ID специальных токенов	tiny

TensorRT на Jetson

TensorRT engine архитектурно-зависим — собранный на T4/L4 не запустится на Jetson Orin.
Конвертируй ONNX в engine непосредственно на Jetson:

# Скачать ONNX
huggingface-cli download RockMan256/needle-onnx-lfm encoder.onnx --local-dir .
huggingface-cli download RockMan256/needle-onnx-lfm decoder_step.onnx --local-dir .

# TensorRT 10.3 на Jetson (НЕ использовать --fp16!)
TRTEXEC=/usr/src/tensorrt/bin/trtexec

# Encoder (dynamic seq 1..4096)
$TRTEXEC --onnx=encoder.onnx --saveEngine=encoder.engine   --minShapes=input_ids:1x1   --optShapes=input_ids:1x512   --maxShapes=input_ids:1x4096

# Decoder step (enc_seq 1..4096, past_seq 0..256)
$TRTEXEC --onnx=decoder_step.onnx --saveEngine=decoder_step.engine   --minShapes='decoder_input_ids:1x1,encoder_out:1x1x512,past_self_kv:8x2x1x4x0x64'   --optShapes='decoder_input_ids:1x1,encoder_out:1x256x512,past_self_kv:8x2x1x4x128x64'   --maxShapes='decoder_input_ids:1x1,encoder_out:1x4096x512,past_self_kv:8x2x1x4x256x64'

Важно:

• --fp16 ломает bfloat16 → нули на выходе. Только float32.
• Размеры past_self_kv считаются по формуле: (num_decoder_layers, 2, batch, num_kv_heads, past_seq, head_dim)
  Для 26M: layers=8, batch=1, kv_heads=4, head_dim=64

API сервер (Jetson)

После сборки engines запусти сервер:

python3 api_server.py 39971
# или через PM2:
pm2 start api_server.py --name needle-api --interpreter python3 -- 39971

Сервер принимает POST /v1/generate с prompt, tools (плоский JSON), max_new_tokens, temperature=0.0.

Архитектура

Параметр	Значение
Параметры	26.2M
vocab_size	8192
d_model	512
num_heads / kv_heads	8 / 4 (GQA)
encoder / decoder layers	12 / 8
head_dim	64
activation	SwiGLU
norm	ZCRMSNorm
FFN	Нет
dtype	bfloat16 (weights), float32 (inference)