Getting an LSTM Sentiment Model Onto SiMa's Modalix MLA

A step-by-step walkthrough - and exact terminal commands - for taking a PyTorch LSTM from training all the way to real-time inference on-device.

The MLA on SiMa's Modalix chip is fast, but it only speaks a fairly narrow dialect of ONNX - and the ONNX LSTM operator isn't in that dialect. Getting a trained sentiment model running on-device means taking it apart: splitting the embedding lookup off onto the CPU, unrolling the LSTM recurrence into a plain directed graph of MatMul/Add/Sigmoid nodes, approximating Tanh with a piecewise-linear stand-in, and then quantizing and packaging the result for the accelerator.

Below is the full path from a PyTorch checkpoint to a running inference loop on the DevKit, with every command written out exactly as you'd type it into a terminal - no orchestration layer, no agent, just the commands themselves and which machine each one runs on.

The Six Phases

Two decomposition tricks make this possible: the embedding table stays on the host CPU (a Gather over a 20,000×200 table isn't MLA-friendly), and the LSTM itself gets unrolled across all 150 timesteps into a static graph small enough for the compiler to import.

Phase 1: Train the Model & Export to ONNX

The model is a 2-layer unidirectional LSTM (vocab 20,000, sequence length 150, embedding dim 200, hidden dim 128) trained on IMDB reviews. Install the training dependencies, run the training script, then split the exported ONNX graph so the embedding lookup stays separate from the LSTM/classifier.

# 1. Install training dependencies
pip install torch>=2.0 onnx>=1.14 onnxruntime>=1.16 numpy>=1.24 datasets>=2.14 scikit-learn>=1.3

# 2. Train the LSTM on IMDB reviews -> best_sentiment_model.pt
python3 sentiment_lstm.py

# 3. Split the exported ONNX into a host embedding model + device LSTM/classifier
python3 decompose_sentiment.py

Phase 2: Graph Surgery - Swap LSTM for Supported Ops

The MLA's supported-operator list has no LSTM and no Tanh. The fix: fit a piecewise-linear hinge-ReLU approximation for tanh, then unroll the LSTM recurrence over all 150 timesteps into a static graph built entirely from MatMul, Add, Sigmoid, Mul, Sub, and Relu.

# 1. Fit the tanh(x) ≈ c0 + c1*x + Σ ak·ReLU(x - bk) approximation
#    (least-squares fit over x in [-4, 4], 6 hinge points)
python3 fit_tanh_hinge.py
# -> tanh_hinge_coeffs.npz

# 2. Replace the LSTM nodes with the unrolled, hinge-tanh subgraph
python3 lstm_surgery.py
# -> sentiment_decomposed.onnx (~46,000 nodes - too large to compile as-is)

The naive unroll produces roughly 46,000-49,000 nodes for the stacked LSTM, which is well past what the SiMa Model Compiler's TVM importer can handle (a practical ceiling of ~5,000-10,000 nodes). The working path prunes this down to a compilable sentiment_clean.onnx at 3,905 nodes by removing unused weight initializers, swapping Squeeze/Unsqueeze for Reshape, and exporting at ONNX opset 14:

# Produce the pruned, compilable model
# (lstm_surgery.py adapted with --output to emit the cleaned graph directly)
python3 lstm_surgery.py --output sentiment_clean.onnx

Phase 3: Validate Against the Original Model

Before spending time quantizing and compiling, confirm the decomposed ONNX graph actually reproduces the original PyTorch model's outputs.

python3 validate_sentiment.py

Phase 4: Quantize & Compile for the Modalix MLA

This is the one step that runs inside SiMa's SDK container (image sih2024/sima-neat-quantize), where the Model Compiler lives. Enter the container, activate the compiler, and confirm the versions before running the quantization script.

# Enter the SDK container
docker exec -it sima-neat-quantize bash

# Activate the Model Compiler
source /sdk-ext/model-compiler/bin/activate

# Sanity-check the SDK version
python3 -c "from afe.apis.release_v1 import get_model_sdk_version; print(get_model_sdk_version())"
# Expected: afe 2.1.0, jax 0.4.30, numpy 1.26.4

# Run quantization + compilation
# (expects sentiment_clean.onnx and sentiment_host_embedding.onnx in /workspace/graphsurgery_lstm/)
python3 /workspace/graphsurgery_lstm/quantize_clean.py

# Leave the container once it finishes
exit

The script quantizes to int8 (per-tensor asymmetric activations, per-channel symmetric weights, MSE calibration) and compiles, writing everything to build_quantized_clean_v2/. The two files that matter for deployment are sentiment_clean_mpk.tar.gz (the deployable package) and sentiment_clean_int8.sima (the raw quantized model inside it).

Phase 5: MPK Packaging (usually automatic)

The MPK is normally produced automatically as part of Phase 4's compile step. You only need the commands below if you're repackaging an existing .sima model manually, or want to double-check what's inside the archive before shipping it to the DevKit.

# Only needed if sima-cli isn't already available in the container
pip install sima-cli

# Manually repackage a .sima model into an MPK, if needed
sima-cli mpk create \
    --model sentiment_clean_int8.sima \
    --name sentiment_clean \
    --output sentiment_clean_mpk.tar.gz

# Verify the archive contents
tar -tzf sentiment_clean_mpk.tar.gz
# Expected: manifest.json, share/sentiment_clean_stage1_mla.elf, share/pipeline_config.json

Phase 6: Transfer to the DevKit & Set Up the Runtime

Three files need to land on the DevKit: the MPK package, the host embedding ONNX model, and the vocabulary. Copy them over with scp, then SSH in to set up the Python environment.

# Copy the model, embedding, and vocab over to the DevKit
# (enter the password "edgeai" when prompted for each file)
scp sentiment_clean_mpk.tar.gz sima@192.168.0.28:~/sentiment_model/
scp sentiment_host_embedding.onnx sima@192.168.0.28:~/sentiment_model/
scp word2idx.json sima@192.168.0.28:~/sentiment_model/

# SSH into the DevKit (password: edgeai)
ssh sima@192.168.0.28

# One-time environment setup
python3 -m venv ~/pyneat
source ~/pyneat/bin/activate
pip install pyneat onnxruntime numpy

# Verify the versions match what the pipeline expects
python3 -c "import pyneat; print(pyneat.__version__)"          # expected: 0.2.2
python3 -c "import onnxruntime; print(onnxruntime.__version__)" # expected: 1.26.0

Phase 7: Run Inference on the DevKit

With the environment set up, run the inference script directly. Text goes in on the CPU (tokenize → embed via onnxruntime), the embedding tensor crosses over to the MLA for the LSTM + classifier, and a single logit comes back.

# Run the end-to-end test script
python3 ~/sentiment_model/test_real_text.py

Optional: Improve Quantization With Real Calibration Data

If the deployed int8 model outputs a near-constant logit for every input, the random calibration data from Phase 4 is the likely cause. Generating calibration samples from real IMDB embeddings instead usually fixes it - real train-time activations produce meaningfully better scale and zero-point values than uniform noise.

# Generate calibration embeddings from real IMDB training data
# (run your calibration-generation script - it loads sentiment_host_embedding.onnx,
#  encodes 256 real IMDB reviews, and saves the resulting embeddings)
python3 generate_calibration_data.py
# -> calib_embeddings.npy

# Re-run quantization, now pointed at the real calibration data
# (update quantize_clean.py first so it loads calib_embeddings.npy instead of random noise)
python3 quantize_clean.py

Troubleshooting Common Failure Modes

Command Quick Reference

All seven phases, one command block per machine, in the order you'd actually run them.

# ── HOST ──────────────────────────────────────────────
pip install torch>=2.0 onnx>=1.14 onnxruntime>=1.16 numpy>=1.24 datasets>=2.14 scikit-learn>=1.3
python3 sentiment_lstm.py
python3 decompose_sentiment.py
python3 fit_tanh_hinge.py
python3 lstm_surgery.py --output sentiment_clean.onnx
python3 validate_sentiment.py
python3 build_vocab.py

# ── HOST, ENTER CONTAINER ────────────────────────────
docker exec -it sima-neat-quantize bash

# ── INSIDE THE CONTAINER ─────────────────────────────
source /sdk-ext/model-compiler/bin/activate
python3 /workspace/graphsurgery_lstm/quantize_clean.py
exit

# ── HOST -> DEVKIT ───────────────────────────────────
scp sentiment_clean_mpk.tar.gz sima@192.168.0.28:~/sentiment_model/
scp sentiment_host_embedding.onnx sima@192.168.0.28:~/sentiment_model/
scp word2idx.json sima@192.168.0.28:~/sentiment_model/
ssh sima@192.168.0.28

# ── ON THE DEVKIT ────────────────────────────────────
source ~/pyneat/bin/activate
python3 ~/sentiment_model/test_real_text.py