I tried improving RAG accuracy of Nemotron 9B-v2 with NTA FAQ × RAFT

Introduction

Hello, I'm Morishige from Classmethod's Manufacturing Business Technology Department.

In a previous article, when I tried RAG (Retrieval-Augmented Generation) with Nemotron 9B-v2-Japanese, I ran a test where I passed a fictional internal policy document as a "reference document" and had the model answer questions. Even the base model picked up context to a reasonable degree, but there were instances where it returned hallucinations (plausible-sounding falsehoods) when the reference document didn't contain the answer.

So, how much would accuracy improve if we trained on actual domain data?

To investigate that question, I tried fine-tuning specialized for RAG using the National Tax Agency's FAQ dataset. It's also tax filing season, so I thought tax FAQs would be easy to visualize both as a verification exercise and as a practical application. The method is RAFT (Retrieval Augmented Fine Tuning), and the training data consists of 1,000 items.

With the growing trend toward sovereign AI, there's an increasing need to make models domain-specific without sending internal data to external parties. This time, by utilizing NGC containers, I was able to complete everything from training to inference evaluation entirely on a DGX Spark. No cloud GPU was required, and the only cost was $2 for CoT generation. I hope this serves as a reference for those who own a DGX Spark or want to improve RAG accuracy with their own domain data.

What is RAFT

RAFT (Retrieval Augmented Fine Tuning) is a fine-tuning method that improves the "answer generation" part of a RAG pipeline (arXiv:2403.10131).

In standard RAG, documents retrieved by the search engine are passed directly to the LLM to generate answers. Whether the LLM can accurately locate the correct answer within the documents depends on the model's general-purpose reading comprehension. RAFT directly addresses this by training the model on the skill of "citing the correct answer from relevant documents."

The approach to creating training data is as follows:

1. Prepare a correct document (oracle) for the question, along with distractor documents that are related but don't contain the answer
2. Include the oracle in 80% of samples, and exclude it from 20% of samples
3. For samples with an oracle, attach a CoT (Chain-of-Thought) style answer that reasons after citing the correct passage
4. For samples without an oracle, attach a response of "cannot answer"

The model simultaneously learns both the task of "finding the correct answer within documents" and the task of "being honest when there is no correct answer." By becoming accustomed to situations where distractors are mixed in, the goal is to enable robust answers even when search precision in the actual RAG pipeline is imperfect.

Creating a RAFT Dataset from National Tax Agency FAQs

Overview of JaGovFaqs-22k

For training data, I used JaGovFaqs-22k (CC BY 4.0). This is a dataset of approximately 22,800 FAQs collected from various Japanese government ministries, with question-and-answer pairs.

https://huggingface.co/datasets/matsuxr/JaGovFaqs-22k

Exploring the data, most answer texts fall within the range of 50–2,000 characters, which is just the right size for RAG reference documents. The copyright field contains ministry names, so this can be used to narrow down by domain.

Sampling

There's no need to use all 22,800 items. The RAFT paper confirms effectiveness at a scale of a few thousand items, and since there are also LoRA coverage constraints described later, I narrowed it down to 1,000 items.

This time, I concentrated on the National Tax Agency's FAQs. Rather than spreading thinly across 41 ministries, by concentrating 1,000 items on the single domain of taxation, the RAFT distractors (dummy documents that don't contain the answer) naturally become hard negatives that are "on the same tax topic but with different answers." After filtering, there are approximately 2,476 National Tax Agency FAQs, providing ample margin for sampling 1,000 items.

For testing, 200 separate items were reserved from the National Tax Agency, and an additional 100 items from other ministry FAQs were set aside as test-unseen. This was to see how a tax-specialized model would behave in other administrative domains.

Converting to RAFT Format

I wrote a script to convert the sampled FAQs into RAFT format. The processing flow is simple.

1. For each question, select 4 other FAQ answers from the same National Tax Agency as distractors (documents from the same tax domain are topically similar, making them harder negative samples)
2. Mix in the correct FAQ answer (oracle) for 80% of samples, and use only distractors for 20%
3. For samples with an oracle, have the Claude API (Haiku) generate CoT answers with citations
4. Mix in an additional 100 "cannot answer" samples, creating a total of 1,100 training data items

