I tried fine-tuning NVIDIA Cosmos-Reason2-8B on DGX Spark for PPE detection

Introduction

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

You can use NVIDIA's VLM (Vision-Language Model) as-is to analyze images, but what if you want to specialize it for your own domain? For example, when you want to accurately detect "no helmet worn" in factory safety monitoring, general-purpose models often miss many cases.

This time, I tried LoRA fine-tuning Cosmos-Reason2-8B on DGX Spark to specialize it for PPE (Personal Protective Equipment) detection. I'll introduce the process of improving violation detection accuracy from 46.7% to 90.0% with 300 steps of training.

In video AI pipelines like VSS (Video Search and Summarization), there are configurations where a VLM is placed after object detection to visually verify detection results. This article focuses on the procedure for strengthening that VLM component with your own data.

Goal of This Article

Using a single DGX Spark (128GB unified memory), we'll teach Cosmos-Reason2-8B knowledge about PPE detection. The goal is to specialize a general-purpose VLM for our own domain and see how much we can improve violation detection accuracy.

Here's a rough summary of what we did, in this flow:

1. Convert the SH17 PPE dataset (YOLO format) into VLM training data
2. Run LoRA fine-tuning with TRL + PEFT
3. Check whether accuracy improved compared to the base model

Converting the PPE Dataset into VLM Training Data

SH17 Dataset

SH17 is a dataset annotating PPE wearing status at construction sites with 17 classes. The original images are construction site photos from Pexels, and the dataset includes images like these.

Image source: Pexels (Mikael Blomkvist / Abdulrhman Alkady)

The VLM looks at these images to determine who is wearing or not wearing which protective equipment.

Creating VQA Conversation Pairs from YOLO Annotations

For VLM fine-tuning, pairs of images and conversations are needed. I created a script to read PPE wearing status from YOLO bounding box coordinates and convert them into VQA pairs.

The conversion flow is as follows. The annotations for each image are parsed, and the counts of worn and unworn items are tallied from PPE class pairs (safety-helmet / no-helmet, etc.). These are then formatted into structured answers using <think> tags.

# Example of converted data
{
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "Please evaluate the PPE wearing status of the workers in this image."}
            ]
        },
        {
            "role": "assistant",
            "content": [{"type": "text", "text":
                "<think>\n3 workers can be confirmed in the image."
                " Confirmed worn: Safety glasses worn: 3 instances."
                " Violations detected: Safety jacket not worn: 4 instances.\n</think>\n"
                "There are PPE compliance violations. Safety jacket not worn: 4 instances."
            }]
        }
    ],
    "images": [image]  # PIL.Image object
}

5 question templates were prepared to add diversity to the data. Images are resized to 1280px on the long side to match VLM input requirements.

The conversion results were saved in HuggingFace Datasets format so they can be passed directly to TRL's SFTTrainer. The final data amounts are as follows.

Images that don't contain PPE annotations (person only, face only, etc.) are skipped, so the count is slightly less than the original number of images.

Running LoRA SFT on DGX Spark

Environment Setup

The environment was built based on the official Cosmos-Reason2 repository. Since DGX Spark uses CUDA 13.0 / ARM64 (aarch64), we switched to cu130 dependencies.

git clone https://github.com/nvidia-cosmos/cosmos-reason2
cd cosmos-reason2/examples/notebooks

# After changing cu128_torch28 → cu130_torch29 in pyproject.toml
uv sync

The versions of the main libraries are as follows.

Full Parameter SFT Doesn't Have Enough Memory

When I first calculated the memory estimate for full parameter SFT, it came to 110-120GB total: model weights 17GB + gradients 17GB + AdamW optimizer states 70GB (m and v in FP32). This is tight against DGX Spark's 128GB unified memory, and I judged there was a high risk of OOM when considering activation memory.

Why We Chose BF16 LoRA

So we adopted a method of training only the LoRA adapter while keeping the BF16 model frozen.

The official trl_sft.py sample uses QLoRA (4-bit quantization + LoRA), but since there was no guarantee that the quantization library bitsandbytes would work on ARM64 + CUDA 13.0, we used BF16 LoRA without quantization. Since memory is sufficient, there's no problem without quantization.

# Not using BitsAndBytesConfig
model = Qwen3VLForConditionalGeneration.from_pretrained(
    "nvidia/Cosmos-Reason2-8B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa",  # No aarch64 wheel for flash-attn
)

peft_config = LoraConfig(
    r=32,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                     "gate_proj", "up_proj", "down_proj"],
    task_type="CAUSAL_LM",
)

Increasing Batch Size Doesn't Speed Things Up

