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

Introduction

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

In a previous article, when I tried RAG (Retrieval-Augmented Generation) with Nemotron 9B-v2-Japanese, I ran a test passing fictional internal company regulations as "reference documents" and having the model answer questions. Even the base model picked up on the context reasonably well, but there were instances where it returned hallucinations (plausible-sounding falsehoods) when the reference documents didn't contain the answer.

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

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

With the trend toward sovereign AI, the need to specialize models for specific domains without sending proprietary data externally is growing. This time, by leveraging NGC containers, I was able to complete everything from training to inference evaluation entirely on DGX Spark. No cloud GPU was needed, and the only cost incurred was $2 for CoT generation. I hope this is useful 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 find the correct location within those documents depends on the model's general reading comprehension ability. RAFT intervenes at this point to directly teach the skill of "citing the correct answer from relevant documents."

The training data is constructed as follows:

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

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." Because it gets accustomed to situations where distractors are mixed in, the goal is for it to be able to respond robustly even when search precision in an actual RAG pipeline is not perfect.

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 collecting approximately 22,800 FAQs published by various Japanese government ministries and agencies, with question-answer pairs.

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

Exploring the data, most answer texts fall within the 50–2,000 character range, which is a good size for RAG reference documents. Since the copyright field contains the ministry/agency name, you can use this to narrow down by field.

Sampling

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

This time, I focused on concentrating on National Tax Agency FAQs. Rather than spreading thinly across 41 ministries, by concentrating 1,000 samples in the single domain of taxation, the RAFT distractors (dummy documents that don't contain the answer) naturally become "same tax topic but different answer" hard negatives. There are approximately 2,476 National Tax Agency FAQs after filtering, so there's more than enough room to sample 1,000.

For testing, I separately secured 200 samples from the National Tax Agency, and additionally carved out 100 FAQ samples from other ministries as test-unseen. This is to see how the tax-specialized model behaves 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 within the same National Tax Agency as distractors (documents from the same tax domain have closely related topics, 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. Additionally mix in 100 "cannot answer" samples, creating a total of 1,100 training samples

The CoT answer generation cost was approximately $2 with Claude Haiku. After generation, I randomly extracted 50 samples for manual checking to verify citation accuracy and format consistency. Local models had somewhat unstable accuracy, so I used cloud models only for creating the 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 include...' Based on this provision, ...Therefore, the eligible scope is..."
    }
  ]
}

Mamba2 + LoRA Constraints

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 trying to train with LoRA, a problem arises here. The in_proj and out_proj of Mamba-2 layers are not supported by PEFT library's LoRA (PEFT Issue #2274). The reason is that Mamba-2 is designed with proprietary CUDA kernels that directly reference weight tensors, bypassing the forward hooks that LoRA inserts.

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

Whether RAFT's effects would appear with nearly half of the parameters frozen. Honestly, at this point I was skeptical. However, even a negative result would be an answer to the question "Does RAG FT work with Mamba2 frozen + LoRA?", so I decided to proceed.

QLoRA Is Incompatible with Mamba-2

If you want to save memory, you'd want to use QLoRA (NF4 quantization + LoRA), but QLoRA also doesn't work with Nemotron-H. The root cause is the same as before — Mamba-2's CUDA kernels are designed to directly reference weights. Bitsandbytes' NF4 quantization converts 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.

Furthermore, even if 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 pass. In other words, QLoRA is a fundamental incompatibility of "cannot load the model in the first place," not just an issue of "there are layers without LoRA applied."

Fortunately, DGX Spark has 128GB of unified memory, so loading the model in BF16 only requires about 18GB. In this case, the decision was: if QLoRA can't be used, let's go with BF16 LoRA.

Breaking Through the SM 12.1 Problem with NGC Containers

Another obstacle is the GPU architecture. On the DGX Spark's GB10 GPU (Compute Capability 12.1, hereafter SM 12.1), with PyTorch installed via pip, the Mamba-2 fused kernels are unsupported and output becomes corrupted. Initially I was prepared for a hybrid configuration where "training is on the cloud and inference is at the edge," but the NGC container solved this.

BF16 LoRA Training with NGC Container

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. Inside the container, fused kernels also work correctly, so training through inference can be completed entirely on DGX Spark. The only additional requirement is trl (SFTTrainer); peft, datasets, transformers, etc. are 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 the start of training, including the 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 and FFN layers.

