Tried training full-layer LoRA including Mamba-2 on Nemotron 9B × Megatron-Bridge using NVIDIA Brev H100

Introduction

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

In my previous article, I fine-tuned Nemotron 9B-v2-Japanese with RAFT using the NTA FAQ and improved RAG answer accuracy. However, I was always bothered by the fact that LoRA couldn't pass through the Mamba-2 layers, meaning only 53% of all parameters were trained. The results themselves weren't bad — F1 improved by +8.9 points and answer refusals dropped dramatically from 47 to 3 — but still...

Of the 56 layers in Nemotron 9B-v2-Japanese, 27 are Mamba-2 (SSM). The in_proj and out_proj of these layers have no means of applying LoRA in the HuggingFace PEFT implementation at the time of writing. Results were achieved with only the remaining Attention layers and FFN layers, but the question "would results change if the remaining 47% were also trained?" never left my mind.

In this article, I used Megatron-Bridge v0.2.0 included in the NGC NeMo container to achieve 100% LoRA coverage including Mamba-2, and ran training on NVIDIA Brev's cloud H100. I'll cover the entire pipeline from bringing the trained adapter back to DGX Spark, converting to GGUF, and running inference with Ollama. The conclusion: it didn't turn out the way I expected.

What is Megatron-Bridge

At the time of writing, HuggingFace PEFT had a limitation where there was no way to apply LoRA to in_proj and out_proj of Mamba-2. This is a problem stemming from the HuggingFace implementation, while Megatron's training pipeline natively supports these layers. In other words, you convert the model to Megatron format, train it, and then convert it back to HuggingFace format. Megatron-Bridge is what bridges that gap.

Megatron-Bridge is included in the NeMo 25.11 series containers, and with 3 steps — import_ckpt → finetune → export_ckpt — you can train HuggingFace ecosystem models using Megatron's native pipeline and bring the results back in HuggingFace format.

The NeMo container includes a dedicated recipe (nemotron_nano_9b_v2_finetune_config) with pre-tuned LoRA target modules and hyperparameters. The LoRA targets included in the recipe are the following 6 types.

HF PEFT covers only the top 4 types at 53%. Megatron-Bridge covers all 6 types at 100%.

The DGX Spark Wall and Brev

The SM 12.1 Wall

In the previous article I broke through the SM 12.1 issue with the NGC NeMo container on DGX Spark, but with Megatron-Bridge's import_ckpt, I hit another wall where multiprocessing.Manager().Queue() causes an EOFError inside a Docker container. So this time I decided it would be faster to rent a cloud H100 rather than struggle locally.

NVIDIA Brev

https://developer.nvidia.com/brev

NVIDIA Brev is a cloud GPU instance service provided by NVIDIA. You can use H100s and A100s by the hour.

With a single H100 80GB PCIe configuration, at the time of writing the price was around $2.26/hr (Hyperstack provider), which is within reach for personal verification purposes. Since I was doing LoRA training with TP=1 (no tensor parallelism) this time, one instance was sufficient.

https://dev.classmethod.jp/articles/nvidia-brev-introduction/

Choosing VM Mode

Brev offers several ways to configure GPU environments.

I initially tried to directly specify the NeMo container in Custom Container mode, but pip install wouldn't work because the Python environment inside the container was treated as externally managed. Switching to VM Mode and running docker run myself turned out to be more flexible, and I could work the same way as locally.

Brev Environment Setup

The entire pipeline is compiled into a shell script (n6-brev-run.sh).

SSH into the VM Mode instance and start the NeMo container.

# Pull and start the NeMo container
docker pull nvcr.io/nvidia/nemo:25.11.01
docker run --gpus all --ipc=host --ulimit memlock=-1 \
  --ulimit stack=67108864 -v /ephemeral:/workspace \
  nvcr.io/nvidia/nemo:25.11.01 bash /workspace/n6-brev-run.sh

/ephemeral is the temporary storage of the Brev instance, and training data and scripts are transferred here in advance using scp. --ipc=host and --ulimit memlock=-1 are flags required for Megatron's distributed training.

The H100's nvidia-smi output is as follows.

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.195.03             Driver Version: 570.195.03     CUDA Version: 13.0     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
|   0  NVIDIA H100 PCIe               On  |   00000000:00:07.0 Off |                    0 |
| N/A   28C    P0             67W /  350W |       0MiB /  81559MiB |      0%      Default |
+-----------------------------------------+------------------------+----------------------+

Converting from HF to Megatron Format

Use AutoBridge.import_ckpt() to convert the HuggingFace format model to Megatron format.

from megatron.bridge import AutoBridge

AutoBridge.import_ckpt(
    "nvidia/NVIDIA-Nemotron-Nano-9B-v2-Japanese",
    "/workspace/megatron-ckpt",
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="cpu",
)

trust_remote_code=True is required. Since Nemotron 9B-v2-Japanese is a custom model, without this it gets rejected at can_handle(). This was a subtle gotcha I noticed during a dry run.

mp.Queue() Patch

