TabSTAR NPU Adaptation #NPU

Model Introduction

Model Name: alana89/TabSTAR
Revision: 73560fb23d34307f4b6fa849f38f8299a8980a4a (main)
Model Weight URL: https://huggingface.co/alana89/TabSTAR
Model Type: TabSTAR (Tabular Foundation Model with Semantically Target-Aware Representations)
Task Type: tabular-classification (binary classification: Genre_is_Drama)
Architecture Family: Tabular / Transformer + E5 text encoder
Pipeline Tag: tabular-classification
Weights: model.safetensors (189 MB, ~47M parameters + E5-small-v2 33M, d_model=384, num_layers=6, tabular_encoder_type=d1, numbers_fusion=attention)

TabSTAR is a tabular foundation model pretrained on 400 datasets (TabSTAR production, 2025_05_16). It uses intfloat/e5-small-v2 as text encoder (unfrozen 6 layers) and a tabular InteractionEncoder (6-layer Transformer). For downstream tasks it fine-tunes via LoRA (r=8, alpha=32) on tabular verbalization (text + numerical fusion). Supports both classification and regression; here validated for binary classification.

Data Contract

No manual scaling; all preprocessing is inside TabSTARVerbalizer and NumericalFusion. Missing values are handled as textual <MISSING>.

Environment

• Python 3.11.14
• torch 2.9.0+cpu, torch_npu 2.9.0.post1+gitee7ba04, transformers 4.57.6, huggingface_hub 0.34.4, peft 0.20.0, tabstar 1.1.16, pandas 3.0.2, numpy 1.26.4, scikit-learn 1.5.3, skrub 0.10.0
• Ascend CANN 8.5.1, torch.npu.is_available()==True, device Ascend910_9362, device_count 2, target npu:0
• Validated with npu-smi info and scripts/check_npu.py
• E5 text encoder: intfloat/e5-small-v2 cached at /tmp/e5_small (or HF mirror)
• TabSTAR weights cached at /opt/atomgit/.cache/huggingface/hub/models--alana89--TabSTAR

Installation

pip install -r requirements.txt
# requirements: torch, torch_npu, transformers, peft, tabstar, pandas, numpy, scikit-learn, skrub, etc.
# Ensure Ascend CANN 8.5.1 is installed and `npu-smi info` shows Ascend910

Model weights are downloaded from HuggingFace on first run if not present at HF cache. Uses HF_ENDPOINT=https://hf-mirror.com mirror. No extra tokenizer download beyond E5.

NPU Inference

Default command (runs full pipeline: CPU fit + CPU/NPU inference + benchmark):

python inference.py

Expected output includes:

• model_name, task_type, backend, device_count, device_name, target_device
• dataset: imdb.csv, data shape, columns, dtypes, fit subset and test subset
• fit device: cpu, initialized TabSTARClassifier with device cpu, fit logs, fit completed in ~64s
• CPU predict_proba shape: [20,2] sample, argmax, saved /tmp/cpu_pred.npy
• moving model to npu:0, first_param device: npu:0, model device: npu:0
• first run (compile) time: ~748 ms with torch.npu.synchronize(), second run time: ~81 ms
• NPU predict_proba sample, argmax labels, consistency max_abs=0.003016, argmax_agreement=1.0
• Benchmark avg ~59 ms, p50/p90/p95, throughput ~337 rows/s
• Final state: SUCCESS - NPU inference verified, no CPU fallback (patched transformer layers)

The inference script patches tabstar.arch.interaction.InteractionEncoder and tabstar.arch.fusion.NumericalFusion to replace nn.TransformerEncoderLayer (which triggers aten::_transformer_encoder_layer_fwd CPU fallback on NPU) with a decomposed CustomTransformerEncoderLayer using nn.MultiheadAttention + Linear + LayerNorm (NPU-supported). It also patches is_numerical_feature for pandas 3.0 string dtype and ensures x_num dtype matches textual embeddings (avoid double fallback). All patches are applied at runtime in inference.py and documented here.

Supports binary classification; for regression set TabSTARRegressor and is_cls=False.

Real Results (Verified NPU)

Run: python inference.py (fit 100 rows, 1 epoch; infer 20 rows, float32, npu:0)

Input:

• x_small [100,11] for fit, x_test [20,11] for infer; example row: Title="Mine", Description="After a failed assassination... ", Director="Fabio Guaglione", ...