The memory estimate is as follows:

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

With about 30GB against DGX Spark's 128GB unified memory, there's plenty of room for training.

Training Results

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

Loss bottomed out around step 120, then rose slightly. Since it's 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

Converting the trained LoRA adapter to GGUF format. Using llama.cpp's convert_lora_to_gguf.py, but since Nemotron 9B-v2-Japanese is a custom model requiring trust_remote_code=True, the key point is specifying the local HuggingFace cache directly with the --base option. Using --base-model-id causes an interactive confirmation prompt that blocks 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 (the version with thinking disabled). Because there was a problem where the content of <think> tags became noise during RAG evaluation and EM (exact match) became 0, I decided that using the nothink version as a base would yield more accurate metrics.

When writing custom templates in Modelfile, attention is needed for handling special tokens. In Nemotron 9B-v2-Japanese, single tokens like <SPECIAL_10> are the correct representation rather than <extra_id_0>, and using incorrect tokens degrades accuracy. By using FROM nemotron-9b-jp-nothink and inheriting the correct template already configured in a previous article, this is handled properly.

Evaluation Results

Quantitative Evaluation (F1)

Token-level F1 (match rate based on word overlap between prediction and correct answer) was measured on 200 test samples. Exact Match (EM) is 0.0 for both models. After RAFT training, the model answers in a citation-based CoT format of "Reference Document X states ~," so exact match with the correct answer text is structurally impossible. F1 better reflects actual accuracy.

F1 improved by +8.9 points. However, what's even more impactful is the drastic reduction in false positives for "cannot answer." While the Baseline incorrectly refused to answer 47 out of 200 cases (23.5%) with "Cannot answer from the reference documents," RAFT FT reduced this to just 3 cases (1.5%).

Looking at the top 5 cases with the largest improvements, all of them had an F1 of 0.0 (answer refusal) in the Baseline, recovering to around 0.8 with RAFT FT. The effect appears to be particularly pronounced for questions containing many specialized tax terminology or questions with complex conditions.

On the other hand, regressions were also confirmed. In the top 5 regression cases, F1 dropped by 0.3–0.4 points. These are cases where the F1 of questions that the Baseline could originally answer straightforwardly decreased because RAFT training changed the answer format (to citation + reasoning CoT format).

JCQ Regression Check

To confirm that the model's general-purpose capabilities hadn't degraded due to domain-specific training, I compared against the baseline using JCommonsenseQA (1,119 questions). The baseline was measured under the same conditions as a previous article (Ollama + BF16 GGUF, thinking OFF). Since a 7.7-point difference due to inference engine is confirmed as noted in N1, the premise is that pre/post-FT comparisons 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

I extracted characteristic cases from the test data to compare the responses of Baseline and RAFT FT.

Improvement Example: Baseline Refuses to Answer → RAFT FT Answers Correctly

For questions where the Baseline refused with "Cannot answer from the provided information," RAFT FT accurately responds by citing the reference documents.

Cannot answer from the provided information.

Reference Document 2 states: "Cargo subject to export declarations entrusted to AEO
customs brokers falls under special export cargo as defined in Article 17, Paragraph 2,
Item 4 of the Consumption Tax Law Enforcement Order, and only services provided at bonded
areas, etc., at the location of loading onto ships or aircraft for export of special export
cargo, and transportation between bonded areas are exempt from consumption tax as
export-equivalent transactions."

3 Questions Familiar During Tax Filing Season

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

Of Reference Documents 1 through 5, the one containing information related to the
medical expense deduction is Reference Document 4.

Reference Document 4 states the following:
"Expenses for beautifying one's appearance or changing one's looks do not qualify as
expenses for treating illness, so mole removal costs are not eligible for the medical
expense deduction."

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

However, if the mole is caused by skin cancer or other diseases 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 "expenses for beautifying one's appearance or changing one's
looks do not qualify as expenses for treating illness." Mole removal is performed for the
purpose of improving one's appearance or looks, and since it does not fall under "disease
treatment" eligible for the medical expense deduction, 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 claim if income tax
withheld at source or advance tax payments have been overpaid (Reference Document 5).

Reference Document 5 states: "Even those who do not need to file a tax return, if they
fall under any of the following or similar cases and income tax withheld at source or
advance tax payments have been overpaid, a tax refund can be received through a refund
claim (return for refund)."

