gpt-oss-20b (Data Parallelism)¶

Installation¶

Install vllm-rbln as described in RBLN vLLM. The example scripts on this page live under examples/ in the vllm-rbln repository, so following the source-code install leaves them in place as well.

Activate RBLN Profiler¶

Setting RBLN_PROFILER=1 enables the RBLN Profiler.

1

export RBLN_PROFILER=1

No additional profiler-specific settings are required for Rebel100™.

Extract Profiled Data with Inference¶

See data_parallel_offline.py in vllm-rbln for an example of data-parallel execution. To profile a model, keep its execution settings and pass profiler_config when creating the LLM. After warm-up, use start_profile() and stop_profile() to mark the inference range.

The RBLN Profiler runs on the vllm model implementation, and model_impl="vllm" on LLM selects it.

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

llm = LLM(
    model=model,
    model_impl="vllm",
    tensor_parallel_size=tp_size,
    enable_expert_parallel=enable_ep,
    block_size=block_size,
    max_num_batched_tokens=128,
    max_num_seqs=1,
    profiler_config={
        "profiler": "torch",
        "torch_profiler_dir": "./profile",
    },
)
llm.generate(prompts, SamplingParams(temperature=0.0, max_tokens=2))
llm.start_profile()
outputs = llm.generate(prompts, sampling_params)
llm.stop_profile()

--dp-size sets the data-parallel size and --enable-expert-parallel turns on Expert Parallelism. The example below brings up two ranks.

1
2
3
4
5
6

$ cd vllm-rbln
$ RBLN_PROFILER=1 \
    python3 examples/vllm/features/data_parallel/data_parallel_offline.py \
    --model openai/gpt-oss-20b \
    --dp-size 2 \
    --enable-expert-parallel

The run produces as many artifacts as there are ranks. The RBLN Profiler's traces go under a per-rank directory, and the PyTorch profiler's chrometrace goes to the directory named by torch_profiler_dir, one per rank.

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

# RBLN Profiler — one directory per rank
$ ls rank_0/session_0
rbln_{date}_{time}.pb            # the whole trace; this is the file you open in Perfetto
rbln_{date}_{time}_0.pb          # one per execution
rbln_{date}_{time}_1.pb
...
$ ls rank_1/session_0
rbln_{date}_{time}.pb
rbln_{date}_{time}_0.pb
...

# PyTorch profiler — the per-rank chrometrace files share one directory
$ ls profile
dp0_pp0_tp0_dcp0_ep0_rank0.{timestamp}.pt.trace.json.gz
dp1_pp0_tp0_dcp0_ep1_rank0.{timestamp}.pt.trace.json.gz

Collective Communication in a MoE Model¶

When a model runs across several devices, the ranks exchange data through collective communication (CCL). In a MoE model run with Expert Parallelism, ccl_all_gather_into_tensor and ccl_reduce_scatter_tensor form a pair that repeats layer by layer through the model, and the screenshot zooms in on one such layer.

One layer appears as two Device Workload slices and two collectives. In the screenshot, i1_0_Device Workload_24 → i1_0_ccl_all_gather_into_tensor_1 → i1_0_Device Workload_25 → i1_0_ccl_reduce_scatter_tensor_1 is one layer, and the next one begins at i1_0_Device Workload_26 and i1_0_ccl_all_gather_into_tensor_2. ccl_all_gather_into_tensor collects the tokens, the Device Workload that follows runs the routing and the experts each rank owns, and ccl_reduce_scatter_tensor splits the results back out.

The first Device Workload holds the attention commands, such as input_layernorm and attn_qkv_proj_weight, and the second the routing and expert commands, such as mlp_router_weight and custom_moe_glu_mxfp4. The screenshot marks off those four stages with every chiplet expanded, showing that both the attention and the MoE run split across the four chiplets.

Collective communication appears in the trace at two levels. The CCL Workload track, which sits under Host next to Device Workload, holds one slice per collective covering its whole execution window, and the CCL Command track inside the expanded CCL group holds a slice of the same name covering only the portion in which the command stream runs.

A collective is one operation shared by every rank, so the same slice name appears in each rank's trace. Because the slice covers a rank's whole window for that collective, comparing the same name across the traces shows how long each rank spent waiting for the others.

Per-Rank Trace Merging¶

Opening the per-rank traces one at a time makes the relationship between ranks hard to see. The Perfetto UI can merge several chrometrace files onto one timeline. A chrometrace timestamps against the host clock and carries the same reference time, baseTimeNanoseconds, in every file, so files from different ranks land on one time axis.

In the merged view, each rank's tracks stay apart while they share the time axis. In the screenshot, the collectives of the same name finish in the same time window on both ranks.

The procedure, including enabling the plugin, is described in Multiple Traces on One Timeline.