🚀 NVIDIA Nemotron-3-Nano LoRA Fine-Tuning Guide with Megatron Bridge

This notebook walks you through fine-tuning the NVIDIA Nemotron-3-Nano-30B model from start to finish—using LoRA (Low-Rank Adaptation) so you train only a small set of parameters. In this notebook you will train with Megatron Bridge, part of the NeMo framework.

📋 What You're Working With

✅ Prerequisites

💻 Hardware

• 8× GPUs — NVIDIA H100 or A100
• 250 GB free storage (minimum)

📦 Software

• OS: Ubuntu 22.04
• GPU driver: 580 or newer
• CUDA: 12.8 or newer
• NVIDIA Container Toolkit (for Docker + GPU)

🗺️ Workflow at a Glance

Follow the steps below in order. Let's go! 👇

Step 1: Download the NeMo Docker Container

Let's begin by obtaining the official NVIDIA NeMo container, which comes preloaded with everything needed for training. Before proceeding, make sure to set your NGC_API_KEY.

Setup NGC_API_KEY

NGC offers a wide variety of public images, models, and datasets and you'll need to generate an API key and authenticate with NGC.

To create your API key, visit: https://org.ngc.nvidia.com/setup/api-keys

When generating your NGC key, make sure to enable the "NGC Catalog" under "Services Included".

Now let's download the NVIDIA Nemotron-3 Nano container from NGC

Step 2: Launch the Docker Container

Launch the Docker container with GPU capabilities and mount your current directory as the workspace.

Key options explained:

• --gpus all: Grants the container access to all available GPUs.
• --ipc=host: Shares the host’s IPC namespace for improved multi-GPU support.
• --network host: Uses the host machine's network settings.
• -v $(pwd):/workspace: Mounts your present directory to /workspace inside the container.
• -p 8080:8080 -p 8088:8088: Opens essential ports for monitoring and service access.

Note: Run this command in your terminal (not in Jupyter). The command will start an interactive shell session inside the container.

Note : The following steps should be executed inside the Docker container.

Step 3: Set Environment Variables

Set the HuggingFace model ID and specify the destination for the Megatron checkpoint.

Step 4: Convert HuggingFace Model to Megatron Format

NeMo uses Megatron-LM format for training. We need to convert the HuggingFace checkpoint to Megatron format.

This step:

• Downloads the model from HuggingFace Hub
• Converts model weights to Megatron-compatible format
• Saves to the specified output path

⏱️ Note: The first run downloads ~60GB+ and can take 15–60+ minutes depending on your connection. The progress bar may stay at 0% for a while before moving—this is normal. Do not interrupt the cell.

Step 5: Fine-tune with LoRA

In this step, you will fine-tune the model using LoRA (Low-Rank Adaptation), which is an efficient technique for adapting large models with fewer trainable parameters. We'll use the SQuAD dataset for this example. The SQuAD dataset is a popular benchmark for question answering, containing over 100,000 question-and-answer pairs across more than 500 diverse articles.

Key training parameters:

• --peft lora: Activates LoRA for efficient fine-tuning.
• train.global_batch_size=128: Sets the total batch size combining all GPUs.
• train.train_iters=50: Determines the number of training iterations.
• scheduler.lr_warmup_iters=10: Number of iterations to gradually increase (warm up) the learning rate at the start of training.
• checkpoint.pretrained_checkpoint: Specifies the path to the pre-converted Megatron checkpoint to start from.

Hardware note: The sample command uses 8 GPUs (--nproc-per-node=8). If you have a different number of GPUs, adjust this parameter as needed.

What to expect: After training, you will see output similar to: validation loss at iteration 50 on validation set | lm loss value: 1.261660E-01 | lm loss PPL: 1.134470E+00 | This indicates the loss and perplexity on the validation set after 50 iterations. For this setup, the training should complete in about 15 minutes.

Optional Step 5a – Fine-Tune with a Custom Training Script and Dataset

To use your own dataset rather than the default SQuAD dataset, you can write a custom Python training script.
This approach allows you to customize the training setup to fit your specific needs.
Run the steps below directly on the machine.

Prepare Your Dataset

This step uses the BIRD SQL dataset (a text-to-SQL benchmark with schema, question, evidence, and SQL pairs) some helper scripts. These scripts apply the Nemotron chat template, filters by sequence length, and writes training.jsonl into a dataset directory.