According to Reference Document 3, the premium for "cancer insurance" is 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 "the premium for 'cancer insurance' is eligible for the
life insurance premium deduction." This means that cancer insurance, like ordinary life
insurance, qualifies as "life insurance" under the Income Tax Act, and therefore can
receive the life insurance premium deduction...

Regression Example: RAFT FT's Summary Is Overly Simplified

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

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

Discussion

Did RAFT Work with 53% Coverage LoRA?

With Mamba-2 layers frozen, can changing only the Attention and FFN layers alter RAG's response patterns? That was the core question of this article.

To state the conclusion: even with nearly half the parameters frozen, RAFT's effects were clearly apparent. Compared to the inference engine difference confirmed in a previous article (7.7pp for JCQ between Ollama vs vLLM), today's F1 improvement of 8.9pp is comparable in impact. This means that training data quality can move accuracy by about as much as the choice of inference engine.

That said, looking at the breakdown of the improvement, it's not so much "F1 was raised across the board" but rather that cases where "F1 was 0 due to answer refusal recovered to around 0.8" contributed significantly. By mixing 100 "cannot answer" samples into the RAFT training data, the model was guided toward "answering as long as the reference documents contain an answer." Conversely, the discovery that the Baseline refused to answer 47 out of 200 cases (23.5%) is itself significant — it shows that RAG prompts alone cannot control Nemotron 9B-v2-Japanese's excessive caution.

Flipping this around, for general questions that the Baseline was already answering correctly, there is almost no improvement in F1. The improvements in this case are concentrated on "suppression of answer refusal" and "correction to citation-based answers," and these are strongly dependent on the composition of the training data, particularly the mixing ratio of "cannot answer" samples and the citation style of CoT. To target a different domain or a different issue (such as hallucination suppression), the training data design would need to be revisited from scratch.

Cost-Effectiveness of 1,000-Sample RAFT

Here's a breakdown of the costs for this implementation:

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

However, the $2 is only for the API cost (CoT generation) and computational resources. In practice, preparing the training data takes considerable effort. Steps like filtering source data, distractor selection logic, CoT prompt tuning, and quality checking of generated results can be partially automated, but there are many situations that require judgment based on domain knowledge. When thinking about RAFT's cost-effectiveness, it's realistic to include the effort for data preparation in the estimate.

Note that while I used a cloud AI API (Claude Haiku) for CoT generation this time, there will be cases where sending data externally is not possible. In such cases, running a 70B-class open model on DGX Spark to generate CoT is a conceivable approach. With 128GB of unified memory, a Q4-quantized 70B model (approximately 40GB) fits comfortably, so it's possible to build a pipeline closed locally from training data generation through training and inference evaluation. Since CoT quality directly translates to training data quality, thorough 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

This verification has several constraints.

First, as long as Mamba-2 layers are frozen, the model's SSM-based long-range memory capability cannot be touched. The 53% LoRA coverage only adjusts the Attention layer's query-key-value computation and FFN layer's representational transformation. To apply LoRA to Mamba-2's in_proj/out_proj, you'd need to wait for an HF PEFT update or switch to NeMo's training pipeline. The NGC NeMo container includes Megatron-Bridge v0.2.0, which enables model conversion between HuggingFace and Megatron formats. Converting to Megatron format and using NeMo's SFT recipe may enable 100% LoRA coverage including Mamba-2 layers — but I'd like to verify this in a separate article.

Another thing: while 1,000 samples is an intentional constraint, it's untested how much a 53% coverage LoRA scales when increased to 5,000 or 10,000 samples. There's a lesson from the previous 4B model where generalizability degraded at 10,000 samples, so I'm curious whether the same trend would appear with the 9B model.

Summary

I performed RAFT fine-tuning of Nemotron 9B-v2-Japanese using National Tax Agency FAQs (JaGovFaqs-22k) and evaluated RAG accuracy.

Thanks to the NGC container, being able to complete everything from training to evaluation entirely on DGX Spark was a major personal achievement. From a sovereign AI and data sovereignty perspective as well, being able to run domain-specialized fine-tuning entirely within your own environment offers real practical peace of mind. While preparing training data is no easy task, RAFT is worth trying for those who want to improve RAG accuracy with their own domain data.

The data preparation scripts, training scripts, and evaluation scripts are published in the repository below.

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