Skip to content

Cosmos-Reason2 (VLM)

Overview

Cosmos-Reason2 is the next generation of Cosmos-Reason1, a reasoning vision-language model (VLM) for Physical AI that understands space, time, and fundamental physics in video and image inputs. It is post-trained from Qwen3-VL-Instruct with physical common sense and embodied reasoning data, and emits a long chain-of-thought trace before its final answer, which lets it serve as a planning model for an embodied agent. The model family consists of 2B, 8B, and 32B variants, and supports a multimodal context of up to 256K tokens. This tutorial demonstrates how to compile Cosmos-Reason2-2B for multiple RBLN NPUs using RBLN Optimum (optimum-rbln), and answer a question about an input video.

Info

Rebellions Scalable Design (RSD) is available on ATOM™+ (RBLN-CA22) and ATOM™-Max (RBLN-CA25). You can check the type of your current RBLN NPU using the rbln-stat command.

Setup & Installation

Before you begin, ensure that your system environment is properly configured and that all required packages are installed. This includes:

Note

  • RBLN Compiler requires an RBLN Portal account.
  • The commands above are intended for a default pip install on Debian-based Linux such as Ubuntu. For all other configurations, refer to the Installation Guide for the supported install matrix and the applicable commands.

Note

nvidia/Cosmos-Reason2-2B on HuggingFace is a gated model. Once access is granted by agreeing to the NVIDIA Open Model License, you can log in using the hf auth (huggingface-cli) command as shown below:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
$ hf auth login

    _|    _|  _|    _|    _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|_|_|_|    _|_|      _|_|_|  _|_|_|_|
    _|    _|  _|    _|  _|        _|          _|    _|_|    _|  _|            _|        _|    _|  _|        _|
    _|_|_|_|  _|    _|  _|  _|_|  _|  _|_|    _|    _|  _|  _|  _|  _|_|      _|_|_|    _|_|_|_|  _|        _|_|_|
    _|    _|  _|    _|  _|    _|  _|    _|    _|    _|    _|_|  _|    _|      _|        _|    _|  _|        _|
    _|    _|    _|_|      _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|        _|    _|    _|_|_|  _|_|_|_|

    To login, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens.
Enter your token (input will not be visible):

Running with the RBLN Optimum (optimum-rbln) library

This step demonstrates how to compile the Cosmos-Reason2 model to run on multiple RBLN NPUs.

Model Compilation for Multiple NPUs

To begin, import the RBLNAutoModelForImageTextToText class from optimum-rbln. This class's from_pretrained() method downloads the Cosmos-Reason2 model from the HuggingFace Hub and compiles it using the RBLN Compiler. Cosmos-Reason2 follows the Qwen3-VL structure, in which the vision encoder is declared as the visual submodule of the language model, so the rbln_config below is nested accordingly. To compile the model, set the export argument to True. The example below compiles the vision encoder and the language model for 8 NPUs each, and sets create_runtimes to False at both levels so that compilation does not allocate NPU runtimes.

After compilation, save the model artifacts to disk using the save_pretrained() method. This creates a directory (e.g., rbln-cosmos-reason2) containing the compiled model.

Tip

Each field of the nested rbln_config, along with bucketing options that cover several input resolutions or decoder batch sizes in a single compilation, is documented in Qwen3-VL-2B (VLM).

Note

Cosmos-Reason2-8B shares the same architecture and the same compilation settings, so change model_id to compile it. Cosmos-Reason2-32B has a larger language model, so set its num_devices to 16 while the vision encoder keeps 8, and pass a matching 16-NPU device list when loading it.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from optimum.rbln import RBLNAutoModelForImageTextToText

# Define the HuggingFace model ID
model_id = "nvidia/Cosmos-Reason2-2B"