Run the cells below to download/prepare the dataset and save it to dataset/training.jsonl.

Create Custom Training Script

Create a Python script that configures the training with your custom dataset.
This also should be run directly on the host machine, but the paths used will be the mount points in the container.
No need to change paths in the script.

Key Configuration Options

You can customize training by setting environment variables:

Paths:

• BASE_MODEL_PATH: Path to converted Megatron checkpoint
• DATASET_DIR: Directory containing training.jsonl
• CHECKPOINT_DIR: Where to save training checkpoints

LoRA Parameters:

• LORA_RANK: LoRA rank (default: 16, higher = more parameters)
• LORA_ALPHA: LoRA alpha scaling (default: 32)
• LORA_DROPOUT: LoRA dropout rate (default: 0.05)

Launch Custom Training

Now run your custom training script with torchrun

Training Parameters:

• N_DEVICES: Number of GPUs (must match --nproc-per-node)
• GLOBAL_BS: Global batch size across all GPUs
• PER_DEVICE_BS: Batch size per GPU
• EPOCHS: Number of training epochs
• LR: Learning rate (default: 5e-5)
• WARMUP_RATIO: Fraction of steps for learning rate warmup (default: 0.03)

Run the steps below from inside the container. Your data should already be at /workspace/dataset

Step 6: Check Training Outputs

Once training is finished, check that the checkpoints have been saved successfully.

Expected output:

iter_0000050
latest_checkpointed_iteration.txt
latest_train_state.pt

If you used custom data (5a) you will see more iterations and you can view latest_checkpointed_iteration.txt to view which folder you should use.
Note the name of the folder of this checkpoint for the next step

Step 7: Merge LoRA Weights

To obtain a standalone fine-tuned model, merge the LoRA adapters into the base model.

In this step:

• The base model weights are combined with the LoRA adapter weights
• The result is a single, merged checkpoint

Again, make sure to change --nproc_per_node=8 to your GPU count.

You may need to change --lora-checkpoint if you used your own data
to the checkpoint to your latest checkpoint you would like to use.

There is no need to change the --output directory, but if you do make sure
the path you choose is used in the remainder in the steps.

Step 8: Export to HuggingFace Format

Finally, convert the merged Megatron checkpoint back to HuggingFace format for easy deployment and inference.

This creates a standard HuggingFace model that can be loaded with transformers library.

Expected output:

Converting to HuggingFace ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% (6231/6231)

Success: All tensors from the original checkpoint were written.
✅ Successfully exported model to: /workspace/models/merged_0050-hf

Step 9: Deploy the Fine-tuned Model with Docker Compose

Now that you have your fine-tuned model, you can deploy it for inference via NVIDIA NIM (NVIDIA Inference Microservices) or vLLM.

You can exit the NeMo container that you ran commands in.

Deployment Options

The docker-compose configuration below provides two deployment options:

1. NVIDIA NIM

   • Uses your merged weights
   • Supports tensor parallelism
   • Optimized for production inference

2. vLLM - Open-source inference

   • Fast and memory-efficient
   • Supports LoRA adapters
   • PagedAttention for throughput

Prerequisites

Before deploying, ensure:

• Your fine-tuned model is accessible on the host machine
• You have NGC API key (for NIM services - requires an)
• Docker Compose is installed
• NVIDIA Docker runtime is configured

The prerequisites listed above pertain to your local environment. The following commands should be executed from this notebook, outside the NeMo container, directly on your host machine.

Create docker-compose.yml

Specify the HOST_MODELS_DIR where your models are stored on your host machine.

Create a docker-compose.yml file with the following configuration:

After running this cell the docker-compose.yml file should be created in the current directory.

Start the Inference Service

Choose your deployment backend by commenting or uncommenting the relevant lines below—start either the NIM or vLLM service as needed.
Reminder: if you would like to use NVIDIA NIM a license is required. See documentation on how to obtain a free license here.

Deployment Configuration Tips

GPU Allocation:

• Adjust device_ids based on available GPUs
• Use nvidia-smi to check GPU availability
• Avoid overlapping GPU assignments between services

Tensor Parallelism:

• NIM_TENSOR_PARALLEL_SIZE=2 splits model across 2 GPUs
• Adjust based on model size and GPU memory
• Higher TP = more GPUs, better for large models