You might want to increase the batch size to speed things up when there's memory to spare... but with VLMs, it actually had the opposite effect.

With batch=4, memory usage increased 2.3x, but speed was actually slower. This is because VLMs have different token counts per image, leading to large padding overhead within a batch.

The correct approach to increase speed was to reduce the number of gradient accumulation steps, not increase batch size.

Training Configuration and Execution

Based on the batch size validation results, we settled on this final configuration.

uv run python train_lora_sft.py \
  --max-steps 300 \
  --batch-size 1 \
  --grad-accum 4 \
  --lr 2e-4 \
  --lora-rank 32 \
  --lora-alpha 32

Training Results

The loss dropped from 19.87 to 7.40 over 300 steps. Only 27% of the 128GB memory was used, so there's still quite a bit of headroom.

It dropped sharply in the first 20 steps, then gradually converged — a typical LoRA SFT curve.

Comparison Before and After Fine-Tuning

30 images were randomly selected from the test set and inference was compared with the same question.

Quantitative Evaluation

The base model had no PPE domain knowledge, so misclassifications like "No workers are shown" or "No violations are observed" were noticeable. After SFT, violation detection accuracy improved to 90%.

Inference time increased by about +13% due to the LoRA adapter overhead, but I think this is at a level that won't be a concern in practical use.

Qualitative Changes

The format of responses changed significantly. The base model answered in free-form text, but after SFT it began structuring the reasoning process with <think> tags and reporting specific violation items and counts.

The worker in the image is not wearing a helmet, but is wearing an ear-clip headset.
A helmet is an essential protective device for head protection during work, and since
it is not worn in this image, it is evaluated as a violation. On the other hand, the
ear-clip headset is a protective device necessary for voice communication, and since
it is properly worn, it is not evaluated as a violation.

<think>
3 workers can be confirmed in the image. Confirmed worn: Safety shoes worn: 1 instance, Safety glasses worn: 3 instances.
Violations detected: Safety jacket not worn: 3 instances.
</think>
There are PPE compliance violations. Safety jacket not worn: 3 instances.

The SFT model clearly separates and reports worn and unworn protective equipment. This structured output is easy to work with in downstream systems like alert notifications and monitoring dashboards.

Tendency for False Positives

The SFT model tends to lean toward "violation present." Among the 30 samples, 3 were over-detections of "actually no problem → violation present." This was influenced by the fact that 90% of the test set (27 out of 30 images) were samples with violations.

In actual operation, this is not a bad direction for safety monitoring. "False alarms" carry lower risk than "missed detections." However, it should be possible to further improve by increasing the number of compliant (no problem) samples to create better balance.

Pain Points

flash-attention Doesn't Work

DGX Spark uses ARM64 (aarch64), but official flash-attention wheels are only for x86_64. In the cosmos-reason2 dependency definition, it is also excluded with platform_machine != 'aarch64'. By specifying attn_implementation="sdpa", you can fall back to PyTorch's standard Scaled Dot-Product Attention.

Note that I confirmed in separate testing that flash-attn 2.7.4 can be installed on DGX Spark using the sbsa-targeted wheel published by Jetson AI Lab. It worked without issues for inference, but stability during long training runs with SFTTrainer has not been verified, so we adopted SDPA this time. If flash-attn can be used, training time reduction can be expected, and we'd like to try it in the future.

torch_dtype Deprecation Warning

The torch_dtype parameter has been deprecated, and a warning appears to use dtype instead. This is a change from transformers 4.57 onwards. It doesn't affect operation, but for new code it's better to use dtype.

Summary

Using a single DGX Spark, we ran LoRA SFT on Cosmos-Reason2-8B and improved PPE violation detection accuracy from 46.7% to 90.0%. Only 27% of the 128GB unified memory was used, leaving plenty of room for experiments with larger models or batch sizes.

The fine-tuning pipeline itself is simple, following the flow of YOLO annotations → VQA conversion → TRL SFTTrainer. VLM-specific considerations include the fact that increasing batch size doesn't speed things up (padding overhead) and handling the lack of flash-attention support in ARM64 environments.

The LoRA adapter from this time can be incorporated into the second stage of a two-stage pipeline, for example: DeepStream object detection (1st stage) → VLM visual verification (2nd stage). I hope this serves as a reference for a workflow of gradually growing a general-purpose VLM with your own on-site data.

Reference Links

• nvidia-cosmos/cosmos-reason2 — Official repository (SFT / GRPO recipes)
• Cosmos Cookbook: Post-training — Official Cookbook
• SH17 Dataset — PPE detection dataset
• TRL SFTTrainer — TRL official documentation
• DGX Spark Playbooks: VLM Fine-tuning — NVIDIA official Playbook