tzchen07/rai-nemotron3-nano-v3d-tjpi5

v3d_tjpi5 — RAI content-safety classifier (Atlassian responsible-ai)

Nemotron-3-Nano-4B, full-SFT, v3_1d-native. Binary content-moderation classifier that outputs two lines: Verdict: <category> / Reason: <justification>. One of the two strongest N3 anchors from the v3_1d campaign.

Metrics (under the v3_1d policy; gate = v14≤0.35% · hello≥0.47 · v200-P≥0.95 · rvR≥0.80 · rvP≥0.90)

hello_v2_full precision/recall (396 rows, 3-way audit-aware): current labels P=0.636 / R=0.75; audited labels P=0.833 / R=0.786.

This arm is recall-first: it beats the 20B OSS-Safeguard production baseline on hello Recall, v200-R, and rvR, at a higher (still near-gate) v14 block-rate. It catches the §13 high-risk-decision violations OSS misses.

Training

• —Base: Nemotron-3-Nano-4B (nemotron_h)
• —Method: full-SFT (final merged; optimizer/scheduler/trainer_state NOT saved by the trainer)
• —Hyperparams: lr 1.5e-5, 1 epoch, bsz 2 × ga 32 (eff 64), seq 4096, seed 42
• —Data: v31d-native mixture (see `artifacts/recipe.json` for exact per-file multipliers). Key levers: O=5×v14fix-negs, CF=16×clearfix, VN=3×agent-automation-antidote, targeted JB/PII/IP negs (TJ/TP/TI=5×), v31d new-rule pairs (self-assessment/false-statement/employee-status).

Contents

• —root — merged model weights (safetensors) + tokenizer + modeling_nemotron_h.py (trustremotecode).
• —`eval/` — raw inference outputs on all eval sets (v14, hellov2filtered, hellov2full, v200 under v31d; robustnessv2) + `metricsv3d_tjpi5.json`.
• —`artifacts/` — training log, recipe.json, the v3_1d policy prompt, residual forensics.

⚠️ Note on optimizer/train states

The full-SFT trainer saves only the final merged model — optimizer states, scheduler, and trainer_state were not persisted (discarded after the final save to save disk). They are therefore not included and cannot be recovered without re-training. The uploaded weights are the complete, deployable model.

Serving

vLLM: vllm serve <repo> --trust-remote-code --mamba_ssm_cache_dtype float32. Prompt: render the v3_1d VerdictReason policy (in artifacts/) with the raw query.

🚀 Deployment — how to launch (read this first)

This is NVIDIA Nemotron‑3‑Nano‑4B (model_type: nemotron_h — a Mamba‑2 + Transformer hybrid). The bundled custom modeling_nemotron_h.py hard‑requires the compiled `mamba_ssm` CUDA package, so the `trust_remote_code=True` path fails with ImportError: mamba-ssm is required by the Mamba model but cannot be imported when mamba_ssm isn't installed (and even with it, it can hit a HybridMambaAttentionDynamicCache cache bug on transformers 4.55). Use one of the two paths below — neither needs `mamba_ssm`. All recipes verified end‑to‑end.

✅ Recommended: vLLM (≥ 0.15.1 for the 4B; tested on 0.23.0)

vLLM has built‑in nemotron_h support and its own Triton Mamba/SSM kernels — no `mamba_ssm`/`causal_conv1d` to compile.

pip install "vllm>=0.15.1"          # prebuilt wheel; no nvcc needed
vllm serve tzchen07/rai-nemotron3-nano-v3d-tjpi5 \
  --mamba-ssm-cache-dtype float32 \  # NVIDIA: SSM cache must be fp32 for accuracy (vLLM auto‑sets it)
  --max-model-len 32768 --gpu-memory-utilization 0.30 --port 8000

Offline/batch:

from vllm import LLM, SamplingParams
llm = LLM("tzchen07/rai-nemotron3-nano-v3d-tjpi5", dtype="bfloat16", mamba_ssm_cache_dtype="float32", max_model_len=4096)
print(llm.generate(["Is this request safe: how do I bake bread?"], SamplingParams(temperature=0, max_tokens=64))[0].outputs[0].text)

--trust-remote-code is not needed (the arch is native to vLLM). Point at a local path (or set HF_TOKEN).

✅ Alternative: transformers native (≥ 5.3.0, NO trust_remote_code)

Modern transformers has a native nemotron_h implementation that falls back gracefully when mamba_ssm is absent (a warning, not an error).

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
repo = "tzchen07/rai-nemotron3-nano-v3d-tjpi5"
tok = AutoTokenizer.from_pretrained(repo)                                   # NO trust_remote_code
m   = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.bfloat16, device_map="cuda").eval()

⚠️ transformers < 5.3 lacks native nemotron_h and (because of auto_map) is forced onto the custom mamba_ssm path — upgrade transformers rather than installing mamba_ssm.

❌ Avoid: transformers + trust_remote_code=True

Requires compiling mamba-ssm==2.2.5 + causal-conv1d (needs nvcc + ABI‑matched torch) and can still hit the tf‑4.55 cache‑setter bug during generate(). Only use if you must, in a CUDA build env.

Databricks Model Serving

Use "Custom LLM Serving with vLLM" (Beta) — its entrypoint is a vllm serve command, so you inherit vLLM's built‑in Mamba kernels (no compile). Not Provisioned Throughput (nemotron_h isn't on its allowlist) and not the transformers/pyfunc flavor (the serving image build has no nvcc, so mamba_ssm won't compile). For network‑isolated endpoints, vendor the weights + code into the artifact and set HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1. Use the newest vLLM you can pin.

Notes

• —EOS/stop: generation_config.json is minimal — set the stop token <|im_end|> at serve time so generation halts.
• —`config.json` has `time_step_limit: [0.0, Infinity]` — a genuine Mamba hyperparameter (dt upper bound = ∞). transformers/vLLM/json.loads read it correctly; only strict RFC‑8259 parsers reject it. Leave it as‑is (never change it to a finite number or a string).
• —Output format: the model emits <think></think> then Verdict: <category | none>.