Fine-Grained Device Assignment¶

Overview¶

The transformers (optimum) model implementation supports assigning different sets of RBLN NPU devices to model components called submodules, such as a vision encoder or a language model. Using disjoint device sets distributes submodule memory across separate NPUs.

For tensor-level device assignment through torch.device, use the vllm model implementation.

Configure Device Assignment¶

Assign device IDs through additional_config["optimum_overrides"] in Python or --rbln-optimum-overrides on the CLI. Configure the device count and device IDs separately.

Submodule names depend on the model architecture. Use the configuration keys defined by your model's RBLN Optimum API. In the table below, <submodule> stands for one of these keys.

The Python column shows keys within additional_config.

On the CLI, use --rbln-num-devices-per-local-rank for the top-level device count and --rbln-optimum-overrides for the remaining settings, passed as JSON.

• Submodules without an explicit device setting use the top-level device setting.
• To assign different devices to a specific submodule, set device under that submodule's configuration key.
• When compiling a model, the submodule's num_devices is automatically set to the number of IDs in its explicit device list, so it does not need to be specified separately.

For a precompiled model, pass the artifact directory as model. Only device entries in optimum_overrides apply, and each device list must match the corresponding submodule's compiled device count.

For model loading or device memory issues, see Troubleshooting Multi-Module Models.

Example¶

Qwen3-VL consists of a vision encoder and a language model. The vision encoder has the configuration key visual. In this model, the top-level device applies to the language model, and visual.device applies to the vision encoder.

The following examples use 16 ATOM™ devices, with the language model on devices 0–7 and the vision encoder on devices 8–15. In this example, num_devices_per_local_rank=8 sets the language model's device count, while visual.num_devices=8 sets the vision encoder's device count. visual.num_devices is specified explicitly here to make the device count clear. Run rbln-smi to check the devices available on your system.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

from vllm import LLM

llm = LLM(
    model="Qwen/Qwen3-VL-2B-Instruct",
    model_impl="transformers",
    block_size=16_384,
    max_model_len=262_144,
    max_num_seqs=8,
    additional_config={
        "num_devices_per_local_rank": 8,
        "optimum_overrides": {
            "device": [0, 1, 2, 3, 4, 5, 6, 7],
            "visual": {
                "device": [8, 9, 10, 11, 12, 13, 14, 15],
                "num_devices": 8,
                "max_seq_len": 16_384,
            },
        },
    },
)

1
2
3
4
5
6
7

vllm serve Qwen/Qwen3-VL-2B-Instruct \
  --model-impl transformers \
  --block-size 16384 \
  --max-model-len 262144 \
  --max-num-seqs 8 \
  --rbln-num-devices-per-local-rank 8 \
  --rbln-optimum-overrides '{"device": [0, 1, 2, 3, 4, 5, 6, 7], "visual": {"device": [8, 9, 10, 11, 12, 13, 14, 15], "num_devices": 8, "max_seq_len": 16384}}'

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

from vllm import LLM

llm = LLM(
    model="rbln-Qwen3-VL-2B-Instruct",
    model_impl="transformers",
    additional_config={
        "optimum_overrides": {
            "device": [0, 1, 2, 3, 4, 5, 6, 7],
            "visual": {
                "device": [8, 9, 10, 11, 12, 13, 14, 15],
            },
        },
    },
)