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 fine-tuning Cosmos-Reason2-8B with LoRA on DGX Spark to specialize it for PPE (Personal Protective Equipment) detection. I'll introduce the process by which violation detection accuracy improved 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 downstream of object detection to visually verify detection results. This article focuses on the procedure for reinforcing that VLM component with your own data.

Goal for This Time

Using a single DGX Spark (128GB unified memory), we 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 violation detection accuracy can be improved.

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 with 17-class annotations of PPE wearing status at construction sites. The original images are construction site photos from Pexels, and the dataset includes images like this:

Image source: Pexels (Mikael Blomkvist / Abdulrhman Alkady)

The VLM looks at these kinds of images and determines who is wearing or not wearing each piece of protective equipment.

Creating VQA Conversation Pairs from YOLO Annotations

For VLM fine-tuning, pairs of images and conversations are needed. I created a script that reads PPE wearing status from YOLO bounding box coordinates and converts it into VQA pairs.

The conversion flow is as follows: parse the annotations for each image, count the number of worn and unworn items from PPE class pairs (safety-helmet / no-helmet, etc.), and format them into answers structured with <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 wearing: 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 template patterns were prepared to add diversity to the data. Images are resized to 1280px on the long side to match the VLM's input requirements.

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

Images without PPE annotations (person only, face only, etc.) are skipped, so the count is slightly lower 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), the dependencies were switched to cu130.

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, the total came to 110–120GB: model weights 17GB + gradients 17GB + AdamW optimizer states 70GB (m and v in FP32). Given DGX Spark's 128GB unified memory, this was cutting it close, and I judged the risk of OOM to be high when considering activation memory.

Why I Chose BF16 LoRA

So I went with a method of freezing the BF16 model and only training the LoRA adapter.

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, I used BF16 LoRA without quantization. Since there's plenty of memory, this doesn't pose a problem.

# 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 batch size to speed things up when memory is available, but for VLMs this actually backfires.

With batch=4, memory usage increased 2.3x, but speed actually got slower. This is because VLMs have varying token counts per image, making padding overhead within a batch significant.

If you want to improve speed, the right approach is to reduce the number of gradient accumulation steps rather than increasing the batch size.

Training Settings and Execution

Based on the batch size verification results, I settled on these final settings:

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

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's a typical LoRA SFT curve where it drops sharply in the first 20 steps and then gradually converges.

Comparison Before and After Fine-Tuning

30 images were randomly selected from the test set and inference was compared using the same questions.

Quantitative Evaluation

The base model had no knowledge of the PPE domain, so misclassifications like "No workers are visible" 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 a negligible level for practical use.

Qualitative Changes

The response format 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 over-the-ear 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 over-the-ear
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 wearing: 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 handle for integration with downstream systems such as alert notifications and monitoring dashboards.

Trends in False Detections

The SFT model tends to lean toward "violation present." Of 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" are lower risk than "missed detections." However, further improvement should be possible by adding more compliant (no problem) samples to balance the dataset.

Stumbling Points

flash-attention Cannot Be Used

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

Note that in separate testing, I confirmed 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 since stability during long-duration training with SFTTrainer was unverified, I adopted SDPA this time. If flash-attn can be used, it would be expected to reduce training time, so I'd like to try it in the future.

Deprecation Warning for torch_dtype

The torch_dtype parameter has been deprecated, and a warning is displayed 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, I ran LoRA SFT on Cosmos-Reason2-8B and was able to improve PPE violation detection accuracy from 46.7% to 90.0%. Since only 27% of the 128GB unified memory was used, there's still room to experiment with larger models and batch sizes.

The fine-tuning pipeline itself is simple, following the flow of YOLO annotation → VQA conversion → TRL SFTTrainer. VLM-specific points to note are that increasing batch size does not speed things up (due to padding overhead), and dealing with flash-attention incompatibility in ARM64 environments.

The LoRA adapter from this time can be incorporated as the second stage of a two-stage pipeline, for example: DeepStream object detection (stage 1) → VLM visual verification (stage 2). I hope this serves as a useful reference for a workflow of gradually training 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