needle-onnx-lfm

Needle — ONNX export (finetuned, LFM dataset)

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

Файлы

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.

Архитектура