tzchen07/rai-nemotron3-nano-dpo-w2

dpo_w2 — RAI content-safety classifier (DPO precision-tuned)

What this is: a full, standalone model — complete merged BF16 weights, a drop-in replacement for v3d_tjpi3o6. It is NOT a LoRA adapter. LoRA was only the training mechanism for the DPO stage; the adapter was merged into the base before release, so this repo ships complete model-*.safetensors and needs no base model and no adapter to run.

Lineage: Nemotron-3-Nano-4B → v3d_tjpi3o6 (full-parameter SFT) → DPO (LoRA r32, merged). RECOMMENDED shipping arm — best v14 ≤ 0.25% / recall balance.

Training method — artifact vs. method (read this to avoid the "is it LoRA?" confusion)

The two stages used different training methods, but both ship as full merged models:

stage	how it was trained	what the repo contains
SFT → v3d_tjpi3o6	full-parameter SFT (every weight updated)	full weights (separate repo)
DPO → dpo_w2 (this repo)	LoRA — r=32, α=64, target_modules=all-linear; DPO β=0.05, lr cosine, 1 epoch, on top of tjpi3o6	full merged weights — the LoRA was folded in with merge_and_unload()

• Reference model (free): DPO needs a frozen reference policy. Because the DPO update was a LoRA, the reference was obtained by simply disabling the adapter (PEFT adapter-off) on the same weights — no second model copy in memory.
• Merged for release: after training, merge_and_unload() → save_pretrained(). Therefore there is no adapter_config.json, no adapter_model.safetensors, and no base-model dependency — the repo root is the complete, deployable model (same file layout as tjpi3o6).
• Why LoRA for DPO but full-FT for SFT? N3 is a hybrid Mamba/SSM model. The big capability injection (SFT) needs full-parameter training because LoRA can't move the SSM state modules — that's why tjpi3o6 is full-SFT. The DPO stage here is a narrow calibration (pushing down confident benign false-blocks) that lives in the linear/projection layers where LoRA suffices, and LoRA is what enables the free adapter-disable reference. DPO's loss also has a non-saturating gradient on high-confidence tokens, which is exactly what moved the confident false-blocks that full-SFT (100+ arms) walled at v14 ≈ 0.36%.

How the DPO data was created

The preference set is failure-mode-driven — built from tjpi3o6's own mistakes, entirely from the training side (zero eval-set leakage):

1. Mine the errors. Run tjpi3o6 over the training and development data and collect its confident benign false-blocks — benign requests it wrongly flagged. Dominant clusters: jailbreak / own-agent-control (§07), IP / own-content transforms (§08), synthetic/dummy PII (§09), and reading-existing-feedback (§13). (Own-agent JB alone drove the plurality of false-blocks.). Evaluation data is untouched and remains blind.
2. fix-false-block pairs (the majority). For each wrongly-blocked benign request: prompt = the request rendered in the deployment prompt (with the <think></think> reasoning scaffold — deploy render, not the SFT render); chosen = the correct Verdict: none (allow) with a v3_1d-policy-cited reason; rejected = the block verdict tjpi3o6 actually produced. → teaches "prefer allow over block" precisely on the misfire patterns.
3. recall-preserve (contrastive) pairs. For genuine violations: chosen = block, rejected = allow. → stops DPO from over-correcting into under-blocking; protects recall.
4. The mix ratio is the knob. wave-2 ≈ 800 pairs, roughly 3 : 1 fix-false-block : recall-preserve. The recall-preserve share is the v14 ↔ recall dial: 3:1 → v14 0.23 / recall 0.759; 2:1 → v14 0.27 / higher recall. β was ~invariant (0.03 / 0.05 / 0.07 land on the same point) → the operating point is data-determined, not β-determined; denser data, not more epochs (which over-fit and drop recall), is the lever.

Two waves were run: wave-1 (450 pairs, 2:1) and wave-2 (800 denser pairs, 3:1, JB-heavy + feedback-contrastive). dpo_w2 = wave-2. A later 36-arm carve-out data wave did not beat it, so wave-2 is the practical N3 DPO frontier. (The wave-2 preference JSONL, run log, recipe, and raw eval inferences were bundled in an earlier revision of this repo and were removed in a size cleanup; they can be re-attached on request.)

Metrics (v3_1d policy; gate = v14<=0.35 . hello>=0.47 . v200-P>=0.95 . rvR>=0.80 . rvP>=0.90) — 5/5

v14 Block%	hello R	v200-P	v200-R	rvR	rvP
0.23%	0.759	0.999	0.889	0.862	0.931

Audited hello_v2_full (396 ex): P 0.831 / R 0.771.

vs the tjpi3o6 base (v14 0.36 / rvP 0.914 / v200-P 0.990): v14 down, precision up (rvP, v200-P), recall held — DPO reduced benign false-blocks, which lifts every precision metric at once.

Contents (this repo = 13 files, the complete merged model)

• merged final weights — model-00001-of-00002.safetensors, model-00002-of-00002.safetensors, model.safetensors.index.json (~7.95 GB, 4B BF16).
• tokenizer — tokenizer.json, tokenizer_config.json, special_tokens_map.json, chat_template.jinja.
• config — config.json, generation_config.json.
• trust-remote-code modules — modeling_nemotron_h.py, configuration_nemotron_h.py.

There are no checkpoint-* dirs, no optimizer.pt/scheduler.pt/trainer_state.json, and no separate LoRA adapter — the DPO trainer ran with save_strategy="no" and saved only the merged weights. The repo root is the complete deployable model.

Serve

vllm serve <repo> --trust-remote-code --mamba_ssm_cache_dtype float32

---

🚀 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-dpo-w2 \
  --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-dpo-w2", 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-dpo-w2"
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>.