Memory Settings:

• shm_size: 128GB provides shared memory for IPC
• Increase if you encounter "Bus error" or shared memory issues
• Must be large enough for model weights and KV cache

LoRA Configuration:

• NIM_MAX_LORA_RANK=64 sets maximum adapter rank
• NIM_PEFT_REFRESH_INTERVAL controls adapter reload frequency
• Place adapters in NIM_PEFT_SOURCE directory

Performance Tuning:

• Monitor with docker stats or nvidia-smi
• Adjust batch sizes via environment variables
• Enable NIM_KV_CACHE_REUSE for better throughput

Test the Deployed Model

Once the service is running, you can test it using the OpenAI-compatible API:

Send a request to NVIDIA NIM - Nemotron-3-Nano (Custom fine-tuned model) which is running on port 8007.

Alternative: Test with cURL

You can also test using curl from the command line to the NVIDIA NIM - Nemotron-3-Nano customized model.

If you want to test the vLLM service, uncomment the following cell and run it.

Monitor and Manage Services

Export Structure

The exported HuggingFace model contains:

📁 /workspace/models/merged_0050-hf/
   📄 config.json                         # Model configuration
   📄 generation_config.json             # Generation parameters
   📄 tokenizer.json                     # Tokenizer vocabulary
   📄 tokenizer_config.json              # Tokenizer configuration
   📄 special_tokens_map.json            # Special tokens mapping
   📄 chat_template.jinja                # Chat template
   📄 model.safetensors.index.json       # Model sharding index
   📄 model-00001-of-00013.safetensors   # Model weights (sharded)
   📄 model-00002-of-00013.safetensors
   ... (13 shard files total)
   📄 modeling_nemotron_h.py             # Custom model code
   📄 configuration_nemotron_h.py        # Custom config code

Additional Resources

• Model Collection: NVIDIA Nemotron V3 on HuggingFace
• Base Model: NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
• NeMo Framework: NVIDIA NeMo Documentation

Tips and Best Practices

Training Configuration

• Batch Size: Adjust train.global_batch_size based on GPU memory
• Iterations: Increase train.train_iters for better convergence
• Learning Rate: Tune via optimizer.lr and scheduler.lr_warmup_iters

LoRA Parameters

• Default LoRA rank is typically 8-16 (configured in recipe)
• Lower rank = fewer trainable parameters, faster training
• Higher rank = more expressivity, potentially better results

GPU Requirements

• This example uses 8 GPUs
• For fewer GPUs: adjust --nproc-per-node and reduce batch size
• Monitor GPU memory with nvidia-smi

Checkpoint Management

• Checkpoints are saved to /opt/Megatron-Bridge/nemo_experiments/
• Use checkpoint.save_interval to control checkpoint frequency
• Keep at least 2-3 checkpoints for rollback

Troubleshooting

Out of Memory (OOM)

• Reduce train.global_batch_size
• Enable gradient checkpointing
• Use fewer GPUs with tensor parallelism

Slow Training

• Check GPU utilization with nvidia-smi
• Verify data loading isn't a bottleneck
• Ensure --ipc=host flag is used

Conversion Errors

• Verify HuggingFace model ID is correct
• Check disk space for checkpoint storage
• Ensure --trust-remote-code is set for custom models

Conclusion

You now have a complete end-to-end workflow for fine-tuning and deploying NVIDIA Nemotron-3-Nano models!

What You've Accomplished:

✅ Fine-tuning: Trained a custom model using LoRA for efficient adaptation ✅ Model Export: Converted to HuggingFace format at /workspace/models/merged_0050-hf ✅ Deployment: Set up production-ready inference with NVIDIA NIM or vLLM

Next Steps:

Your fine-tuned model can now be:

• Deployed: Already configured with docker-compose for immediate use
• Integrated: OpenAI-compatible API for easy integration
• Shared: Upload to HuggingFace Hub for team collaboration
• Improved: Further fine-tune with additional domain-specific data
• Scaled: Deploy across multiple nodes for high-throughput serving

API Endpoints:

Once deployed, your services are available at:

• NIM Nemotron-3-Nano: http://localhost:8007/v1/chat/completions
• vLLM: http://localhost:8006/v1/chat/completions

Happy training and deploying! 🚀