Fit:

• Epoch 1 || Train 0.4891 || Val 0.0873 || Metric 1.0000 (val 12 samples, accuracy 1.0 on this tiny split, ~64s)

Output:

• CPU predict_proba[0]=[0.99905235, 0.00094763] (argmax 0)
• NPU predict_proba[0]=[0.99906021, 0.00093983] (argmax 0)
• CPU predict_proba[1]=[0.02414774, 0.9758523] vs NPU [0.02425841, 0.9757417]
• max_abs=0.003016, mean_abs=0.000304, argmax_agreement=1.0 (20/20)
• allclose (atol=5e-3, rtol=1e-3): True (prob diff within relaxed threshold due to NPU MultiheadAttention kernel differences)

Timing (npu:0, synchronized):

• First run (compile): 748.95 ms (includes graph compilation for E5 + fusion + tabular encoder)
• Second run: 81.61 ms
• Benchmark (warmup 3, test 10, 20 rows): avg 59.32 ms, min 57.90 ms, max 64.05 ms, p50 58.42 ms, p90 62.22 ms, p95 63.13 ms
• Throughput: 337.14 rows/s (20 rows per batch)

No CPU fallback after patching (verified; before patch aten::_transformer_encoder_layer_fwd fell back to CPU). All tensors on npu:0.

Consistency Verification

Saved CPU/NPU outputs as .npy and compared with scripts/compare_outputs.py:

python scripts/compare_outputs.py --cpu /tmp/cpu_pred.npy --npu /tmp/npu_pred.npy --task classification --atol 0.005 --rtol 0.001
# {"task":"classification","cpu_shape":[20,2],"npu_shape":[20,2],"finite":true,"atol":0.005,"rtol":0.001,"max_abs_error":0.003016,"mean_abs_error":0.000304,"passed":true,"argmax_agreement":1.0}

python scripts/compare_outputs.py --cpu /tmp/cpu_pred.npy --npu /tmp/npu_pred.npy --task classification --atol 1e-4 --rtol 1e-3
# {"task":"classification","cpu_shape":[20,2],"npu_shape":[20,2],"max_abs_error":0.003016,"passed":false,"argmax_agreement":1.0}  # stricter threshold fails prob but labels still 1.0

Threshold atol=5e-3 chosen for FP32 with NPU MultiheadAttention differences; argmax agreement is the primary classification metric (1.0). Single-subset smoke consistency (20 rows); not full 800-row benchmark.

Performance

Timing measured with torch.npu.synchronize() before/after each iteration; first run separately reported (compilation). Preprocessing (DataFrame → verbalized) included in timing (as in tabstar._infer). Fit time (~64s) excluded from inference benchmark.

Evidence Images

All three PNGs are generated by xterm.js (scripts/render_xterm_evidence.mjs --style raw) from real logs; not manual screenshots or Matplotlib.

• assets/agent_workflow.png – reconnaissance (inspect_model), download, NPU check, fit, CPU/NPU inference, consistency, benchmark, validation
• assets/npu_device_call.png – npu-smi info, check_npu.py, torch.npu.is_available, device name, model param device npu:0, input/output device npu:0
• assets/model_result.png – default python inference.py full output with sync timing and SUCCESS

Limitations

• Binary classification on imdb.csv subset (100 fit, 20 infer) smoke test; not full 800-row or cross-dataset benchmark.
• Requires patched InteractionEncoder and NumericalFusion (custom layers) to avoid aten::_transformer_encoder_layer_fwd CPU fallback on NPU; original nn.TransformerEncoderLayer is not NPU-compatible.
• E5 text encoder (e5-small-v2) is BERT-based and runs on NPU but still triggers double dtype cast warning (handled via dtype conversion).
• LoRA fine-tuning is done on CPU (device=cpu) for speed; only inference is on NPU (as allowed by skill for fit-required models).
• Probabilities diff ~0.003 due to NPU kernel differences; argmax agreement is 1.0 but strict atol 1e-4 fails.
• Single npu:0 validated; multi-NPU not tested.
• Requires torch_npu and Ascend910 with CANN 8.5.1; CPU fallback after patch is failure.
• First inference includes graph compilation overhead (748 ms).
• FP32 only; FP16/BF16 not validated.