I encountered a problem where multiprocessing.Manager().Queue() causes an EOFError during the conversion save step. It seems related to the /dev/shm size limit inside Docker containers, but I prioritized a workaround over investigating the root cause.

import multiprocessing as mp
import megatron.core.dist_checkpointing.strategies.filesystem_async as fs_async

# Use plain mp.Queue() instead of Manager().Queue()
fs_async._get_write_results_queue = lambda: mp.Queue()

For single-node configurations, mp.Queue() is functionally equivalent. With this patch applied on the H100, the conversion completed in 71 seconds. The checkpoint size is 16.6GB.

100% LoRA Training

Recipe Configuration

Megatron-Bridge has training recipes for each model. For Nemotron 9B-v2-Japanese, nemotron_nano_9b_v2_finetune_config can be used.

from megatron.bridge.recipes.nemotronh import nemotron_nano_9b_v2_finetune_config
from megatron.bridge.peft.lora import LoRA

lora_config = LoRA(
    target_modules=[
        "linear_qkv", "linear_proj",    # Attention
        "linear_fc1", "linear_fc2",     # FFN
        "in_proj", "out_proj",          # Mamba-2 ← Not possible with HF PEFT
    ],
    dim=32,
    alpha=32,
)

config = nemotron_nano_9b_v2_finetune_config(
    peft=lora_config,
    pretrained_checkpoint="/workspace/megatron-ckpt",
    train_iters=500,
    micro_batch_size=1,
    global_batch_size=8,
    seq_length=2048,
    finetune_lr=1e-4,
)

Here is a summary of the configuration differences from the previous run (HF PEFT).

Eval Batch Size Issue

Megatron-Bridge's eval crashes during batch splitting if the number of data items is less than global_batch_size. This time, with 100 eval items and global_batch_size=8, there shouldn't be a problem in principle, but it got caught by a bug in the remainder handling. Since the final evaluation is done on Ollama, I worked around it by making eval a no-op.

import megatron.bridge.training.train as _train_module
_train_module.evaluate_and_print_results = lambda *a, **kw: None

Loss Progress

Here is the lm loss progression over 500 iterations.

Loss decreased from an initial 0.858 to 0.095. Compared to the previous run's (HF PEFT) final loss of 6.64, the order of magnitude is different. Since the loss calculation methods differ between training frameworks, a direct comparison isn't possible, but we can see that training itself converged smoothly.

Training took approximately 17 minutes. You really feel the speed of the H100. Training on the same dataset took 55 minutes (138 steps) on DGX Spark, yet this is the speed even with 500 iterations. GPU memory was 23.3GB / 80GB, with plenty of headroom even at TP=1.

Bringing Back the Adapter

export_ckpt source_path

Once training is complete, export to HuggingFace format. The training checkpoint (iter_0000500/) only contains LoRA delta weights (.distcp), so you pass this to export_ckpt's source_path and merge with the base checkpoint.

bridge = AutoBridge.from_hf_pretrained(
    "nvidia/NVIDIA-Nemotron-Nano-9B-v2-Japanese",
    trust_remote_code=True,
)

bridge.export_ckpt(
    "/workspace/megatron-ckpt",      # Base Megatron checkpoint
    "/workspace/hf-export",           # Export destination (HF format)
    source_path="/workspace/output/n6-megatron-lora/checkpoints/iter_0000500",
)

Note that the bridge instance should be initialized with from_hf_pretrained. Using from_hf_config results in an error because save_artifacts is not implemented. I was fortunate to catch this during the dry run.

Export completed in 50 seconds. The total work time on Brev up to this point was about 20 minutes, costing approximately $2.26/hr × 1 hour ≈ $2.26.

GGUF Conversion and Ollama Registration

Transfer the exported HF model to DGX Spark via scp and convert to GGUF.

In the previous run, I converted the LoRA adapter alone to GGUF and combined it using Ollama's ADAPTER directive. That approach can't be used this time. llama.cpp's LoRA adapter GGUF doesn't support in_proj and out_proj of Mamba-2, and these tensors are ignored during conversion.

So this time I'm directly converting the fully merged HF model to GGUF.

# Convert fully merged HF model to GGUF Q4_K_M
python3 llama.cpp/convert_hf_to_gguf.py ./hf-export/ \
    --outtype q4_k_m \
    --outfile nemotron-9b-n6-Q4_K_M.gguf

The Ollama Modelfile is as follows.

FROM ./nemotron-9b-n6-Q4_K_M.gguf

TEMPLATE """{{- range $i, $_ := .Messages }}
{{- if eq .Role "system" }}<extra_id_0>System
{{ .Content }}
{{ end }}
{{- if eq .Role "user" }}<extra_id_1>User
{{ .Content }}
{{ end }}
{{- if eq .Role "assistant" }}<extra_id_1>Assistant
{{ .Content }}
{{ end }}
{{- end }}<extra_id_1>Assistant
"""

PARAMETER stop "<extra_id_1>"
PARAMETER num_ctx 8192
PARAMETER temperature 0.6
PARAMETER top_p 0.95