# Compile the model for 8 RBLN NPUs
model = RBLNAutoModelForImageTextToText.from_pretrained(
    model_id,
    export=True,
    attn_implementation="sdpa",
    rbln_config={
        "visual": {
            "max_seq_len": 16_384,
            "num_devices": 8,
            "create_runtimes": False,
        },
        "num_devices": 8,
        "kvcache_partition_len": 16_384,
        "max_seq_len": 262_144,
        "create_runtimes": False,
    },
)
# Save all compiled results to disk
model.save_pretrained("rbln-cosmos-reason2")

Load the Compiled RBLN Model

Load the compiled RBLN model using RBLNAutoModelForImageTextToText.from_pretrained(). Pass the saved directory path as model_id and set the export argument to False.

Cosmos-Reason2 consists of a vision encoder and a language model, and loading both on a single NPU can exceed its memory capacity. Therefore, assign the NPUs each submodule runs on with the device parameter of rbln_config. In the example code below, the vision encoder and the language model share the same 8 NPUs.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
from optimum.rbln import RBLNAutoModelForImageTextToText

# Load the compiled RBLN model from the specified directory
model = RBLNAutoModelForImageTextToText.from_pretrained(
    model_id="rbln-cosmos-reason2",
    export=False,
    rbln_config={
        "visual": {
            "device": [0, 1, 2, 3, 4, 5, 6, 7],
        },
        "device": [0, 1, 2, 3, 4, 5, 6, 7],
    },
)

Prepare Inputs

Use the AutoProcessor from the transformers library to build the model inputs from a chat message list. The apply_chat_template() method decodes the video referenced in the message, samples its frames, and tokenizes the whole prompt in one call.

Note

Sample the input video at 4 FPS to match the setup Cosmos-Reason2 was trained with. The answer format appended to the text prompt is what elicits the model's chain-of-thought trace, wrapped in a <think> block.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained("rbln-cosmos-reason2")

video_messages = [
    {
        "role": "system",
        "content": [{"type": "text", "text": "You are a helpful assistant."}],
    },
    {
        "role": "user",
        "content": [
            {
                "type": "video",
                "url": "https://raw.githubusercontent.com/nvidia-cosmos/cosmos-reason1/main/assets/sample.mp4",
                "fps": 4,
            },
            {
                "type": "text",
                "text": (
                    "Is it safe to turn right? Answer the question using the following format:\n\n<think>\nYour reasoning.\n</think>\n\nWrite your final answer immediately after the </think> tag."
                ),
            },
        ],
    },
]

# Process inputs
inputs = processor.apply_chat_template(
    video_messages,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
    return_tensors="pt",
    fps=4,
)

Model Inference

You can now generate a response using the generate() method. Trim the prompt tokens from the generated sequence, then decode the remaining tokens to read the answer.

Note

Keep max_new_tokens at 4096 or above. A chain-of-thought answer is long, and a smaller budget truncates the response before the final answer.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
# autoregressively complete prompt
generated_ids = model.generate(**inputs, max_new_tokens=4096)
generated_ids_trimmed = [
    out_ids[len(in_ids) :]
    for in_ids, out_ids in zip(inputs.input_ids, generated_ids, strict=False)
]

output_text = processor.batch_decode(
    generated_ids_trimmed,
    skip_special_tokens=True,
    clean_up_tokenization_spaces=False,
)
# Show text and result
print(f"Result: {output_text}")

Example Output: 

1
2
3
4
5
6
7
8
<think>
Okay, let's try to figure out if it's safe to turn right based on the video. First, the video shows a residential street with parked cars on both sides. The ego vehicle is approaching an intersection where there are parked cars on the left and right.
...
Since the ego is waiting at a stop sign, they need to check for any oncoming traffic, pedestrians, or cyclists before turning. The video doesn't show any traffic coming from the opposite direction, but in a residential area, there could be cars backing out of driveways or pedestrians. The presence of a stop sign means they must yield to any cross traffic or pedestrians before turning right.
</think>

...
Once the crosswalk is confirmed clear, the turn would be safe.

The full example code is available in the RBLN Model Zoo GitHub repository.

References