silvic-neuron

silvic-neuron trains and serves LightGBM-compatible models on AWS Neuron accelerators. It runs on Linux with Trainium, Trainium2 or Inferentia2, using NKI kernels through the Neuron Runtime. Serving uses float32 arithmetic, and training and serving require explicit device selection.

Setup

Execution requires the Neuron driver and runtime on the accelerator host and a separate Linux x86-64 Python 3.11 environment for the compiler. The Rust crate builds without the Neuron SDK or PyTorch. Create the compiler environment and set the runtime paths before starting the process:

python3.11 -m venv .venv-neuron
.venv-neuron/bin/pip install --extra-index-url https://pip.repos.neuron.amazonaws.com \
    'neuronx-cc==2.21.18209.0+043b1bf7' 'numpy==1.26.4' 'scipy==1.12.0'
export SILVIC_NEURON_PYTHON="$PWD/.venv-neuron/bin/python"
export SILVIC_NEURON_LIBRARY=/opt/aws/neuron/lib/libnrt.so
export NEURON_LOGICAL_NC_CONFIG=1

Trainium2 requires logical-core configuration 1 to match the single-core compiled kernels. The compiler adapter requires the pinned neuronx-cc version. silvic caches compiled kernels by source, shapes, constants, target and compiler version under $XDG_CACHE_HOME/silvic/neuron or ~/.cache/silvic/neuron. SILVIC_NEURON_CACHE overrides the directory.

Training

import silvic

booster = silvic.train(
    {"objective": "lambdarank", "device_type": "neuron"},
    silvic.Dataset(X, label=y, group=query_sizes),
    num_boost_round=200,
)

During training, binned rows, partitions and scores stay on the device, where NKI kernels build histograms, partition rows, reduce leaf gradients and update scores. The host computes objectives, quantizes gradients, selects bags and searches for splits. Bounded integer tiles preserve exact quantized histograms, and signed radix-2¹⁶ limbs preserve fixed-point leaf sums during device reductions. Float histograms are approximate.

The gpu_device_id parameter selects a logical NeuronCore within NEURON_RT_VISIBLE_CORES. The auto selector considers CUDA, Metal and CPU, while gpu considers CUDA and Metal.

For distributed training, silvic.ray.train reserves one neuron_cores resource per worker by default when device_type="neuron". Each backend uses one logical core, with Ray setting core visibility and workers reducing histograms over the TCP ring. Combining Neuron selection with use_gpu=True or GPU resource reservations raises a ValueError.

Serving

import silvic.serve

predictor = silvic.serve.Predictor(booster, device="neuron", precision="float32")
scores = predictor.predict(X_test)

The predictor retains the forest and reusable buffers on the device, and it adds float32 leaf values in tree order. It supports numerical and categorical splits, missing values and tree prefixes. Float64 inputs are converted to float32 before prediction; thresholds, leaves, accumulation and output also use float32, so predictions are approximate.

silvic.serve.Predictor requires explicit precision="float32" and num_threads=1 when device="neuron". An explicit CPU kernel selection raises a ValueError. CPU and CUDA serving default to float64.

Rust API

NeuronBackend::new(ordinal)

Implements silvic_core::Backend. None selects logical core 0. Oversized datasets or models fail explicitly at the signed 32-bit device-address limit.

NeuronPredictor::new(&model, ordinal)

predict(&rows, raw_score, num_trees) returns float32 scores; predict_raw(&rows, num_trees, &mut out) reuses a caller’s float32 output buffer. Both accept row- or column-major float32 or float64 matrices.

is_available()

Probes visible NeuronCores without initializing the runtime or reserving workers’ cores. Training and serving return runtime, device and compiler errors.