ollama create nemotron-9b-n6 -f Modelfile

This enables inference with ollama run nemotron-9b-n6.

Evaluation Results

I used the exact same evaluation scripts and test data as in the previous article to compare under equal conditions.

JCQ Regression Check

I checked for regression in general-purpose capabilities using JCommonsenseQA (1,119 questions).

RAFT Domain F1 Evaluation

I measured token-level F1 on 200 items from the NTA FAQ test set.

Honestly, this result was unexpected. Despite achieving 100% LoRA coverage, F1 fell below the baseline. Answer refusals (FP) also reached 69 cases, the exact opposite of the previous run's dramatic drop from 47 to 3.

Qualitative Samples

I dug into some samples. What stood out first was "the F1 limited to samples where this model actually answered."

The accuracy when this model does answer is actually higher for Megatron-Bridge (0.7457 vs 0.6636). The problem lies in "the decision of whether to answer," where this model has become overly conservative.

Analysis

Why Accuracy Dropped Despite 100% LoRA

As seen in the qualitative samples, when this model answers, F1 is higher (0.7457 vs 0.6636). In other words, the quality of the answers themselves has improved. The reason overall F1 degraded is that 69 cases occurred where the model refused to answer despite being able to. The RAFT training data includes samples where "if the answer is not in the reference documents, respond that you cannot answer," and the model appears to have overfit to this refusal pattern.

So why didn't this happen with the same data in the previous run? I think the biggest factor is the difference in the number of epochs. In the previous run, training was stopped at 138 steps (1 epoch), but this time I ran 500 iterations using the Megatron-Bridge recipe defaults as-is. This corresponds to approximately 3.6 epochs for the 1,100-item dataset. It might be natural for the threshold to drop excessively after repeatedly training on the refusal pattern 3.6 times.

There are other possible factors. The GGUF conversion method is different (the previous run layered LoRA adapter GGUF on top of the base, while this time it's Q4_K_M quantization after full merging), and the subtle weight changes learned during training may have been rounded off by quantization. Alternatively, applying LoRA to Mamba-2's SSM layers may have prioritized memorizing training data patterns over the RAG behavior of "searching for answers in reference documents." Verifying the latter would require a comparative experiment with only the Mamba-2 layers frozen, but I wasn't able to go that far this time.

For a fair comparison with the previous run, an experiment with rank=16, 1 epoch, and the same GGUF conversion method would be needed.

Cost-Effectiveness of Brev

Here is a summary of Brev usage costs this time.

There's a possibility I could have solved it on DGX Spark with more time, but renting an H100 for $2.26/hr and breaking through the problem in one go seems like it was a reasonable decision. The training itself finished in 17 minutes, and most of the time was spent on trial and error like environment verification.

The pattern of "train in the cloud, infer at the edge" feels like a good fit with DGX Spark. Cloud H100 shines for training, while DGX Spark's 128GB unified memory shines for inference. The workflow of renting an H100 on Brev for spot training and bringing results back locally seems useful going forward, and I'm also excited about the coming soon integration between DGX Spark and NVIDIA Brev, which seems designed exactly with this workflow in mind.

Impressions of Megatron-Bridge

Honestly, Megatron-Bridge has sparse documentation. Most of the time I had to read the source code and figure things out through trial and error. The 3-step flow of import_ckpt → finetune → export_ckpt itself is simple, which is a good impression.

However, each step had its own pitfalls. import_ckpt had the Manager().Queue() EOFError, finetune had the eval data batch size mismatch, and export_ckpt had the issue where not passing source_path would corrupt the model. Each workaround was a one-liner, so I hope this article helps anyone who gets stuck in the same places. This is the experience as of NeMo 25.11.01, so future versions may have improvements. (I also suspect there may be issues with my usage as well...)

Summary

Using Megatron-Bridge v0.2.0, I performed 100% LoRA training including Nemotron 9B-v2-Japanese's Mamba-2 layers on Brev H100.

Technically, I was able to achieve 100% LoRA training including Mamba-2 using Brev's cloud H100 with Megatron-Bridge. The workflow of "train in the cloud and bring back to the edge" also worked without issues.

However, the crucial RAFT accuracy degraded compared to the previous run (53% LoRA). The lesson here is that increasing LoRA coverage doesn't necessarily improve accuracy. Since the F1 for the cases where the model did answer actually improved, if I can suppress the overfitting to the refusal pattern by aligning the number of epochs in a re-experiment, the results might change.

Megatron-Bridge itself, while having many pitfalls, felt like an interesting bridge connecting the HuggingFace ecosystem with Megatron's training pipeline. I look forward to improvements in documentation, and I hope this is helpful to anyone trying the same thing.

The full scripts are published in the following repository.

https://github.com/himorishige/dgx-spark-blog/tree/main/n6-megatron-bridge

References

• Megatron-Bridge (NGC NeMo Container)
• NVIDIA Brev
• PEFT Issue #2274: LoRA not working with Mamba2
• Nemotron 9B-v2-Japanese Model Card
• RAFT: Adapting Language Model to Domain Specific RAG