The cost of generating CoT answers with Claude Haiku was approximately $2. After generation, 50 items were randomly extracted for manual checking, confirming citation accuracy and format consistency. Local models weren't stable enough in accuracy, so cloud models were used only for creating training data. (It's possible this could have been resolved with prompt engineering.)

The final output is in JSONL format for SFT (Supervised Fine-Tuning), with a system / user / assistant messages structure.

{
  "messages": [
    {
      "role": "system",
      "content": "You are an AI assistant that answers questions about taxation. Please answer accurately based on the provided reference documents...."
    },
    {
      "role": "user",
      "content": "Please answer the question based on the following reference documents.\n\nReference Document 1:\n...\n\nReference Document 2:\n...\n\nQuestion: What is the scope of medical expenses eligible for the medical expense deduction?"
    },
    {
      "role": "assistant",
      "content": "Reference Document 3 states that 'Medical expenses eligible for the medical expense deduction are...' Based on this provision, ...therefore, the eligible scope is..."
    }
  ]
}

Constraints of Mamba2 + LoRA

Only 53% of All Parameters Can Be Trained

Nemotron 9B-v2-Japanese adopts a hybrid architecture of Mamba-2 and Transformer. Of the 56 layers, 27 are Mamba-2 (SSM), 4 are Self-Attention, and 25 are FFN (MLP).

