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 Department.

In my previous article, I fine-tuned Nemotron 9B-v2-Japanese with RAFT on the NTA FAQ to improve 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 being trained. The results themselves weren't bad—F1 improved by +8.9 points and answer refusals dropped dramatically from 47 to 3—but still...

In Nemotron 9B-v2-Japanese, 27 out of 56 layers are Mamba-2 (SSM). At the time of writing, the HuggingFace PEFT implementation provides no means to apply LoRA to their in_proj and out_proj. Results were achieved with just the remaining Attention and FFN layers, but the question "would results change if we trained the remaining 47% too?" 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 then brought the trained adapter back to DGX Spark, converted it to GGUF, and ran inference with Ollama—all in one end-to-end pipeline. To cut to the chase: the results didn't turn out as expected.

What is Megatron-Bridge?

At the time of writing, HuggingFace PEFT had a limitation where there was no way to apply LoRA to Mamba-2's in_proj and out_proj. This is an issue stemming from the HuggingFace implementation—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 this gap.

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

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

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

The DGX Spark Wall and Brev

The SM 12.1 Wall

Last time I broke through the SM 12.1 issue on DGX Spark using the NGC NeMo container, but with Megatron-Bridge's import_ckpt, I hit another wall where multiprocessing.Manager().Queue() throws an EOFError inside a Docker container. So this time I decided it would be faster to rent a cloud H100 rather than persist 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 on an hourly basis.

With a single H100 80GB PCIe configuration, at the time of writing it costs 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), one instance was sufficient.

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

Choosing VM Mode

Brev offers several ways to configure the GPU environment.

I initially tried to specify the NeMo container directly in Custom Container mode, but pip install didn'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 felt the same as working locally.

Brev Environment Setup

The entire pipeline was consolidated into a shell script (n6-brev-run.sh).

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

# Pull and launch 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 temporary storage on the Brev instance, where training data and scripts are pre-transferred via 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 a 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. Nemotron 9B-v2-Japanese is a custom model, and 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() throws an EOFError during the conversion save step. It seems related to the /dev/shm size limit inside the Docker container, 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 H100, conversion completed in 71 seconds. The checkpoint size is 16.6GB.

100% LoRA Training

Recipe Configuration

Megatron-Bridge provides 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 available 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 configuration differences from the previous run (HF PEFT).

Eval Batch Size Issue

Megatron-Bridge's eval crashes on batch splitting when the number of data points is less than global_batch_size. With 100 eval samples and global_batch_size=8, this shouldn't normally be a problem, but I got caught by a bug in remainder handling. Since the final evaluation is done on Ollama, I worked around this 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 Progression

Here is the lm loss progression over 500 iterations.

Loss dropped from an initial 0.858 to 0.095. Compared to the final loss of 6.64 in the previous run (HF PEFT), the magnitude is completely different. Since the loss calculation methods differ between training frameworks, a direct comparison isn't possible, but it's clear that training converged smoothly.

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

Bringing the Adapter Back

export_ckpt source_path

Once training is complete, export to HuggingFace format. The training checkpoint (iter_0000500/) only contains LoRA delta weights (.distcp), so pass this to export_ckpt's source_path while merging 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 able to catch this during a dry run, which was a relief.

Export completed in 50 seconds. Total work time on Brev up to this point was about 20 minutes, with a cost of 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.

Previously, I converted only the LoRA adapter 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 Mamba-2's in_proj and out_proj, and these tensors get 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 makes it possible to run inference with ollama run nemotron-9b-n6.

Evaluation Results

The exact same evaluation script and test data from the previous article were used to ensure comparable conditions.

JCQ Regression Check

General capability regression was verified with JCommonsenseQA (1,119 questions).

RAFT Domain F1 Evaluation

Token-level F1 was measured on the NTA FAQ test set of 200 samples.

Honestly, this result was unexpected. F1 fell below the baseline despite achieving 100% LoRA. Answer refusals (FP) also increased to 69—the exact opposite of the previous run, which dramatically reduced them from 47 to 3.

Qualitative Samples

I dug into some samples. One thing I wanted to look at was "F1 limited to samples where this model actually answered."

The accuracy when this model does answer is actually higher with Megatron-Bridge. The problem is the "whether to answer" judgment—this model is being overly conservative.

Discussion

Why Accuracy Dropped Despite 100% LoRA

As seen in the qualitative samples, when the model does answer, its F1 is higher (0.7457 vs 0.6636). That means the quality of answers itself improved. The reason overall F1 degraded is the 69 cases of refusing to answer when it could have. The RAFT training data includes samples that say "if the answer is not in the reference documents, say you cannot answer," and it seems the model over-learned 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. Previously I stopped training at 138 steps (1 epoch), but this time I ran 500 iterations using the Megatron-Bridge recipe defaults. This corresponds to approximately 3.6 epochs for the 1,100-sample dataset. It's natural that repeatedly training on the refusal pattern 3.6 times would lower the threshold excessively.

Other factors are also conceivable. The GGUF conversion method differs (previous: stacking LoRA adapter GGUF on base, this time: Q4_K_M quantization after full merge), and the subtle weight changes from training may have been rounded away by quantization. Alternatively, by passing LoRA through the Mamba-2 SSM layers, pattern memorization of training data may have taken precedence over the RAG behavior of "finding 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 aligning rank=16, 1 epoch, and the same GGUF conversion method would be needed.

Cost-Effectiveness of Brev

Here is a breakdown of Brev costs for this run.

There's a possibility the DGX Spark issues could have been resolved with more time, but I think renting an H100 for $2.26/hr to break through the problem at once was a reasonable call. Training itself finished in 17 minutes, and most of the time was trial and error for environment setup.

The "train in the cloud, infer at the edge" pattern feels like a good fit with DGX Spark. Training is where H100 speed shines, and inference is where DGX Spark's 128GB unified memory shines. A workflow of renting an H100 on Brev for spot training and bringing the results back locally is something I'll keep using, and I'm also looking forward to the coming soon integration between DGX Spark and NVIDIA Brev, which seems to be targeting exactly this workflow.

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 by trial and error. The 3-step flow of import_ckpt → finetune → export_ckpt is simple in itself, which is a positive.

However, each step had its own pitfall. import_ckpt had the Manager().Queue() EOFError, finetune had the eval data batch size mismatch, and export_ckpt would corrupt the model if source_path wasn't passed. All the workarounds fit in a single line, so I hope this article helps anyone who gets stuck in the same places. This was my experience as of NeMo 25.11.01, so future versions may have improvements. (I also suspect there may be issues with my own usage...)

Summary

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

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

However, the RAFT accuracy that matters actually degraded compared to the previous run (53% LoRA). The lesson here is that increasing LoRA coverage doesn't necessarily improve accuracy. Since the F1 when the model did answer actually improved, re-running with matched epoch counts to suppress over-learning of the refusal pattern might yield different results.

Megatron-Bridge itself has many pitfalls, but I found its positioning as a bridge connecting the HuggingFace ecosystem and Megatron's training pipeline to be interesting. I look forward to better documentation, and I hope this serves as a reference for anyone trying to do the same thing.

The full scripts are published in the repository below.

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