From dd129c70332adac323c52d639e1484b34e9faf10 Mon Sep 17 00:00:00 2001 From: Zhengyang Liu Date: Fri, 6 Nov 2026 11:49:28 +0800 Subject: [PATCH] fix(verl): preserve routed expert IDs above 255 --- agentlightning/verl/rollout_adapter.py | 14 +++++++++-- tests/verl/test_rollout_adapter.py | 32 ++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/agentlightning/verl/rollout_adapter.py b/agentlightning/verl/rollout_adapter.py index 89b23d03c..9f27e81d4 100644 --- a/agentlightning/verl/rollout_adapter.py +++ b/agentlightning/verl/rollout_adapter.py @@ -243,15 +243,25 @@ def _build_routed_experts_batch( batch: torch.Tensor | None = None for index, (payload, original_prompt_length, prompt_length, response_length) in enumerate(rows): routes = np.load(io.BytesIO(base64.b64decode(payload)), allow_pickle=False) + if not np.issubdtype(routes.dtype, np.integer) or routes.min(initial=0) < 0: + raise ValueError("R3 routed_experts must contain non-negative integer expert IDs") + max_expert_id = int(routes.max(initial=0)) + if max_expert_id > torch.iinfo(torch.int32).max: + raise ValueError("R3 routed_experts expert IDs exceed the int32 range") + # Keep the compact representation for <=256 experts without wrapping + # larger IDs (e.g. 256 -> 0) and changing the routes replayed in training. + dtype = torch.uint8 if max_expert_id <= 255 else torch.int32 if batch is None: batch = torch.zeros( (len(rows), max_prompt_length + max_response_length, *routes.shape[1:]), - dtype=torch.uint8, + dtype=dtype, device=device, ) + elif dtype == torch.int32 and batch.dtype == torch.uint8: + batch = batch.to(dtype=torch.int32) if len(routes) < original_prompt_length + response_length - 1: raise RuntimeError("R3 routed_experts is shorter than its token sequence") - routes = torch.from_numpy(routes).to(device=device, dtype=torch.uint8) + routes = torch.from_numpy(routes).to(device=device, dtype=batch.dtype) prompt_routes = min(prompt_length, len(routes)) prompt_start = max_prompt_length - prompt_length batch[index, prompt_start : prompt_start + prompt_routes] = routes[:prompt_routes] diff --git a/tests/verl/test_rollout_adapter.py b/tests/verl/test_rollout_adapter.py index abd914615..65e06d6b7 100644 --- a/tests/verl/test_rollout_adapter.py +++ b/tests/verl/test_rollout_adapter.py @@ -60,6 +60,38 @@ def test_routed_experts_batch_preserves_padding_and_truncation(dtype: type[np.ge torch.testing.assert_close(batch, expected, rtol=0, atol=0) +@pytest.mark.parametrize("dtype", [np.uint16, np.int32, np.int64]) +@pytest.mark.parametrize("high_routes_first", [False, True]) +def test_routed_experts_batch_preserves_expert_ids_above_255(dtype: type[np.generic], high_routes_first: bool) -> None: + low_routes = np.array([[[0, 255]], [[128, 1]]], dtype=dtype) + high_routes = np.array([[[256, 511]], [[300, 257]], [[510, 400]], [[1, 2]]], dtype=dtype) + low_row = (_encode_routes(low_routes), 1, 1, 2) + high_row = (_encode_routes(high_routes), 2, 2, 2) + rows = [high_row, low_row] if high_routes_first else [low_row, high_row] + + batch = _build_routed_experts_batch(rows, 3, 3, torch.device("cpu")) + + low_expected = [[[0, 0]], [[0, 0]], [[0, 255]], [[128, 1]], [[0, 0]], [[0, 0]]] + high_expected = [[[0, 0]], [[256, 511]], [[300, 257]], [[510, 400]], [[1, 2]], [[0, 0]]] + expected_rows = [high_expected, low_expected] if high_routes_first else [low_expected, high_expected] + expected = torch.tensor(expected_rows, dtype=torch.int32) + torch.testing.assert_close(batch.to(torch.int32), expected, rtol=0, atol=0) + assert batch.dtype == torch.int32 + + +@pytest.mark.parametrize( + "routes", + [ + np.array([[[-1]]], dtype=np.int32), + np.array([[[0.5]]], dtype=np.float32), + np.array([[[2**31]]], dtype=np.int64), + ], +) +def test_routed_experts_batch_rejects_ids_that_cannot_be_represented(routes: np.ndarray) -> None: + with pytest.raises(ValueError, match="R3 routed_experts"): + _build_routed_experts_batch([(_encode_routes(routes), 1, 1, 1)], 1, 1, torch.device("cpu")) + + def test_routed_experts_batch_releases_decoded_rows(monkeypatch: pytest.MonkeyPatch) -> None: payload = _encode_routes(np.arange(48, dtype=np.uint8).reshape(12, 2, 2)) decoded: list[weakref.ReferenceType[np.ndarray]] = []