When attempting to train with LoRA, a problem arises here. The in_proj and out_proj of Mamba-2 layers are not compatible with LoRA in the PEFT library (PEFT Issue #2274). The reason is that Mamba-2 is designed so that its proprietary CUDA kernels directly reference weight tensors, bypassing the forward hooks that LoRA inserts.

As a workaround, LoRA was applied only to the Attention layers and FFN layers. The parameter coverage rate is 53%.

Whether RAFT would be effective with nearly half of the parameters frozen — honestly, at this point I was skeptical. However, even a negative result would answer the question of "Does RAG fine-tuning work with Mamba2 frozen + LoRA?", so I decided to proceed.

QLoRA is Incompatible with Mamba-2

If you want to save memory, you'd naturally want to use QLoRA (NF4 quantization + LoRA), but QLoRA doesn't work with Nemotron-H either. The root cause is the same as before: Mamba-2's CUDA kernels are designed to directly reference weights. bitsandbytes' NF4 quantization transforms the data type and memory layout of weights, causing a mismatch with the memory representation expected by the Mamba-2 kernel, resulting in a crash.

Moreover, even when LoRA is not applied to Mamba-2 layers, loading the model with NF4 quantizes the weights of all layers, causing the same error in the inference path. This means QLoRA is a fundamental incompatibility of "the model cannot even be loaded" — not just the issue of "there are layers where LoRA isn't applied."

Fortunately, DGX Spark is equipped with 128GB of unified memory, so even loading the model in BF16 requires only about 18GB. The decision this time was: if QLoRA can't be used, go with BF16 LoRA.

Overcoming the SM 12.1 Problem with NGC Containers

Another hurdle was GPU architecture. On the DGX Spark's GB10 GPU (Compute Capability 12.1, hereafter SM 12.1), Mamba-2's fused kernels were unsupported with pip-installed PyTorch, causing output corruption. I initially resigned myself to a hybrid configuration of "training in the cloud, inference at the edge," but this was resolved with NGC containers.

BF16 LoRA Training with NGC Containers

Container Configuration

The NGC NeMo container (nvcr.io/nvidia/nemo:25.11.01) includes SM 12.1-compatible PyTorch and CUDA 13.0, and also supports ARM64. Within the container, fused kernels work correctly, so training through inference can be completed entirely on DGX Spark. The only additional requirement is trl (SFTTrainer); peft, datasets, transformers, and others come pre-installed.

docker run --gpus all --rm --ipc=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /home/username/works:/workspace \
  -v /home/username/.cache/huggingface:/root/.cache/huggingface \
  nvcr.io/nvidia/nemo:25.11.01 \
  bash -c '
    pip install -q trl
    python3 /workspace/.../n3-nemo-train.py \
      --backend hf-peft \
      --data-file /workspace/.../train.jsonl \
      --output-dir /workspace/.../ngc-adapter
  '

Since the HuggingFace cache is mounted, there's no need to re-download the model. From container startup to training start, including trl installation, it took about 2 minutes.

BF16 LoRA Configuration

Instead of QLoRA, the model is loaded in BF16 (no quantization), and LoRA adapters are applied to all projections in the Attention layers and FFN layers.

The memory estimate breaks down as follows:

• Model (BF16): ~18GB
• LoRA adapter: ~0.5GB
• Optimizer state: ~1GB
• KV cache + activations: ~10GB
• Total: ~30GB

At approximately 30GB against DGX Spark's 128GB unified memory, there's plenty of headroom for training.

Training Results

Training completed in 138 steps over approximately 55 minutes with 1,100 samples × 1 epoch. Loss decreased from 10.71 to 6.64, and Mean Token Accuracy rose from 68.7% to 78.4%.

Around step 120, the loss bottomed out and then rose slightly. Since this is 1 epoch, it's not overfitting — it's likely the effect of the cosine learning rate schedule entering its final phase.

GGUF Conversion and Ollama Registration

The trained LoRA adapter is converted to GGUF format. llama.cpp's convert_lora_to_gguf.py is used, but since Nemotron 9B-v2-Japanese is a custom model requiring trust_remote_code=True, the key point is to directly specify the local HuggingFace cache using the --base option. Using --base-model-id causes an interactive confirmation prompt to appear, blocking script execution.

# Convert LoRA adapter to GGUF
python llama.cpp/convert_lora_to_gguf.py \
    ./ngc-adapter/adapter \
    --outfile nemotron-9b-raft-lora-ngc.gguf \
    --base ~/.cache/huggingface/hub/models--nvidia--NVIDIA-Nemotron-Nano-9B-v2-Japanese/snapshots/<hash>

The converted adapter was 36.1MB (132 tensors, BF16). This is registered to Ollama using the ADAPTER method.

# Register to Ollama using ADAPTER method
cat <<'EOF' > Modelfile-raft
FROM nemotron-9b-jp-nothink
ADAPTER nemotron-9b-raft-lora-ngc.gguf
EOF

ollama create nemotron-9b-jp-raft -f Modelfile-raft

The base model uses nemotron-9b-jp-nothink (a version with thinking disabled). There was a problem where the content of <think> tags became noise in RAG evaluation, causing EM (exact match) to become 0, so I determined that using the nothink version as the base would yield more accurate metrics.

When writing a custom template in a Modelfile, care is needed with special token handling. In Nemotron 9B-v2-Japanese, the correct representation uses single tokens like <SPECIAL_10> rather than <extra_id_0>, and using incorrect tokens reduces accuracy. By using FROM nemotron-9b-jp-nothink, the correct template already configured in a previous article is used as-is.

Evaluation Results

Quantitative Evaluation (F1)

Token-level F1 (match rate based on word overlap between prediction and ground truth) was measured on 200 test items. Exact Match (EM) was 0.0 for both models. The post-RAFT trained model answers in a citation-based CoT format of "Reference Document X states...", making exact match with the ground truth text structurally impossible. F1 better reflects practical accuracy.

F1 improved by +8.9 points. However, the even larger impact is the dramatic reduction in false positives for "cannot answer." Baseline incorrectly refused to answer 47 out of 200 items (23.5%) with "Cannot answer based on the reference documents," while RAFT FT reduced this to just 3 items (1.5%).

Looking at the top 5 items with the largest improvement, all had F1 of 0.0 in Baseline (answer refusal), recovering to around 0.8 in RAFT FT. The effect was particularly pronounced for questions with many specialized tax terms and questions involving complex conditions.

Regressions were also confirmed. The top 5 regressions show F1 decreasing by 0.3–0.4 points. These are cases where RAFT training changed the response format (to citation + reasoning CoT format), reducing the F1 for questions that the model could previously answer straightforwardly.

JCQ Regression Check

To confirm that the model's general capabilities hadn't degraded from domain-specific training, a comparison with the baseline was performed on JCommonsenseQA (1,119 questions). The baseline was measured under the same conditions as a previous article (Ollama + BF16 GGUF, thinking OFF). As confirmed in N1, there's a 7.7-point difference between inference engines (Ollama vs vLLM), so the prerequisite for comparing before and after FT is to use the same engine.

The difference is -0.5 points, essentially within the margin of error. It was confirmed that even after RAFT training specialized for the tax domain, Japanese commonsense reasoning ability was barely impaired.

Qualitative Evaluation

Characteristic cases were extracted from the test data to compare responses from Baseline and RAFT FT.

Improvement Example: Baseline Refuses to Answer → RAFT FT Gives Correct Answer

For a question where Baseline refused with "Cannot answer from the provided information," RAFT FT accurately cites the reference document and responds.

Cannot answer from the provided information.

Reference Document 2 states that "Cargo related to export declarations entrusted to AEO
customs brokers qualifies as special export cargo stipulated under Article 17, Paragraph 2,
Item 4 of the Consumption Tax Act Enforcement Order, and consumption tax is exempted as
export-equivalent transactions only for provision of services at bonded areas etc. and at
locations for loading onto vessels or aircraft for export of special export cargo, and
transportation between bonded areas etc."

3 Questions Familiar During Tax Filing Season

Since it's tax filing season, I compared some questions that might be of interest from among the FAQs.

Of Reference Documents 1 through 5, Reference Document 4 contains information relevant
to the medical expense deduction.

Reference Document 4 states the following:
"Costs for beautifying one's appearance or changing one's features do not constitute
costs for disease treatment, so the cost of mole removal is not eligible for the
medical expense deduction."

According to this statement, surgery to remove a mole is judged to fall under the
purpose of "beautifying one's appearance or changing one's features," and therefore
does not constitute medical treatment costs, making it ineligible for the medical
expense deduction.

However, if a mole is caused by skin cancer or other disease and is clearly recognized
by a physician as being for "treatment purposes," there is a possibility it could
exceptionally be eligible for the medical expense deduction, but the provided
information does not mention such cases (abbreviated below)

Reference Document 4 states that "costs for beautifying one's appearance or changing
one's features do not constitute costs for disease treatment." Mole removal is performed
for the purpose of improving one's appearance or features, and does not fall under
"disease treatment" eligible for the medical expense deduction, so it is not eligible
for the medical expense deduction.

Therefore, the cost of surgery to remove a mole is not eligible for the medical
expense deduction.

Even those who don't need to file a tax return can file a refund application if taxes
withheld at source or estimated tax payments have been overpaid (Reference Document 5).

Reference Document 5 states that "Even those who don't need to file a tax return,
if they fall under any of the following conditions and taxes withheld at source or
estimated tax payments have been overpaid, taxes will be refunded through a refund
application (refund filing)."

According to Reference Document 3, premiums for "cancer insurance" are eligible for
the life insurance premium deduction.
[Basis] Based on Income Tax Act Article 76, Paragraphs 1, 2, 6, 7, and Income Tax
Act Enforcement Order Article 208-6, Item 2...

Reference Document 3 states that "'Cancer insurance' premiums are eligible for the
life insurance premium deduction." This means that cancer insurance, like general life
insurance, qualifies as "life insurance" under the Income Tax Act and can therefore
receive the life insurance premium deduction...

Regression Example: RAFT FT Summary Is Overly Simplified

Professional athletes are treated as falling under the fifth category of business.
[Basis] Reference Document 1: "Professional athletes are treated as falling under
the fifth category of business.
[Relevant Laws and Regulations] Consumption Tax Act Enforcement Order Article 57,
Paragraph 5"

Reference Document 1 states that "professional athletes are treated as falling under
the fifth category of business." This means that in the transaction classification
under the simplified taxation system, activities performed by professional athletes
fall under the 5th category of business (supplementary to others' business activities)...

Discussion

Did RAFT Work with 53% Coverage LoRA?

With Mamba-2 layers frozen, could training only Attention and FFN change RAG response patterns? That was the core question of this article.

Based on the results, RAFT's effect was clearly apparent even with nearly half of the parameters frozen. Comparing with the inference engine difference confirmed in a previous article (7.7pp difference in JCQ between Ollama and vLLM), the F1 improvement of 8.9pp in this case is a comparable impact. This means that accuracy can be moved by the quality of training data to roughly the same degree as by the choice of inference engine.

However, looking at the breakdown of improvements, the major contribution was not so much "F1 being raised across the board" as "cases where F1 that had been zero recovered to around 0.8 as a result of drastically reduced answer refusals." This is likely the result of mixing 100 "cannot answer" samples into the RAFT training data, guiding the model toward "answering as long as the reference document contains the answer." Conversely, the discovery itself that Baseline refused to answer 47 out of 200 items is important — it shows that RAG prompts alone cannot fully control Nemotron 9B-v2-Japanese's excessive caution.

Turning this around, for general questions where Baseline was already answering correctly, the F1 improvement is almost nonexistent. The improvements in this case are concentrated on "suppression of answer refusals" and "correction toward citation-based answers," and these depend strongly on the composition of training data — especially the mixing ratio of "cannot answer" samples and the citation style of CoT. If targeting a different domain or a different issue (for example, suppressing hallucinations), the training data design would need to be revisited from scratch.

Cost-Effectiveness of 1,000-Item RAFT

Here's a summary of the costs this time.

Using cloud GPU instances would cost around $10 for equivalent training, so the benefit of being able to complete everything locally with NGC containers was significant. The electricity cost for 55 minutes of training is negligible.

However, $2 is purely the API cost (CoT generation) and compute resource cost. In practice, there's a fair amount of work involved in preparing the training data. Steps like source data filtering, distractor selection logic, CoT prompt tuning, and quality checking of generated results can be partially automated, but there are also many moments requiring judgment based on domain knowledge. When considering the cost-effectiveness of RAFT, it's realistic to include the effort of data preparation in the estimate.

Note that this time I used a cloud AI API (Claude Haiku) for CoT generation, but there will be cases where data cannot be sent externally. In those cases, running an open model of the 70B class on DGX Spark to generate CoT is a viable approach. With 128GB of unified memory, a Q4-quantized 70B model (about 40GB) fits comfortably, making it possible to build a fully local pipeline from training data generation to training to inference evaluation. Since CoT quality directly translates to training data quality, rigorous verification of generated results is necessary, but from the perspective of sovereign AI and data governance, it's an option worth considering.

Limitations and Future Work

There are several limitations to this verification.

First, since Mamba-2 layers are frozen, the model's SSM-based long-range memory capability has not been adjusted. The 53% LoRA coverage only adjusts the query-key-value computation in Attention layers and the representation transformation in FFN layers. Applying LoRA to Mamba-2's in_proj/out_proj requires either waiting for an update to HF PEFT or switching to NeMo's training pipeline. The NGC NeMo container includes Megatron-Bridge v0.2.0, which enables model conversion between HuggingFace format and Megatron format. By converting to Megatron format and using NeMo's SFT recipe, 100% LoRA coverage including Mamba-2 layers may be achievable — but I'd like to verify this in a separate article.

Another point: 1,000 items is an intentional constraint, but how far the 53%-coverage LoRA scales when increased to 5,000 or 10,000 items is unverified. With the 4B model in the previous experiment, there was a lesson that generality degraded at 10,000 items, so whether the same trend appears with the 9B model is something I'm curious about.

Summary

Nemotron 9B-v2-Japanese was RAFT fine-tuned on the National Tax Agency FAQ (JaGovFaqs-22k), and RAG accuracy was evaluated.

Thanks to the NGC container, being able to complete everything from training to evaluation on DGX Spark was a significant personal gain. From the perspective of sovereign AI and data sovereignty as well, there's a practical reassurance in being able to run domain-specific fine-tuning entirely within one's own environment. While preparing training data is no easy task, for those who want to improve RAG accuracy with their own domain data, RAFT may well be worth trying.

Data preparation scripts, training scripts, and evaluation scripts are published in the following repository.

https://github.com/himorishige/dgx-spark-blog/tree/main/n3-raft-finetuning

References

• RAFT: Adapting Language Model to Domain Specific RAG
• JaGovFaqs-22k
• PEFT Issue #2274: LoRA not working with Mamba2
• Nemotron 9B-v2-Japanese Model Card
• NGC NeMo Container