diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 1724a96263..b3a1184873 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -15,8 +15,11 @@ import logging import re import socket +import threading from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager +from contextvars import ContextVar +from importlib import import_module from typing import Any, Literal, Optional import torch @@ -43,6 +46,10 @@ ) logger = logging.getLogger(__name__) +_BF16_TRTLLM_LAYOUT_PATCH_LOCK = threading.RLock() +_BF16_TRTLLM_LAYOUT_PATCH_ACTIVE: ContextVar[bool] = ContextVar( + "bf16_trtllm_layout_patch_active", default=False +) try: import vllm # noqa: F401 @@ -173,6 +180,109 @@ def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: return bool(_unquantized_flashinfer_trtllm_modules(model)) +def _convert_bf16_moe_weights_to_trtllm_block_layout_batched( + cache_permute_indices: dict[torch.Size, torch.Tensor], + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + is_gated_act_gemm: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert all BF16 experts to TRTLLM block layout in two gathers.""" + if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16: + raise ValueError( + "Unquantized MoE backend FlashInfer TRTLLM requires bfloat16 weights" + ) + if w13_weight.ndim != 3 or w2_weight.ndim != 3: + raise ValueError( + "TRTLLM BF16 MoE weights must have shape [experts, rows, cols]" + ) + if w13_weight.shape[0] != w2_weight.shape[0]: + raise ValueError("W13 and W2 must contain the same number of experts") + + flashinfer_moe_core = import_module("flashinfer.fused_moe.core") + + epilogue_tile_m = 128 + block_k = 128 + w13_expert_uint8 = w13_weight[0].view(torch.uint8) + w2_expert_uint8 = w2_weight[0].view(torch.uint8) + w13_permute_indices = flashinfer_moe_core._maybe_get_cached_w3_w1_permute_indices( + cache_permute_indices, + w13_expert_uint8, + epilogue_tile_m, + is_gated_act_gemm=is_gated_act_gemm, + ) + if is_gated_act_gemm: + rows = w13_expert_uint8.shape[0] + w13_permute_indices = (w13_permute_indices + rows // 2) % rows + w2_permute_indices = flashinfer_moe_core.get_w2_permute_indices_with_cache( + cache_permute_indices, + w2_expert_uint8, + epilogue_tile_m, + ) + + def _convert(weight: torch.Tensor, source_indices: torch.Tensor) -> torch.Tensor: + weight_uint8 = weight.view(torch.uint8) + num_experts, rows, byte_cols = weight_uint8.shape + if byte_cols % block_k != 0: + raise ValueError( + f"TRTLLM BF16 MoE byte columns must be divisible by {block_k}; " + f"got {byte_cols}" + ) + expert_blocks = weight_uint8.view( + num_experts, rows, byte_cols // block_k, block_k + ).permute(0, 2, 1, 3) + return ( + torch.index_select( + expert_blocks, + 2, + source_indices.to(weight.device), + ) + .contiguous() + .view(torch.bfloat16) + ) + + return ( + _convert(w13_weight, w13_permute_indices), + _convert(w2_weight, w2_permute_indices), + ) + + +@contextmanager +def _use_batched_bf16_trtllm_layout_conversion() -> Iterator[None]: + """Use the batched converter only while vLLM rebuilds TRTLLM MoE state.""" + from vllm.model_executor.layers.fused_moe.oracle import unquantized + + with _BF16_TRTLLM_LAYOUT_PATCH_LOCK: + original_converter = ( + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + ) + + def _dispatch( + cache_permute_indices: dict[torch.Size, torch.Tensor], + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + is_gated_act_gemm: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + converter = ( + _convert_bf16_moe_weights_to_trtllm_block_layout_batched + if _BF16_TRTLLM_LAYOUT_PATCH_ACTIVE.get() + else original_converter + ) + return converter( + cache_permute_indices, + w13_weight, + w2_weight, + is_gated_act_gemm=is_gated_act_gemm, + ) + + active_token = _BF16_TRTLLM_LAYOUT_PATCH_ACTIVE.set(True) + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout = _dispatch + try: + yield + finally: + _BF16_TRTLLM_LAYOUT_PATCH_ACTIVE.reset(active_token) + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout = ( + original_converter + ) def _local_shard_slices(param_info: dict[str, Any], rank: int) -> tuple[slice, ...]: """Return this destination rank's slices in an HF-global tensor.""" from nemo_rl.weight_sync.xferdtensor_python import _compute_shard_slices @@ -1075,22 +1185,24 @@ def _weight_update_lifecycle( reloaded_module_ids = _reload_target_module_ids(reload_targets) def finalize() -> None: - with torch.device(self.device): - finalize_layerwise_reload(model, self.model_config) - _process_mxfp8_modules_after_native_reload( - model, reloaded_module_ids - ) - _refresh_hpc_modules_after_layerwise_reload(model) - self._maybe_process_mtp_drafter_after_loading() + with _use_batched_bf16_trtllm_layout_conversion(): + with torch.device(self.device): + finalize_layerwise_reload(model, self.model_config) + _process_mxfp8_modules_after_native_reload( + model, reloaded_module_ids + ) + _refresh_hpc_modules_after_layerwise_reload(model) + self._maybe_process_mtp_drafter_after_loading() torch.cuda.synchronize() try: with set_current_vllm_config(self.model_runner.vllm_config): - with torch.device(self.device): - for reload_target in reload_targets: - initialize_layerwise_reload(reload_target) - self._nrl_layerwise_reload_active = True - yield finalize + with _use_batched_bf16_trtllm_layout_conversion(): + with torch.device(self.device): + for reload_target in reload_targets: + initialize_layerwise_reload(reload_target) + self._nrl_layerwise_reload_active = True + yield finalize except Exception as error: self._nrl_layerwise_reload_failure = error raise diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 2961764a77..3289c7cdfa 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -249,6 +249,133 @@ def __init__(self): hpc_module.process_weights_after_loading.assert_called_once_with(model) +@pytest.mark.vllm +@pytest.mark.parametrize("is_gated_act_gemm", [False, True]) +def test_batched_bf16_trtllm_layout_matches_vllm_expertwise_converter( + monkeypatch, is_gated_act_gemm +): + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + convert_moe_weights_to_flashinfer_trtllm_block_layout, + ) + + from nemo_rl.models.generation.vllm import vllm_backend + + num_experts = 3 + w13_rows = 4 + w2_rows = 3 + cols = 128 + w13 = torch.arange(num_experts * w13_rows * cols, dtype=torch.bfloat16).view( + num_experts, w13_rows, cols + ) + w2 = torch.arange(num_experts * w2_rows * cols, dtype=torch.bfloat16).view( + num_experts, w2_rows, cols + ) + w13_perm = torch.tensor([2, 0, 3, 1]) + w2_perm = torch.tensor([1, 2, 0]) + + def get_w13_perm(cache, weight, tile_m, *, is_gated_act_gemm): + return w13_perm + + def get_w2_perm(cache, weight, tile_m): + return w2_perm + + monkeypatch.setattr( + "flashinfer.fused_moe.core._maybe_get_cached_w3_w1_permute_indices", + get_w13_perm, + ) + monkeypatch.setattr( + "flashinfer.fused_moe.core.get_w2_permute_indices_with_cache", + get_w2_perm, + ) + + expected_w13, expected_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout( + {}, w13, w2, is_gated_act_gemm=is_gated_act_gemm + ) + actual_w13, actual_w2 = ( + vllm_backend._convert_bf16_moe_weights_to_trtllm_block_layout_batched( + {}, w13, w2, is_gated_act_gemm=is_gated_act_gemm + ) + ) + + torch.testing.assert_close(actual_w13, expected_w13) + torch.testing.assert_close(actual_w2, expected_w2) + + +@pytest.mark.vllm +def test_batched_bf16_trtllm_layout_is_scoped_to_reload_finalize(monkeypatch): + import threading + + from vllm.model_executor.layers.fused_moe.oracle import unquantized + + from nemo_rl.models.generation.vllm import vllm_backend + + original_result = (object(), object()) + batched_result = (object(), object()) + original_converter = MagicMock(return_value=original_result) + batched_converter = MagicMock(return_value=batched_result) + monkeypatch.setattr( + vllm_backend, + "_convert_bf16_moe_weights_to_trtllm_block_layout_batched", + batched_converter, + ) + monkeypatch.setattr( + unquantized, + "convert_moe_weights_to_flashinfer_trtllm_block_layout", + original_converter, + ) + + model = _make_unquantized_moe_model("FlashInfer TRTLLM") + vllm_config = SimpleNamespace(quant_config=None) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=vllm_config) + ext.model_config = object() + ext.device = torch.device("cpu") + ext._maybe_process_mtp_drafter_after_loading = MagicMock() + + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda _: None, + ) + + def finalize_layerwise_reload(_model, _model_config): + assert unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", + finalize_layerwise_reload, + ) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + with ext._weight_update_lifecycle("collective") as finalize: + active_converter = ( + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + ) + assert active_converter({}, object(), object()) == batched_result + + thread_results = [] + thread = threading.Thread( + target=lambda: thread_results.append( + active_converter({}, object(), object()) + ) + ) + thread.start() + thread.join() + assert thread_results == [original_result] + finalize() + + assert ( + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + is original_converter + ) + batched_converter.assert_called_once() + original_converter.assert_called_once() + + class _DeferredReloadLayer(torch.nn.Module): def __init__(self) -> None: super().__init__()