From 7db8244c7290c35b59b7607acee15b01d0c1d739 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 22:22:39 -0700 Subject: [PATCH 01/68] test(vllm): cover native BF16 MoE refit lifecycle Signed-off-by: seonjinn (cherry picked from commit 707968a004c2917debaf15bf0d7dcaef711f038f) --- .../models/generation/test_vllm_backend.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index c7e7969ddc9..a0f5e102b4b 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -119,6 +119,69 @@ def _make_mtp_refit_extension( return ext, drafter_model +@pytest.mark.vllm +def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + call_order = [] + model = object() + model_config = object() + vllm_config = object() + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=vllm_config) + ext.model_config = model_config + ext.device = torch.device("cpu") + ext._maybe_process_mtp_drafter_after_loading = lambda: call_order.append("mtp") + ext._maybe_process_fp8_kv_cache = lambda: call_order.append("kv") + + monkeypatch.setattr(fp8, "is_fp8_model", lambda config: False) + + @contextlib.contextmanager + def set_current_vllm_config(config): + assert config is vllm_config + call_order.append("config_enter") + try: + yield + finally: + call_order.append("config_exit") + + monkeypatch.setattr("vllm.config.set_current_vllm_config", set_current_vllm_config) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda reload_model: call_order.append(("initialize", reload_model)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", + lambda reload_model, config: call_order.append( + ("finalize", reload_model, config) + ), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.process_weights_after_loading", + lambda *_args: pytest.fail( + "unquantized refit must use vLLM's native layerwise reload lifecycle" + ), + ) + + with ext._weight_update_lifecycle("collective") as finalize: + call_order.append("load") + finalize() + + assert call_order == [ + "config_enter", + ("initialize", model), + "load", + ("finalize", model, model_config), + "mtp", + "config_exit", + "kv", + ] + + @pytest.mark.vllm @pytest.mark.parametrize("with_mtp", [False, True]) def test_update_weights_from_collective_processes_weights_after_loading( From e919d263180fc7369f010c8a8a43cb33ab431d19 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 22:49:46 -0700 Subject: [PATCH 02/68] fix(vllm): reload unquantized TRTLLM MoE weights Signed-off-by: seonjinn (cherry picked from commit cbd413867316e46e18c261066c6edeaa6b619551) --- .../models/generation/vllm/vllm_backend.py | 73 +++++++++++++++++- .../models/generation/test_vllm_backend.py | 77 +++++++++++++++++-- 2 files changed, 143 insertions(+), 7 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 8c28d3c19b2..7eef2de6d23 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -71,6 +71,25 @@ class IPCWeightManifestError(RuntimeError): """An IPC transfer did not match the prepared state-dict manifest.""" +def _detach_pending_layerwise_weights( + model: torch.nn.Module, source_storage_ptrs: set[int] +) -> None: + """Detach deferred reload weights from a reusable transport buffer.""" + if not source_storage_ptrs: + return + + from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info + + for module in model.modules(): + info = get_layerwise_info(module) + for _, arguments in info.loaded_weights: + loaded_weight = arguments.arguments.get("loaded_weight") + if not isinstance(loaded_weight, torch.Tensor): + continue + if loaded_weight.untyped_storage().data_ptr() in source_storage_ptrs: + arguments.arguments["loaded_weight"] = loaded_weight.clone() + + class _IPCWeightManifest: """Validate an IPC stream against its prepared state-dict manifest.""" @@ -184,6 +203,7 @@ class VllmInternalWorkerExtension: _mtp_drafter_from_disk: bool = False _sparse_delta_applier: Any = None _nrl_named_parameters: dict[str, torch.nn.Parameter] + _nrl_layerwise_reload_active: bool = False def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: params = getattr(self, "_nrl_named_parameters", None) @@ -195,7 +215,19 @@ def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: def _load_full_hf_weights( self, policy_weights: list[tuple[str, torch.Tensor]] ) -> None: - self.model_runner.model.load_weights(weights=policy_weights) + if not self._nrl_layerwise_reload_active: + self.model_runner.model.load_weights(weights=policy_weights) + return + + source_storage_ptrs = { + tensor.untyped_storage().data_ptr() for _, tensor in policy_weights + } + try: + self.model_runner.model.load_weights(weights=policy_weights) + finally: + _detach_pending_layerwise_weights( + self.model_runner.model, source_storage_ptrs + ) def _load_hf_weights(self, policy_weights: list[tuple[str, torch.Tensor]]) -> None: from nemo_rl.models.generation.vllm.quantization import fp8 @@ -605,12 +637,49 @@ def _get_sparse_delta_applier(self) -> Any: ) return self._sparse_delta_applier + def _uses_unquantized_flashinfer_trtllm(self) -> bool: + vllm_config = self.model_runner.vllm_config + kernel_config = getattr(vllm_config, "kernel_config", None) + if getattr(kernel_config, "moe_backend", None) != "flashinfer_trtllm": + return False + + from nemo_rl.models.generation.vllm.quantization import fp8 + + return not fp8.is_fp8_model(vllm_config) + @contextmanager def _weight_update_lifecycle( self, transport: WeightUpdateTransport ) -> Iterator[WeightUpdateFinalizer]: """Provide setup/finalization around a transport-owned weight update.""" del transport + if self._uses_unquantized_flashinfer_trtllm(): + from vllm.config import set_current_vllm_config + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + ) + + model = self.model_runner.model + + def finalize() -> None: + with torch.device(self.device): + finalize_layerwise_reload(model, self.model_config) + self._maybe_process_mtp_drafter_after_loading() + torch.accelerator.synchronize() + + try: + with set_current_vllm_config(self.model_runner.vllm_config): + with torch.device(self.device): + initialize_layerwise_reload(model) + self._nrl_layerwise_reload_active = True + yield finalize + finally: + self._nrl_layerwise_reload_active = False + + self._maybe_process_fp8_kv_cache() + return + from vllm.config import set_current_vllm_config from vllm.model_executor.model_loader.utils import ( process_weights_after_loading, @@ -630,7 +699,7 @@ def finalize() -> None: def _weight_update_errors_are_fatal(self) -> bool: """Whether transport errors should propagate instead of returning False.""" - return False + return self._uses_unquantized_flashinfer_trtllm() def _synchronize_before_ipc_data_ack(self) -> None: """Fence work consuming one IPC data batch before its acknowledgment.""" diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index a0f5e102b4b..3fc233f7a33 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -127,7 +127,9 @@ def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): call_order = [] model = object() model_config = object() - vllm_config = object() + vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + ) ext = vllm_backend.VllmInternalWorkerExtension.__new__( vllm_backend.VllmInternalWorkerExtension @@ -139,6 +141,7 @@ def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): ext._maybe_process_fp8_kv_cache = lambda: call_order.append("kv") monkeypatch.setattr(fp8, "is_fp8_model", lambda config: False) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) @contextlib.contextmanager def set_current_vllm_config(config): @@ -167,11 +170,13 @@ def set_current_vllm_config(config): ), ) - with ext._weight_update_lifecycle("collective") as finalize: - call_order.append("load") - finalize() + for _ in range(2): + with ext._weight_update_lifecycle("collective") as finalize: + call_order.append("load") + finalize() + assert ext._nrl_layerwise_reload_active is False - assert call_order == [ + expected_cycle = [ "config_enter", ("initialize", model), "load", @@ -180,6 +185,68 @@ def set_current_vllm_config(config): "config_exit", "kv", ] + assert call_order == expected_cycle * 2 + + +@pytest.mark.vllm +def test_layerwise_reload_detaches_deferred_transport_weights(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + model = SimpleNamespace(load_weights=MagicMock()) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model) + ext._nrl_layerwise_reload_active = True + weight = torch.ones(2) + detach = MagicMock() + monkeypatch.setattr(vllm_backend, "_detach_pending_layerwise_weights", detach) + + ext._load_full_hf_weights([("model.weight", weight)]) + + model.load_weights.assert_called_once_with(weights=[("model.weight", weight)]) + detach.assert_called_once_with(model, {weight.untyped_storage().data_ptr()}) + + +@pytest.mark.vllm +def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + model = object() + model_config = object() + vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + ) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=vllm_config) + ext.model_config = model_config + ext.device = torch.device("cpu") + ext._maybe_process_mtp_drafter_after_loading = MagicMock() + ext._maybe_process_fp8_kv_cache = MagicMock() + + monkeypatch.setattr(fp8, "is_fp8_model", lambda config: True) + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + process = MagicMock() + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.process_weights_after_loading", + process, + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda _: pytest.fail("FP8 must not use the unquantized reload lifecycle"), + ) + + with ext._weight_update_lifecycle("collective") as finalize: + finalize() + + process.assert_called_once_with(model, model_config, ext.device) + ext._maybe_process_mtp_drafter_after_loading.assert_called_once_with() + ext._maybe_process_fp8_kv_cache.assert_called_once_with() @pytest.mark.vllm From 2ce743e0c8a1963d1f21efa8d033f2c1a8d43e80 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 22:57:15 -0700 Subject: [PATCH 03/68] fix(vllm): preserve refit error fallback Signed-off-by: seonjinn (cherry picked from commit bda88584c2e8ffbb54a40c1e807f1b410c095047) --- nemo_rl/models/generation/vllm/vllm_backend.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 7eef2de6d23..7b880cb284d 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -638,7 +638,10 @@ def _get_sparse_delta_applier(self) -> Any: return self._sparse_delta_applier def _uses_unquantized_flashinfer_trtllm(self) -> bool: - vllm_config = self.model_runner.vllm_config + model_runner = getattr(self, "model_runner", None) + vllm_config = getattr(model_runner, "vllm_config", None) + if vllm_config is None: + return False kernel_config = getattr(vllm_config, "kernel_config", None) if getattr(kernel_config, "moe_backend", None) != "flashinfer_trtllm": return False From b79256b9666e91f892076a9b29e00f9f9d9e1675 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 23:00:46 -0700 Subject: [PATCH 04/68] fix(vllm): default injected reload state safely Signed-off-by: seonjinn (cherry picked from commit 2aa1570ecce43297262a5fc4d4c49388c778169a) --- nemo_rl/models/generation/vllm/vllm_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 7b880cb284d..c0c69dcfd64 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -215,7 +215,7 @@ def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: def _load_full_hf_weights( self, policy_weights: list[tuple[str, torch.Tensor]] ) -> None: - if not self._nrl_layerwise_reload_active: + if not getattr(self, "_nrl_layerwise_reload_active", False): self.model_runner.model.load_weights(weights=policy_weights) return From e3d561d308db261da031e185c7cfe04740a9099f Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 23:02:46 -0700 Subject: [PATCH 05/68] fix(vllm): constrain native TRTLLM refit scope Signed-off-by: seonjinn (cherry picked from commit 4aece4b71405772ab4de0d82e63a9f1ef573e6db) --- .../models/generation/vllm/vllm_backend.py | 8 +++- .../models/generation/test_vllm_backend.py | 44 ++++++++++++++++++- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index c0c69dcfd64..58a0bf0e500 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -638,6 +638,8 @@ def _get_sparse_delta_applier(self) -> Any: return self._sparse_delta_applier def _uses_unquantized_flashinfer_trtllm(self) -> bool: + if callable(getattr(self, "_is_real_quant_model", None)): + return False model_runner = getattr(self, "model_runner", None) vllm_config = getattr(model_runner, "vllm_config", None) if vllm_config is None: @@ -657,6 +659,11 @@ def _weight_update_lifecycle( """Provide setup/finalization around a transport-owned weight update.""" del transport if self._uses_unquantized_flashinfer_trtllm(): + if self._mtp_drafter_refit_enabled(): + raise RuntimeError( + "Unquantized FlashInfer TRTLLM refit does not yet support " + "a co-trained MTP drafter" + ) from vllm.config import set_current_vllm_config from vllm.model_executor.model_loader.reload import ( finalize_layerwise_reload, @@ -680,7 +687,6 @@ def finalize() -> None: finally: self._nrl_layerwise_reload_active = False - self._maybe_process_fp8_kv_cache() return from vllm.config import set_current_vllm_config diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 3fc233f7a33..0ff2e1ae626 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -138,7 +138,7 @@ def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): ext.model_config = model_config ext.device = torch.device("cpu") ext._maybe_process_mtp_drafter_after_loading = lambda: call_order.append("mtp") - ext._maybe_process_fp8_kv_cache = lambda: call_order.append("kv") + ext._maybe_process_fp8_kv_cache = MagicMock() monkeypatch.setattr(fp8, "is_fp8_model", lambda config: False) monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) @@ -183,9 +183,9 @@ def set_current_vllm_config(config): ("finalize", model, model_config), "mtp", "config_exit", - "kv", ] assert call_order == expected_cycle * 2 + ext._maybe_process_fp8_kv_cache.assert_not_called() @pytest.mark.vllm @@ -249,6 +249,46 @@ def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): ext._maybe_process_fp8_kv_cache.assert_called_once_with() +@pytest.mark.vllm +def test_modelopt_extension_does_not_use_unquantized_reload(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + ) + ) + ext._is_real_quant_model = lambda: False + monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) + + assert ext._uses_unquantized_flashinfer_trtllm() is False + + +@pytest.mark.vllm +def test_unquantized_reload_rejects_cotrained_mtp(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + ) + ) + ext._mtp_drafter_refit_enabled = lambda: True + monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) + + with pytest.raises(RuntimeError, match="co-trained MTP drafter"): + with ext._weight_update_lifecycle("collective"): + pass + + @pytest.mark.vllm @pytest.mark.parametrize("with_mtp", [False, True]) def test_update_weights_from_collective_processes_weights_after_loading( From a81e440e0cc27ccb750c844b4c240ad389517377 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 23:25:46 -0700 Subject: [PATCH 06/68] test(vllm): harden native refit failure coverage Signed-off-by: seonjinn (cherry picked from commit 1292dbe9d17c650a71e840f65c7ae883ba02a034) --- .../models/generation/test_vllm_backend.py | 91 ++++++++++++++++--- 1 file changed, 77 insertions(+), 14 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 0ff2e1ae626..4ae2d6c4811 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -192,20 +192,27 @@ def set_current_vllm_config(config): def test_layerwise_reload_detaches_deferred_transport_weights(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend - model = SimpleNamespace(load_weights=MagicMock()) - ext = vllm_backend.VllmInternalWorkerExtension.__new__( - vllm_backend.VllmInternalWorkerExtension + source = torch.ones(4) + unrelated = torch.full((2,), 7.0) + source_args = SimpleNamespace(arguments={"loaded_weight": source[:2]}) + unrelated_args = SimpleNamespace(arguments={"loaded_weight": unrelated}) + model = SimpleNamespace(modules=lambda: [object()]) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.layerwise.get_layerwise_info", + lambda _module: SimpleNamespace( + loaded_weights=[("source", source_args), ("other", unrelated_args)] + ), ) - ext.model_runner = SimpleNamespace(model=model) - ext._nrl_layerwise_reload_active = True - weight = torch.ones(2) - detach = MagicMock() - monkeypatch.setattr(vllm_backend, "_detach_pending_layerwise_weights", detach) - ext._load_full_hf_weights([("model.weight", weight)]) + vllm_backend._detach_pending_layerwise_weights( + model, {source.untyped_storage().data_ptr()} + ) - model.load_weights.assert_called_once_with(weights=[("model.weight", weight)]) - detach.assert_called_once_with(model, {weight.untyped_storage().data_ptr()}) + detached = source_args.arguments["loaded_weight"] + assert detached.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + assert unrelated_args.arguments["loaded_weight"] is unrelated + source.zero_() + torch.testing.assert_close(detached, torch.ones(2)) @pytest.mark.vllm @@ -251,6 +258,25 @@ def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): @pytest.mark.vllm def test_modelopt_extension_does_not_use_unquantized_reload(monkeypatch): + from nemo_rl.modelopt.models.generation import vllm_quant_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + ext = vllm_quant_backend.VllmQuantInternalWorkerExtension.__new__( + vllm_quant_backend.VllmQuantInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + ) + ) + monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) + + assert ext._uses_unquantized_flashinfer_trtllm() is False + + +@pytest.mark.vllm +@pytest.mark.parametrize("moe_backend", ["auto", "triton", None]) +def test_other_moe_backends_keep_existing_refit_lifecycle(monkeypatch, moe_backend): from nemo_rl.models.generation.vllm import vllm_backend from nemo_rl.models.generation.vllm.quantization import fp8 @@ -259,17 +285,16 @@ def test_modelopt_extension_does_not_use_unquantized_reload(monkeypatch): ) ext.model_runner = SimpleNamespace( vllm_config=SimpleNamespace( - kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + kernel_config=SimpleNamespace(moe_backend=moe_backend) ) ) - ext._is_real_quant_model = lambda: False monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) assert ext._uses_unquantized_flashinfer_trtllm() is False @pytest.mark.vllm -def test_unquantized_reload_rejects_cotrained_mtp(monkeypatch): +def test_unquantized_reload_rejects_cotrained_mtp_during_prepare(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend from nemo_rl.models.generation.vllm.quantization import fp8 @@ -285,6 +310,44 @@ def test_unquantized_reload_rejects_cotrained_mtp(monkeypatch): monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) with pytest.raises(RuntimeError, match="co-trained MTP drafter"): + ext.prepare_refit_info({"model.weight": object()}) + + assert not hasattr(ext, "state_dict_info") + + +@pytest.mark.vllm +def test_failed_unquantized_reload_marks_worker_unusable(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=object(), + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + ), + ) + ext.model_config = object() + ext.device = torch.device("cpu") + ext._mtp_drafter_refit_enabled = lambda: False + monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda _: None, + ) + + failure = RuntimeError("load failed") + with pytest.raises(RuntimeError, match="load failed"): + with ext._weight_update_lifecycle("collective"): + raise failure + + assert ext._nrl_layerwise_reload_failure is failure + with pytest.raises(RuntimeError, match="worker is unusable"): with ext._weight_update_lifecycle("collective"): pass From d791753a3a258d27a800e4d3122092b40adc797b Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 23:42:38 -0700 Subject: [PATCH 07/68] test(vllm): isolate refit capability contract Signed-off-by: seonjinn (cherry picked from commit 56e27a8abc85990caf5a7c5a5e2e3ac3156e62e5) --- tests/unit/models/generation/test_vllm_backend.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 4ae2d6c4811..9bf794af385 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -257,18 +257,19 @@ def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): @pytest.mark.vllm -def test_modelopt_extension_does_not_use_unquantized_reload(monkeypatch): - from nemo_rl.modelopt.models.generation import vllm_quant_backend +def test_extension_capability_can_disable_unquantized_reload(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend from nemo_rl.models.generation.vllm.quantization import fp8 - ext = vllm_quant_backend.VllmQuantInternalWorkerExtension.__new__( - vllm_quant_backend.VllmQuantInternalWorkerExtension + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension ) ext.model_runner = SimpleNamespace( vllm_config=SimpleNamespace( kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") ) ) + ext._supports_unquantized_flashinfer_trtllm_refit = lambda: False monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) assert ext._uses_unquantized_flashinfer_trtllm() is False From 60094522c75c683116ed733720b75fa9dd922adc Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 1 Aug 2026 23:49:44 -0700 Subject: [PATCH 08/68] fix(vllm): harden native layerwise refit failures Signed-off-by: seonjinn (cherry picked from commit 0c33c309b4172d4ef15d8ed272f9710cd011f89c) --- .../models/generation/vllm_quant_backend.py | 3 ++ .../models/generation/vllm/vllm_backend.py | 28 +++++++++++++++---- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py index 20b53757540..56032d9ac66 100644 --- a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py +++ b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py @@ -437,6 +437,9 @@ class VllmQuantInternalWorkerExtension(VllmInternalWorkerExtension): _nrl_w13_num_shards_by_prefix: dict[str, int] _nrl_modelopt_reload_roots: tuple[torch.nn.Module, ...] | None = None + def _supports_unquantized_flashinfer_trtllm_refit(self) -> bool: + return False + def maybe_init_zmq(self) -> None: """Use a longer timeout only for ModelOpt real-quant refits.""" super().maybe_init_zmq() diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 58a0bf0e500..3d6c9b1d36c 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -352,6 +352,7 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: state_dict_info (dict): A dictionary containing the info for refit. e.g. {tensor_name: (shape, dtype)} """ + self._validate_weight_update_compatibility() self.state_dict_info = state_dict_info # pyrefly: ignore[implicitly-defined-attribute] This class does not define __init__ so assignments like this should be ignored def prepare_sparse_delta_refit_info( @@ -637,8 +638,11 @@ def _get_sparse_delta_applier(self) -> Any: ) return self._sparse_delta_applier + def _supports_unquantized_flashinfer_trtllm_refit(self) -> bool: + return True + def _uses_unquantized_flashinfer_trtllm(self) -> bool: - if callable(getattr(self, "_is_real_quant_model", None)): + if not self._supports_unquantized_flashinfer_trtllm_refit(): return False model_runner = getattr(self, "model_runner", None) vllm_config = getattr(model_runner, "vllm_config", None) @@ -652,6 +656,16 @@ def _uses_unquantized_flashinfer_trtllm(self) -> bool: return not fp8.is_fp8_model(vllm_config) + def _validate_weight_update_compatibility(self) -> None: + if ( + self._uses_unquantized_flashinfer_trtllm() + and self._mtp_drafter_refit_enabled() + ): + raise RuntimeError( + "Unquantized FlashInfer TRTLLM refit does not yet support " + "a co-trained MTP drafter" + ) + @contextmanager def _weight_update_lifecycle( self, transport: WeightUpdateTransport @@ -659,11 +673,12 @@ def _weight_update_lifecycle( """Provide setup/finalization around a transport-owned weight update.""" del transport if self._uses_unquantized_flashinfer_trtllm(): - if self._mtp_drafter_refit_enabled(): + self._validate_weight_update_compatibility() + previous_failure = getattr(self, "_nrl_layerwise_reload_failure", None) + if previous_failure is not None: raise RuntimeError( - "Unquantized FlashInfer TRTLLM refit does not yet support " - "a co-trained MTP drafter" - ) + "The vLLM worker is unusable after a failed native layerwise refit" + ) from previous_failure from vllm.config import set_current_vllm_config from vllm.model_executor.model_loader.reload import ( finalize_layerwise_reload, @@ -684,6 +699,9 @@ def finalize() -> None: initialize_layerwise_reload(model) self._nrl_layerwise_reload_active = True yield finalize + except Exception as error: + self._nrl_layerwise_reload_failure = error + raise finally: self._nrl_layerwise_reload_active = False From a73421e57fe5b88e019834ff6e4c4e2c0d448701 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 7 Aug 2026 19:21:22 -0700 Subject: [PATCH 09/68] fix(vllm): detect realized TRTLLM refit backend Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 85 ++++++++-- .../models/generation/test_vllm_backend.py | 155 +++++++++++++++--- 2 files changed, 195 insertions(+), 45 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 3d6c9b1d36c..7d09ca8e616 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -78,6 +78,7 @@ def _detach_pending_layerwise_weights( if not source_storage_ptrs: return + # Keep reload internals off the normal non-layerwise weight-loading path. from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info for module in model.modules(): @@ -90,6 +91,26 @@ def _detach_pending_layerwise_weights( arguments.arguments["loaded_weight"] = loaded_weight.clone() +def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: + """Return whether a model realized the unquantized TRTLLM MoE backend.""" + # Import backend types only when inspecting a constructed vLLM model. + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + ) + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, + ) + + return any( + isinstance( + quant_method := getattr(module, "quant_method", None), + UnquantizedFusedMoEMethod, + ) + and quant_method.unquantized_backend is UnquantizedMoeBackend.FLASHINFER_TRTLLM + for module in model.modules() + ) + + class _IPCWeightManifest: """Validate an IPC stream against its prepared state-dict manifest.""" @@ -204,6 +225,7 @@ class VllmInternalWorkerExtension: _sparse_delta_applier: Any = None _nrl_named_parameters: dict[str, torch.nn.Parameter] _nrl_layerwise_reload_active: bool = False + _nrl_layerwise_reload_failure: Exception | None = None def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: params = getattr(self, "_nrl_named_parameters", None) @@ -222,12 +244,23 @@ def _load_full_hf_weights( source_storage_ptrs = { tensor.untyped_storage().data_ptr() for _, tensor in policy_weights } + load_error: Exception | None = None try: self.model_runner.model.load_weights(weights=policy_weights) + except Exception as error: + load_error = error + raise finally: - _detach_pending_layerwise_weights( - self.model_runner.model, source_storage_ptrs - ) + try: + _detach_pending_layerwise_weights( + self.model_runner.model, source_storage_ptrs + ) + except Exception: + if load_error is None: + raise + logger.exception( + "Failed to detach deferred weights after a weight load failure" + ) def _load_hf_weights(self, policy_weights: list[tuple[str, torch.Tensor]]) -> None: from nemo_rl.models.generation.vllm.quantization import fp8 @@ -579,20 +612,38 @@ def load_mtp_weights_from_disk(self, model_path: str) -> bool: f"include MTP layer weights to run deepseek_mtp speculative decoding." ) - self._load_draft_weights(weights) + draft_model_config = ( + self.model_runner.vllm_config.speculative_config.draft_model_config + ) # The MTP block contains MoE experts whose weights need post-load # processing (e.g. grouped-GEMM layout), matching the main-model path. + # Keep vLLM reload internals off the normal draft-loading path. from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.utils import ( - process_weights_after_loading, - ) - draft_model_config = ( - self.model_runner.vllm_config.speculative_config.draft_model_config - ) - with set_current_vllm_config(self.model_runner.vllm_config): - process_weights_after_loading(draft_model, draft_model_config, self.device) + if self._supports_unquantized_flashinfer_trtllm_refit() and ( + _model_uses_unquantized_flashinfer_trtllm(draft_model) + ): + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + ) + + with set_current_vllm_config(self.model_runner.vllm_config): + with torch.device(self.device): + initialize_layerwise_reload(draft_model) + self._load_draft_weights(weights) + finalize_layerwise_reload(draft_model, draft_model_config) + else: + from vllm.model_executor.model_loader.utils import ( + process_weights_after_loading, + ) + + self._load_draft_weights(weights) + with set_current_vllm_config(self.model_runner.vllm_config): + process_weights_after_loading( + draft_model, draft_model_config, self.device + ) # Mark that the MTP drafter is served from a one-time disk load so refit # does not re-load or re-process these static weights. self._mtp_drafter_from_disk = True @@ -648,13 +699,10 @@ def _uses_unquantized_flashinfer_trtllm(self) -> bool: vllm_config = getattr(model_runner, "vllm_config", None) if vllm_config is None: return False - kernel_config = getattr(vllm_config, "kernel_config", None) - if getattr(kernel_config, "moe_backend", None) != "flashinfer_trtllm": + if getattr(vllm_config, "quant_config", None) is not None: return False - from nemo_rl.models.generation.vllm.quantization import fp8 - - return not fp8.is_fp8_model(vllm_config) + return _model_uses_unquantized_flashinfer_trtllm(self.model_runner.model) def _validate_weight_update_compatibility(self) -> None: if ( @@ -674,11 +722,12 @@ def _weight_update_lifecycle( del transport if self._uses_unquantized_flashinfer_trtllm(): self._validate_weight_update_compatibility() - previous_failure = getattr(self, "_nrl_layerwise_reload_failure", None) + previous_failure = self._nrl_layerwise_reload_failure if previous_failure is not None: raise RuntimeError( "The vLLM worker is unusable after a failed native layerwise refit" ) from previous_failure + # Load vLLM reload internals only for the native layerwise path. from vllm.config import set_current_vllm_config from vllm.model_executor.model_loader.reload import ( finalize_layerwise_reload, diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 9bf794af385..e88318434d7 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -119,16 +119,29 @@ def _make_mtp_refit_extension( return ext, drafter_model +def _make_unquantized_moe_model(moe_backend: str) -> SimpleNamespace: + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + ) + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, + ) + + quant_method = UnquantizedFusedMoEMethod.__new__(UnquantizedFusedMoEMethod) + quant_method.unquantized_backend = UnquantizedMoeBackend(moe_backend) + module = SimpleNamespace(quant_method=quant_method) + return SimpleNamespace(modules=lambda: [module]) + + @pytest.mark.vllm def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend - from nemo_rl.models.generation.vllm.quantization import fp8 call_order = [] - model = object() + model = _make_unquantized_moe_model("FlashInfer TRTLLM") model_config = object() vllm_config = SimpleNamespace( - kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + kernel_config=SimpleNamespace(moe_backend="auto"), quant_config=None ) ext = vllm_backend.VllmInternalWorkerExtension.__new__( @@ -140,7 +153,6 @@ def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): ext._maybe_process_mtp_drafter_after_loading = lambda: call_order.append("mtp") ext._maybe_process_fp8_kv_cache = MagicMock() - monkeypatch.setattr(fp8, "is_fp8_model", lambda config: False) monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) @contextlib.contextmanager @@ -215,15 +227,39 @@ def test_layerwise_reload_detaches_deferred_transport_weights(monkeypatch): torch.testing.assert_close(detached, torch.ones(2)) +@pytest.mark.vllm +def test_layerwise_reload_preserves_weight_load_error(monkeypatch, caplog): + from nemo_rl.models.generation.vllm import vllm_backend + + load_error = RuntimeError("load failed") + model = SimpleNamespace(load_weights=MagicMock(side_effect=load_error)) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model) + ext._nrl_layerwise_reload_active = True + monkeypatch.setattr( + vllm_backend, + "_detach_pending_layerwise_weights", + MagicMock(side_effect=RuntimeError("detach failed")), + ) + + with pytest.raises(RuntimeError, match="load failed") as exc_info: + ext._load_full_hf_weights([("model.weight", torch.ones(1))]) + + assert exc_info.value is load_error + assert "Failed to detach deferred weights" in caplog.text + + @pytest.mark.vllm def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend - from nemo_rl.models.generation.vllm.quantization import fp8 model = object() model_config = object() vllm_config = SimpleNamespace( - kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=object(), ) ext = vllm_backend.VllmInternalWorkerExtension.__new__( vllm_backend.VllmInternalWorkerExtension @@ -234,7 +270,6 @@ def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): ext._maybe_process_mtp_drafter_after_loading = MagicMock() ext._maybe_process_fp8_kv_cache = MagicMock() - monkeypatch.setattr(fp8, "is_fp8_model", lambda config: True) monkeypatch.setattr( "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() ) @@ -257,58 +292,77 @@ def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): @pytest.mark.vllm -def test_extension_capability_can_disable_unquantized_reload(monkeypatch): +def test_extension_capability_can_disable_unquantized_reload(): from nemo_rl.models.generation.vllm import vllm_backend - from nemo_rl.models.generation.vllm.quantization import fp8 ext = vllm_backend.VllmInternalWorkerExtension.__new__( vllm_backend.VllmInternalWorkerExtension ) ext.model_runner = SimpleNamespace( + model=object(), vllm_config=SimpleNamespace( - kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") - ) + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, + ), ) ext._supports_unquantized_flashinfer_trtllm_refit = lambda: False - monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) assert ext._uses_unquantized_flashinfer_trtllm() is False @pytest.mark.vllm -@pytest.mark.parametrize("moe_backend", ["auto", "triton", None]) -def test_other_moe_backends_keep_existing_refit_lifecycle(monkeypatch, moe_backend): +def test_realized_moe_backend_controls_native_refit_lifecycle(): from nemo_rl.models.generation.vllm import vllm_backend - from nemo_rl.models.generation.vllm.quantization import fp8 ext = vllm_backend.VllmInternalWorkerExtension.__new__( vllm_backend.VllmInternalWorkerExtension ) ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("TRITON"), vllm_config=SimpleNamespace( - kernel_config=SimpleNamespace(moe_backend=moe_backend) - ) + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, + ), ) - monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) assert ext._uses_unquantized_flashinfer_trtllm() is False + ext.model_runner.model = _make_unquantized_moe_model("FlashInfer TRTLLM") + ext.model_runner.vllm_config.kernel_config.moe_backend = "auto" + + assert ext._uses_unquantized_flashinfer_trtllm() is True + @pytest.mark.vllm -def test_unquantized_reload_rejects_cotrained_mtp_during_prepare(monkeypatch): +def test_quantized_model_does_not_use_unquantized_refit_lifecycle(): from nemo_rl.models.generation.vllm import vllm_backend - from nemo_rl.models.generation.vllm.quantization import fp8 ext = vllm_backend.VllmInternalWorkerExtension.__new__( vllm_backend.VllmInternalWorkerExtension ) ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM"), + vllm_config=SimpleNamespace(quant_config=object()), + ) + + assert ext._uses_unquantized_flashinfer_trtllm() is False + + +@pytest.mark.vllm +def test_unquantized_reload_rejects_cotrained_mtp_during_prepare(): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM"), vllm_config=SimpleNamespace( - kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") - ) + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, + ), ) ext._mtp_drafter_refit_enabled = lambda: True - monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) with pytest.raises(RuntimeError, match="co-trained MTP drafter"): ext.prepare_refit_info({"model.weight": object()}) @@ -319,21 +373,20 @@ def test_unquantized_reload_rejects_cotrained_mtp_during_prepare(monkeypatch): @pytest.mark.vllm def test_failed_unquantized_reload_marks_worker_unusable(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend - from nemo_rl.models.generation.vllm.quantization import fp8 ext = vllm_backend.VllmInternalWorkerExtension.__new__( vllm_backend.VllmInternalWorkerExtension ) ext.model_runner = SimpleNamespace( - model=object(), + model=_make_unquantized_moe_model("FlashInfer TRTLLM"), vllm_config=SimpleNamespace( - kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm") + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, ), ) ext.model_config = object() ext.device = torch.device("cpu") ext._mtp_drafter_refit_enabled = lambda: False - monkeypatch.setattr(fp8, "is_fp8_model", lambda _: False) monkeypatch.setattr( "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() ) @@ -581,6 +634,54 @@ def test_load_mtp_weights_from_disk_loads_only_mtp_layer(tmp_path, monkeypatch): process_weights.assert_called_once() +@pytest.mark.vllm +def test_load_mtp_weights_from_disk_uses_layerwise_reload_for_trtllm( + tmp_path, monkeypatch +): + """TRTLLM draft weights reload into their preserved runtime storage.""" + model_dir = tmp_path / "ckpt" + _write_sharded_checkpoint( + model_dir, + { + "model-00001-of-00001.safetensors": { + "model.layers.2.mlp.up_proj.weight": torch.randn(4, 4), + } + }, + ) + ext = _make_extension_with_drafter(mtp_start_layer_idx=2, num_mtp_layers=1) + draft_model = ext._get_drafter_model() + draft_model.modules = _make_unquantized_moe_model("FlashInfer TRTLLM").modules + call_order = [] + + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda model: call_order.append(("initialize", model)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", + lambda model, config: call_order.append(("finalize", model, config)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.process_weights_after_loading", + lambda *_: pytest.fail("TRTLLM draft reload must use the layerwise path"), + ) + ext._load_draft_weights.side_effect = lambda _: call_order.append("load") + + assert ext.load_mtp_weights_from_disk(str(model_dir)) is True + assert call_order == [ + ("initialize", draft_model), + "load", + ( + "finalize", + draft_model, + ext.model_runner.vllm_config.speculative_config.draft_model_config, + ), + ] + + @pytest.mark.vllm @pytest.mark.parametrize("is_last_rank", [False, True]) def test_load_mtp_weights_from_disk_without_drafter( From 7180bff8d8fc439ed03de46da325d9c74e0b9edf Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 7 Aug 2026 22:40:44 -0700 Subject: [PATCH 10/68] test(vllm): use module-shaped refit fixtures Signed-off-by: seonjinn --- tests/unit/models/generation/test_vllm_backend.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index e88318434d7..56d0e89ea04 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -33,7 +33,7 @@ def _make_collective_update_extension(backend): state_info = object() ext.state_dict_info = {"model.weight": state_info} ext.model_update_group = object() - ext.model_runner = SimpleNamespace(model=object(), vllm_config=object()) + ext.model_runner = SimpleNamespace(model=torch.nn.Module(), vllm_config=object()) ext.model_config = object() ext.device = object() return ext, state_info @@ -68,7 +68,9 @@ def _make_extension_with_drafter(mtp_start_layer_idx, num_mtp_layers): mtp_start_layer_idx=mtp_start_layer_idx, num_mtp_layers=num_mtp_layers ) ext.model_runner = MagicMock() - ext.model_runner.drafter.model = SimpleNamespace(model=predictor) + draft_model = torch.nn.Module() + setattr(draft_model, "model", predictor) + ext.model_runner.drafter.model = draft_model # Isolate this test from _load_draft_weights internals. ext._load_draft_weights = MagicMock() return ext @@ -415,7 +417,7 @@ def test_update_weights_from_collective_processes_weights_after_loading( call_order = [] process_calls = [] - draft_model = object() if with_mtp else None + draft_model = torch.nn.Module() if with_mtp else None draft_model_config = object() if with_mtp else None def process_weights_after_loading(model, model_config, device): From 95452e91042784df0e41750204f50b4f29d7bb58 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 21:31:48 -0700 Subject: [PATCH 11/68] test(vllm): address native refit self-review Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 23 +++++++-- .../models/generation/test_vllm_backend.py | 49 +++++++++++++++++++ 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 7d09ca8e616..0d8deae010d 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -225,6 +225,8 @@ class VllmInternalWorkerExtension: _sparse_delta_applier: Any = None _nrl_named_parameters: dict[str, torch.nn.Parameter] _nrl_layerwise_reload_active: bool = False + # Initialization detaches parameters, so any later failure leaves this + # worker unsafe to reuse. Keep the original failure for the worker lifetime. _nrl_layerwise_reload_failure: Exception | None = None def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: @@ -237,6 +239,7 @@ def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: def _load_full_hf_weights( self, policy_weights: list[tuple[str, torch.Tensor]] ) -> None: + """Load HF weights and detach any deferred reload tensors from transport storage.""" if not getattr(self, "_nrl_layerwise_reload_active", False): self.model_runner.model.load_weights(weights=policy_weights) return @@ -690,9 +693,11 @@ def _get_sparse_delta_applier(self) -> Any: return self._sparse_delta_applier def _supports_unquantized_flashinfer_trtllm_refit(self) -> bool: + """Whether this worker supports native unquantized TRTLLM refits.""" return True def _uses_unquantized_flashinfer_trtllm(self) -> bool: + """Detect a realized unquantized FlashInfer TRTLLM MoE backend.""" if not self._supports_unquantized_flashinfer_trtllm_refit(): return False model_runner = getattr(self, "model_runner", None) @@ -705,6 +710,7 @@ def _uses_unquantized_flashinfer_trtllm(self) -> bool: return _model_uses_unquantized_flashinfer_trtllm(self.model_runner.model) def _validate_weight_update_compatibility(self) -> None: + """Reject unsupported native layerwise refit combinations.""" if ( self._uses_unquantized_flashinfer_trtllm() and self._mtp_drafter_refit_enabled() @@ -718,7 +724,11 @@ def _validate_weight_update_compatibility(self) -> None: def _weight_update_lifecycle( self, transport: WeightUpdateTransport ) -> Iterator[WeightUpdateFinalizer]: - """Provide setup/finalization around a transport-owned weight update.""" + """Provide setup/finalization around a transport-owned weight update. + + Native reload initialization invalidates the old runtime layout. Any + subsequent exception therefore marks this worker permanently unusable. + """ del transport if self._uses_unquantized_flashinfer_trtllm(): self._validate_weight_update_compatibility() @@ -1224,9 +1234,14 @@ def _recv_one_param(param_info, group, stream): process_weights_after_loading, ) - # Finalize post-load weight processing: dense Linear + attention/MLA, and - # crucially the per-MoE-backend w13 layout (FlashInfer CUTLASS/TRTLLM) that - # the canonical [gate; up] bulk write above defers to here. + # This direct-buffer transport intentionally does not use the native + # layerwise lifecycle. Its mapping is cached against live parameter + # storage at setup, and every refit rewrites canonical fused-weight + # regions before this backend-specific repack. Rebuilding parameters + # would invalidate that mapping. + # + # Finalize dense Linear + attention/MLA state and the per-MoE-backend + # w13 layout (FlashInfer CUTLASS/TRTLLM). with set_current_vllm_config(self.model_runner.vllm_config): process_weights_after_loading( self.model_runner.model, self.model_config, self.device diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 56d0e89ea04..3a80e2a67f1 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -253,6 +253,30 @@ def test_layerwise_reload_preserves_weight_load_error(monkeypatch, caplog): assert "Failed to detach deferred weights" in caplog.text +@pytest.mark.vllm +def test_layerwise_reload_propagates_detach_error_after_successful_load(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + detach_error = RuntimeError("detach failed") + model = SimpleNamespace(load_weights=MagicMock()) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model) + ext._nrl_layerwise_reload_active = True + monkeypatch.setattr( + vllm_backend, + "_detach_pending_layerwise_weights", + MagicMock(side_effect=detach_error), + ) + + with pytest.raises(RuntimeError, match="detach failed") as exc_info: + ext._load_full_hf_weights([("model.weight", torch.ones(1))]) + + assert exc_info.value is detach_error + model.load_weights.assert_called_once() + + @pytest.mark.vllm def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend @@ -350,6 +374,31 @@ def test_quantized_model_does_not_use_unquantized_refit_lifecycle(): assert ext._uses_unquantized_flashinfer_trtllm() is False +@pytest.mark.vllm +@pytest.mark.parametrize( + ("moe_backend", "quant_config", "expected"), + [ + ("FlashInfer TRTLLM", None, True), + ("TRITON", None, False), + ("FlashInfer TRTLLM", object(), False), + ], +) +def test_weight_update_errors_are_fatal_only_for_native_trtllm_refit( + moe_backend, quant_config, expected +): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model(moe_backend), + vllm_config=SimpleNamespace(quant_config=quant_config), + ) + + assert ext._weight_update_errors_are_fatal() is expected + + @pytest.mark.vllm def test_unquantized_reload_rejects_cotrained_mtp_during_prepare(): from nemo_rl.models.generation.vllm import vllm_backend From 8881a6df2a3f8e424b20d5f18464613a3b1f18ff Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 21:52:46 -0700 Subject: [PATCH 12/68] test(vllm): cover refit buffer reuse Signed-off-by: seonjinn --- .../models/generation/test_vllm_backend.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 3a80e2a67f1..9aa9d821214 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -135,6 +135,27 @@ def _make_unquantized_moe_model(moe_backend: str) -> SimpleNamespace: return SimpleNamespace(modules=lambda: [module]) +class _DeferredReloadLayer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.first = torch.nn.Parameter(torch.zeros(2)) + self.second = torch.nn.Parameter(torch.zeros(2)) + + +class _DeferredReloadModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.layer = _DeferredReloadLayer() + + def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> None: + params = dict(self.named_parameters()) + for name, loaded_weight in weights: + param = params[name] + weight_loader = getattr(param, "weight_loader", None) + assert callable(weight_loader) + weight_loader(param, loaded_weight) + + @pytest.mark.vllm def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend @@ -202,6 +223,42 @@ def set_current_vllm_config(config): ext._maybe_process_fp8_kv_cache.assert_not_called() +@pytest.mark.vllm +def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from vllm.model_executor.model_loader.reload import record_metadata_for_reloading + + model = _DeferredReloadModel() + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=object()) + ext.model_config = None + ext.device = torch.device("cpu") + ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._validate_weight_update_compatibility = lambda: None + ext._maybe_process_mtp_drafter_after_loading = MagicMock() + + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + transport_buffer = torch.empty(2) + record_metadata_for_reloading(model) + + with ext._weight_update_lifecycle("collective") as finalize: + transport_buffer.copy_(torch.tensor([1.0, 2.0])) + ext._load_full_hf_weights([("layer.first", transport_buffer)]) + + transport_buffer.copy_(torch.tensor([7.0, 8.0])) + ext._load_full_hf_weights([("layer.second", transport_buffer)]) + finalize() + + torch.testing.assert_close(model.layer.first, torch.tensor([1.0, 2.0])) + torch.testing.assert_close(model.layer.second, torch.tensor([7.0, 8.0])) + + @pytest.mark.vllm def test_layerwise_reload_detaches_deferred_transport_weights(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend From cfd6cab291b250dff2da7d74bf80f2cd0ecfd048 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 22:01:06 -0700 Subject: [PATCH 13/68] test(vllm): satisfy import ordering Signed-off-by: seonjinn --- tests/unit/models/generation/test_vllm_backend.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 9aa9d821214..62a4977e0cc 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -225,9 +225,10 @@ def set_current_vllm_config(config): @pytest.mark.vllm def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypatch): - from nemo_rl.models.generation.vllm import vllm_backend from vllm.model_executor.model_loader.reload import record_metadata_for_reloading + from nemo_rl.models.generation.vllm import vllm_backend + model = _DeferredReloadModel() ext = vllm_backend.VllmInternalWorkerExtension.__new__( vllm_backend.VllmInternalWorkerExtension From 87af15c679729cb73ab6aaed10118065dab47975 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 11:23:55 -0700 Subject: [PATCH 14/68] fix(vllm): support native BF16 TRTLLM refit Signed-off-by: seonjinn --- .../models/generation/vllm_quant_backend.py | 3 + .../models/generation/vllm/vllm_backend.py | 179 ++++++- .../models/generation/test_vllm_backend.py | 450 +++++++++++++++++- 3 files changed, 617 insertions(+), 15 deletions(-) diff --git a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py index 7332339037b..1f9a5abe282 100644 --- a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py +++ b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py @@ -442,6 +442,9 @@ class VllmQuantInternalWorkerExtension(VllmInternalWorkerExtension): _nrl_w13_num_shards_by_prefix: dict[str, int] _nrl_modelopt_reload_roots: tuple[torch.nn.Module, ...] | None = None + def _supports_unquantized_flashinfer_trtllm_refit(self) -> bool: + return False + def maybe_init_zmq(self) -> None: """Use a longer timeout only for ModelOpt real-quant refits.""" super().maybe_init_zmq() diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index de790b642d5..5c166077208 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -71,6 +71,46 @@ class IPCWeightManifestError(RuntimeError): """An IPC transfer did not match the prepared state-dict manifest.""" +def _detach_pending_layerwise_weights( + model: torch.nn.Module, source_storage_ptrs: set[int] +) -> None: + """Detach deferred reload weights from a reusable transport buffer.""" + if not source_storage_ptrs: + return + + # Keep reload internals off the normal non-layerwise weight-loading path. + from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info + + for module in model.modules(): + info = get_layerwise_info(module) + for _, arguments in info.loaded_weights: + loaded_weight = arguments.arguments.get("loaded_weight") + if not isinstance(loaded_weight, torch.Tensor): + continue + if loaded_weight.untyped_storage().data_ptr() in source_storage_ptrs: + arguments.arguments["loaded_weight"] = loaded_weight.clone() + + +def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: + """Return whether a model realized the unquantized TRTLLM MoE backend.""" + # Import backend types only when inspecting a constructed vLLM model. + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + ) + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, + ) + + return any( + isinstance( + quant_method := getattr(module, "quant_method", None), + UnquantizedFusedMoEMethod, + ) + and quant_method.unquantized_backend is UnquantizedMoeBackend.FLASHINFER_TRTLLM + for module in model.modules() + ) + + class _IPCWeightManifest: """Validate an IPC stream against its prepared state-dict manifest.""" @@ -184,6 +224,10 @@ class VllmInternalWorkerExtension: _mtp_drafter_from_disk: bool = False _sparse_delta_applier: Any = None _nrl_named_parameters: dict[str, torch.nn.Parameter] + _nrl_layerwise_reload_active: bool = False + # Initialization detaches parameters, so any later failure leaves this + # worker unsafe to reuse. Keep the original failure for the worker lifetime. + _nrl_layerwise_reload_failure: Exception | None = None def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: params = getattr(self, "_nrl_named_parameters", None) @@ -195,7 +239,31 @@ def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: def _load_full_hf_weights( self, policy_weights: list[tuple[str, torch.Tensor]] ) -> None: - self.model_runner.model.load_weights(weights=policy_weights) + """Load HF weights and detach any deferred reload tensors from transport storage.""" + if not getattr(self, "_nrl_layerwise_reload_active", False): + self.model_runner.model.load_weights(weights=policy_weights) + return + + source_storage_ptrs = { + tensor.untyped_storage().data_ptr() for _, tensor in policy_weights + } + load_error: Exception | None = None + try: + self.model_runner.model.load_weights(weights=policy_weights) + except Exception as error: + load_error = error + raise + finally: + try: + _detach_pending_layerwise_weights( + self.model_runner.model, source_storage_ptrs + ) + except Exception: + if load_error is None: + raise + logger.exception( + "Failed to detach deferred weights after a weight load failure" + ) def _load_hf_weights(self, policy_weights: list[tuple[str, torch.Tensor]]) -> None: from nemo_rl.models.generation.vllm.quantization import fp8 @@ -320,6 +388,7 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: state_dict_info (dict): A dictionary containing the info for refit. e.g. {tensor_name: (shape, dtype)} """ + self._validate_weight_update_compatibility() self.state_dict_info = state_dict_info # pyrefly: ignore[implicitly-defined-attribute] This class does not define __init__ so assignments like this should be ignored def prepare_sparse_delta_refit_info( @@ -546,20 +615,38 @@ def load_mtp_weights_from_disk(self, model_path: str) -> bool: f"include MTP layer weights to run deepseek_mtp speculative decoding." ) - self._load_draft_weights(weights) + draft_model_config = ( + self.model_runner.vllm_config.speculative_config.draft_model_config + ) # The MTP block contains MoE experts whose weights need post-load # processing (e.g. grouped-GEMM layout), matching the main-model path. + # Keep vLLM reload internals off the normal draft-loading path. from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.utils import ( - process_weights_after_loading, - ) - draft_model_config = ( - self.model_runner.vllm_config.speculative_config.draft_model_config - ) - with set_current_vllm_config(self.model_runner.vllm_config): - process_weights_after_loading(draft_model, draft_model_config, self.device) + if self._supports_unquantized_flashinfer_trtllm_refit() and ( + _model_uses_unquantized_flashinfer_trtllm(draft_model) + ): + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + ) + + with set_current_vllm_config(self.model_runner.vllm_config): + with torch.device(self.device): + initialize_layerwise_reload(draft_model) + self._load_draft_weights(weights) + finalize_layerwise_reload(draft_model, draft_model_config) + else: + from vllm.model_executor.model_loader.utils import ( + process_weights_after_loading, + ) + + self._load_draft_weights(weights) + with set_current_vllm_config(self.model_runner.vllm_config): + process_weights_after_loading( + draft_model, draft_model_config, self.device + ) # Mark that the MTP drafter is served from a one-time disk load so refit # does not re-load or re-process these static weights. self._mtp_drafter_from_disk = True @@ -605,12 +692,80 @@ def _get_sparse_delta_applier(self) -> Any: ) return self._sparse_delta_applier + def _supports_unquantized_flashinfer_trtllm_refit(self) -> bool: + """Whether this worker supports native unquantized TRTLLM refits.""" + return True + + def _uses_unquantized_flashinfer_trtllm(self) -> bool: + """Detect a realized unquantized FlashInfer TRTLLM MoE backend.""" + if not self._supports_unquantized_flashinfer_trtllm_refit(): + return False + model_runner = getattr(self, "model_runner", None) + vllm_config = getattr(model_runner, "vllm_config", None) + if vllm_config is None: + return False + if getattr(vllm_config, "quant_config", None) is not None: + return False + + return _model_uses_unquantized_flashinfer_trtllm(self.model_runner.model) + + def _validate_weight_update_compatibility(self) -> None: + """Reject unsupported native layerwise refit combinations.""" + if ( + self._uses_unquantized_flashinfer_trtllm() + and self._mtp_drafter_refit_enabled() + ): + raise RuntimeError( + "Unquantized FlashInfer TRTLLM refit does not yet support " + "a co-trained MTP drafter" + ) + @contextmanager def _weight_update_lifecycle( self, transport: WeightUpdateTransport ) -> Iterator[WeightUpdateFinalizer]: - """Provide setup/finalization around a transport-owned weight update.""" + """Provide setup/finalization around a transport-owned weight update. + + Native reload initialization invalidates the old runtime layout. Any + subsequent exception therefore marks this worker permanently unusable. + """ del transport + if self._uses_unquantized_flashinfer_trtllm(): + self._validate_weight_update_compatibility() + previous_failure = self._nrl_layerwise_reload_failure + if previous_failure is not None: + raise RuntimeError( + "The vLLM worker is unusable after a failed native layerwise refit" + ) from previous_failure + # Load vLLM reload internals only for the native layerwise path. + from vllm.config import set_current_vllm_config + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + ) + + model = self.model_runner.model + + def finalize() -> None: + with torch.device(self.device): + finalize_layerwise_reload(model, self.model_config) + self._maybe_process_mtp_drafter_after_loading() + torch.accelerator.synchronize() + + try: + with set_current_vllm_config(self.model_runner.vllm_config): + with torch.device(self.device): + initialize_layerwise_reload(model) + self._nrl_layerwise_reload_active = True + yield finalize + except Exception as error: + self._nrl_layerwise_reload_failure = error + raise + finally: + self._nrl_layerwise_reload_active = False + + return + from vllm.config import set_current_vllm_config from vllm.model_executor.model_loader.utils import ( process_weights_after_loading, @@ -630,7 +785,7 @@ def finalize() -> None: def _weight_update_errors_are_fatal(self) -> bool: """Whether transport errors should propagate instead of returning False.""" - return False + return self._uses_unquantized_flashinfer_trtllm() def _synchronize_before_ipc_data_ack(self) -> None: """Fence work consuming one IPC data batch before its acknowledgment.""" diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index c7e7969ddc9..62a4977e0cc 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -33,7 +33,7 @@ def _make_collective_update_extension(backend): state_info = object() ext.state_dict_info = {"model.weight": state_info} ext.model_update_group = object() - ext.model_runner = SimpleNamespace(model=object(), vllm_config=object()) + ext.model_runner = SimpleNamespace(model=torch.nn.Module(), vllm_config=object()) ext.model_config = object() ext.device = object() return ext, state_info @@ -68,7 +68,9 @@ def _make_extension_with_drafter(mtp_start_layer_idx, num_mtp_layers): mtp_start_layer_idx=mtp_start_layer_idx, num_mtp_layers=num_mtp_layers ) ext.model_runner = MagicMock() - ext.model_runner.drafter.model = SimpleNamespace(model=predictor) + draft_model = torch.nn.Module() + setattr(draft_model, "model", predictor) + ext.model_runner.drafter.model = draft_model # Isolate this test from _load_draft_weights internals. ext._load_draft_weights = MagicMock() return ext @@ -119,6 +121,400 @@ def _make_mtp_refit_extension( return ext, drafter_model +def _make_unquantized_moe_model(moe_backend: str) -> SimpleNamespace: + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + ) + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, + ) + + quant_method = UnquantizedFusedMoEMethod.__new__(UnquantizedFusedMoEMethod) + quant_method.unquantized_backend = UnquantizedMoeBackend(moe_backend) + module = SimpleNamespace(quant_method=quant_method) + return SimpleNamespace(modules=lambda: [module]) + + +class _DeferredReloadLayer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.first = torch.nn.Parameter(torch.zeros(2)) + self.second = torch.nn.Parameter(torch.zeros(2)) + + +class _DeferredReloadModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.layer = _DeferredReloadLayer() + + def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> None: + params = dict(self.named_parameters()) + for name, loaded_weight in weights: + param = params[name] + weight_loader = getattr(param, "weight_loader", None) + assert callable(weight_loader) + weight_loader(param, loaded_weight) + + +@pytest.mark.vllm +def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + call_order = [] + model = _make_unquantized_moe_model("FlashInfer TRTLLM") + model_config = object() + vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="auto"), quant_config=None + ) + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=vllm_config) + ext.model_config = model_config + ext.device = torch.device("cpu") + ext._maybe_process_mtp_drafter_after_loading = lambda: call_order.append("mtp") + ext._maybe_process_fp8_kv_cache = MagicMock() + + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + @contextlib.contextmanager + def set_current_vllm_config(config): + assert config is vllm_config + call_order.append("config_enter") + try: + yield + finally: + call_order.append("config_exit") + + monkeypatch.setattr("vllm.config.set_current_vllm_config", set_current_vllm_config) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda reload_model: call_order.append(("initialize", reload_model)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", + lambda reload_model, config: call_order.append( + ("finalize", reload_model, config) + ), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.process_weights_after_loading", + lambda *_args: pytest.fail( + "unquantized refit must use vLLM's native layerwise reload lifecycle" + ), + ) + + for _ in range(2): + with ext._weight_update_lifecycle("collective") as finalize: + call_order.append("load") + finalize() + assert ext._nrl_layerwise_reload_active is False + + expected_cycle = [ + "config_enter", + ("initialize", model), + "load", + ("finalize", model, model_config), + "mtp", + "config_exit", + ] + assert call_order == expected_cycle * 2 + ext._maybe_process_fp8_kv_cache.assert_not_called() + + +@pytest.mark.vllm +def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypatch): + from vllm.model_executor.model_loader.reload import record_metadata_for_reloading + + from nemo_rl.models.generation.vllm import vllm_backend + + model = _DeferredReloadModel() + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=object()) + ext.model_config = None + ext.device = torch.device("cpu") + ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._validate_weight_update_compatibility = lambda: None + ext._maybe_process_mtp_drafter_after_loading = MagicMock() + + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + transport_buffer = torch.empty(2) + record_metadata_for_reloading(model) + + with ext._weight_update_lifecycle("collective") as finalize: + transport_buffer.copy_(torch.tensor([1.0, 2.0])) + ext._load_full_hf_weights([("layer.first", transport_buffer)]) + + transport_buffer.copy_(torch.tensor([7.0, 8.0])) + ext._load_full_hf_weights([("layer.second", transport_buffer)]) + finalize() + + torch.testing.assert_close(model.layer.first, torch.tensor([1.0, 2.0])) + torch.testing.assert_close(model.layer.second, torch.tensor([7.0, 8.0])) + + +@pytest.mark.vllm +def test_layerwise_reload_detaches_deferred_transport_weights(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + source = torch.ones(4) + unrelated = torch.full((2,), 7.0) + source_args = SimpleNamespace(arguments={"loaded_weight": source[:2]}) + unrelated_args = SimpleNamespace(arguments={"loaded_weight": unrelated}) + model = SimpleNamespace(modules=lambda: [object()]) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.layerwise.get_layerwise_info", + lambda _module: SimpleNamespace( + loaded_weights=[("source", source_args), ("other", unrelated_args)] + ), + ) + + vllm_backend._detach_pending_layerwise_weights( + model, {source.untyped_storage().data_ptr()} + ) + + detached = source_args.arguments["loaded_weight"] + assert detached.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + assert unrelated_args.arguments["loaded_weight"] is unrelated + source.zero_() + torch.testing.assert_close(detached, torch.ones(2)) + + +@pytest.mark.vllm +def test_layerwise_reload_preserves_weight_load_error(monkeypatch, caplog): + from nemo_rl.models.generation.vllm import vllm_backend + + load_error = RuntimeError("load failed") + model = SimpleNamespace(load_weights=MagicMock(side_effect=load_error)) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model) + ext._nrl_layerwise_reload_active = True + monkeypatch.setattr( + vllm_backend, + "_detach_pending_layerwise_weights", + MagicMock(side_effect=RuntimeError("detach failed")), + ) + + with pytest.raises(RuntimeError, match="load failed") as exc_info: + ext._load_full_hf_weights([("model.weight", torch.ones(1))]) + + assert exc_info.value is load_error + assert "Failed to detach deferred weights" in caplog.text + + +@pytest.mark.vllm +def test_layerwise_reload_propagates_detach_error_after_successful_load(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + detach_error = RuntimeError("detach failed") + model = SimpleNamespace(load_weights=MagicMock()) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model) + ext._nrl_layerwise_reload_active = True + monkeypatch.setattr( + vllm_backend, + "_detach_pending_layerwise_weights", + MagicMock(side_effect=detach_error), + ) + + with pytest.raises(RuntimeError, match="detach failed") as exc_info: + ext._load_full_hf_weights([("model.weight", torch.ones(1))]) + + assert exc_info.value is detach_error + model.load_weights.assert_called_once() + + +@pytest.mark.vllm +def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + model = object() + model_config = object() + vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=object(), + ) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=vllm_config) + ext.model_config = model_config + ext.device = torch.device("cpu") + ext._maybe_process_mtp_drafter_after_loading = MagicMock() + ext._maybe_process_fp8_kv_cache = MagicMock() + + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + process = MagicMock() + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.process_weights_after_loading", + process, + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda _: pytest.fail("FP8 must not use the unquantized reload lifecycle"), + ) + + with ext._weight_update_lifecycle("collective") as finalize: + finalize() + + process.assert_called_once_with(model, model_config, ext.device) + ext._maybe_process_mtp_drafter_after_loading.assert_called_once_with() + ext._maybe_process_fp8_kv_cache.assert_called_once_with() + + +@pytest.mark.vllm +def test_extension_capability_can_disable_unquantized_reload(): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=object(), + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, + ), + ) + ext._supports_unquantized_flashinfer_trtllm_refit = lambda: False + + assert ext._uses_unquantized_flashinfer_trtllm() is False + + +@pytest.mark.vllm +def test_realized_moe_backend_controls_native_refit_lifecycle(): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("TRITON"), + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, + ), + ) + + assert ext._uses_unquantized_flashinfer_trtllm() is False + + ext.model_runner.model = _make_unquantized_moe_model("FlashInfer TRTLLM") + ext.model_runner.vllm_config.kernel_config.moe_backend = "auto" + + assert ext._uses_unquantized_flashinfer_trtllm() is True + + +@pytest.mark.vllm +def test_quantized_model_does_not_use_unquantized_refit_lifecycle(): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM"), + vllm_config=SimpleNamespace(quant_config=object()), + ) + + assert ext._uses_unquantized_flashinfer_trtllm() is False + + +@pytest.mark.vllm +@pytest.mark.parametrize( + ("moe_backend", "quant_config", "expected"), + [ + ("FlashInfer TRTLLM", None, True), + ("TRITON", None, False), + ("FlashInfer TRTLLM", object(), False), + ], +) +def test_weight_update_errors_are_fatal_only_for_native_trtllm_refit( + moe_backend, quant_config, expected +): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model(moe_backend), + vllm_config=SimpleNamespace(quant_config=quant_config), + ) + + assert ext._weight_update_errors_are_fatal() is expected + + +@pytest.mark.vllm +def test_unquantized_reload_rejects_cotrained_mtp_during_prepare(): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM"), + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, + ), + ) + ext._mtp_drafter_refit_enabled = lambda: True + + with pytest.raises(RuntimeError, match="co-trained MTP drafter"): + ext.prepare_refit_info({"model.weight": object()}) + + assert not hasattr(ext, "state_dict_info") + + +@pytest.mark.vllm +def test_failed_unquantized_reload_marks_worker_unusable(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM"), + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + quant_config=None, + ), + ) + ext.model_config = object() + ext.device = torch.device("cpu") + ext._mtp_drafter_refit_enabled = lambda: False + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda _: None, + ) + + failure = RuntimeError("load failed") + with pytest.raises(RuntimeError, match="load failed"): + with ext._weight_update_lifecycle("collective"): + raise failure + + assert ext._nrl_layerwise_reload_failure is failure + with pytest.raises(RuntimeError, match="worker is unusable"): + with ext._weight_update_lifecycle("collective"): + pass + + @pytest.mark.vllm @pytest.mark.parametrize("with_mtp", [False, True]) def test_update_weights_from_collective_processes_weights_after_loading( @@ -128,7 +524,7 @@ def test_update_weights_from_collective_processes_weights_after_loading( call_order = [] process_calls = [] - draft_model = object() if with_mtp else None + draft_model = torch.nn.Module() if with_mtp else None draft_model_config = object() if with_mtp else None def process_weights_after_loading(model, model_config, device): @@ -347,6 +743,54 @@ def test_load_mtp_weights_from_disk_loads_only_mtp_layer(tmp_path, monkeypatch): process_weights.assert_called_once() +@pytest.mark.vllm +def test_load_mtp_weights_from_disk_uses_layerwise_reload_for_trtllm( + tmp_path, monkeypatch +): + """TRTLLM draft weights reload into their preserved runtime storage.""" + model_dir = tmp_path / "ckpt" + _write_sharded_checkpoint( + model_dir, + { + "model-00001-of-00001.safetensors": { + "model.layers.2.mlp.up_proj.weight": torch.randn(4, 4), + } + }, + ) + ext = _make_extension_with_drafter(mtp_start_layer_idx=2, num_mtp_layers=1) + draft_model = ext._get_drafter_model() + draft_model.modules = _make_unquantized_moe_model("FlashInfer TRTLLM").modules + call_order = [] + + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda model: call_order.append(("initialize", model)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", + lambda model, config: call_order.append(("finalize", model, config)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.process_weights_after_loading", + lambda *_: pytest.fail("TRTLLM draft reload must use the layerwise path"), + ) + ext._load_draft_weights.side_effect = lambda _: call_order.append("load") + + assert ext.load_mtp_weights_from_disk(str(model_dir)) is True + assert call_order == [ + ("initialize", draft_model), + "load", + ( + "finalize", + draft_model, + ext.model_runner.vllm_config.speculative_config.draft_model_config, + ), + ] + + @pytest.mark.vllm @pytest.mark.parametrize("is_last_rank", [False, True]) def test_load_mtp_weights_from_disk_without_drafter( From 4edab6833499d915728046d23ed56c7be865d6ef Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 10:07:10 -0700 Subject: [PATCH 15/68] test(vllm): cover TRTLLM NCCL reshard refit Signed-off-by: seonjinn --- .../generation/test_nccl_reshard_backend.py | 142 ++++++++++++++++++ 1 file changed, 142 insertions(+) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index 9fba296e0e7..bfbbbfe9c88 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -23,10 +23,13 @@ skipped where vllm is unavailable. """ +import contextlib from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch +from torch.distributed._tensor import Shard pytest.importorskip("vllm") # module-top `import vllm` in vllm_backend @@ -35,6 +38,7 @@ ) from nemo_rl.weight_sync.nccl_reshard_utils import ( # noqa: E402 HFToLocalParamMap, + MeshInfo, ) pytestmark = pytest.mark.vllm @@ -313,3 +317,141 @@ def test_build_hf_to_local_param_map_specs_and_roundtrip(): egctx.buf.fill_(5.0) eg.post(egctx) assert torch.equal(w13[:, 0:Pl, :], torch.full_like(w13[:, 0:Pl, :], 5.0)) + + +def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): + """Packed TRTLLM storage receives canonical EP-local weights via load_weights.""" + H, E, P = 16, 4, 32 + expert_name = "model.layers.0.mlp.experts.gate_proj.weight" + refit_info = { + "gen_tp_size": 2, + "layer_names": ["model.layers.0"], + "per_layer_params": { + "model.layers.0": [ + { + "name": expert_name, + "global_shape": [E, P, H], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "gate_proj", + "dst_mesh_info": MeshInfo(torch.tensor([8, 9])), + "dst_placements": [Shard(0)], + } + ] + }, + } + packed_w13 = torch.full((128, 16, 24, 64), 7.0) + ext = _make_ext( + { + "model.layers.0.mlp.experts.routed_experts.w13_weight": packed_w13, + } + ) + ext.device = torch.device("cpu") + ext.pp_comm_groups = {0: SimpleNamespace(rank=9)} + ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._load_full_hf_weights = MagicMock() + + spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) + assert spec is not None and spec.pre is not None and spec.post is not None + + ctx = spec.pre(spec.base) + assert ctx.buf.shape == (2, P, H) + assert ctx.buf.dtype == torch.bfloat16 + ctx.buf[0].fill_(2.0) + ctx.buf[1].fill_(3.0) + spec.post(ctx) + + loaded_weights = ext._load_full_hf_weights.call_args.args[0] + assert [name for name, _ in loaded_weights] == [ + "model.layers.0.mlp.experts.2.gate_proj.weight", + "model.layers.0.mlp.experts.3.gate_proj.weight", + ] + torch.testing.assert_close(loaded_weights[0][1], torch.full((P, H), 2.0)) + torch.testing.assert_close(loaded_weights[1][1], torch.full((P, H), 3.0)) + torch.testing.assert_close(packed_w13, torch.full_like(packed_w13, 7.0)) + + +def test_nccl_reshard_lifecycle_initializes_only_trtllm_moe_modules(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + model = SimpleNamespace() + trtllm_moe = SimpleNamespace() + model_config = object() + vllm_config = object() + call_order = [] + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=vllm_config) + ext.model_config = model_config + ext.device = torch.device("cpu") + ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._validate_weight_update_compatibility = lambda: None + ext._maybe_process_mtp_drafter_after_loading = lambda: call_order.append("mtp") + + monkeypatch.setattr( + vllm_backend, + "_unquantized_flashinfer_trtllm_modules", + lambda _model: [trtllm_moe], + raising=False, + ) + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() + ) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda module: call_order.append(("initialize", module)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", + lambda reload_model, config: call_order.append( + ("finalize", reload_model, config) + ), + ) + + with ext._weight_update_lifecycle("nccl_reshard") as finalize: + call_order.append("transfer") + finalize() + + assert call_order == [ + ("initialize", trtllm_moe), + "transfer", + ("finalize", model, model_config), + "mtp", + ] + + +def test_nccl_reshard_refit_runs_transport_lifecycle(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.nccl_reshard_refit_info = { + "layer_names": [], + "per_layer_params": {}, + "misc_meta": {}, + } + ext._receive_and_load_misc_params = MagicMock() + ext._maybe_process_fp8_kv_cache = MagicMock() + finalize = MagicMock() + lifecycle_calls = [] + + @contextlib.contextmanager + def lifecycle(transport): + lifecycle_calls.append(transport) + yield finalize + + ext._weight_update_lifecycle = lifecycle + monkeypatch.setattr(torch.cuda, "Stream", lambda: object()) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 1) + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.process_weights_after_loading", + lambda *_args: pytest.fail("transport lifecycle must own finalization"), + ) + + assert ext.nccl_reshard_refit() is True + assert lifecycle_calls == ["nccl_reshard"] + finalize.assert_called_once_with() From 06b89670cd090fb7cfd39f129bf3c09121f6a3d4 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 10:21:58 -0700 Subject: [PATCH 16/68] fix(vllm): support TRTLLM NCCL reshard refit Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 156 ++++++++++++++---- 1 file changed, 126 insertions(+), 30 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 5c166077208..9a77b22cf8e 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -91,8 +91,10 @@ def _detach_pending_layerwise_weights( arguments.arguments["loaded_weight"] = loaded_weight.clone() -def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: - """Return whether a model realized the unquantized TRTLLM MoE backend.""" +def _unquantized_flashinfer_trtllm_modules( + model: torch.nn.Module, +) -> list[torch.nn.Module]: + """Return modules that realized the unquantized TRTLLM MoE backend.""" # Import backend types only when inspecting a constructed vLLM model. from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, @@ -101,13 +103,43 @@ def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: UnquantizedFusedMoEMethod, ) - return any( - isinstance( + return [ + module + for module in model.modules() + if isinstance( quant_method := getattr(module, "quant_method", None), UnquantizedFusedMoEMethod, ) and quant_method.unquantized_backend is UnquantizedMoeBackend.FLASHINFER_TRTLLM - for module in model.modules() + ] + + +def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: + """Return whether a model realized the unquantized TRTLLM MoE backend.""" + return bool(_unquantized_flashinfer_trtllm_modules(model)) + + +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 import ( + _compute_shard_slices, + _get_mesh_coords, + ) + + dst_mesh = param_info["dst_mesh_info"] + mesh_coords = _get_mesh_coords(dst_mesh, rank) + if mesh_coords is None: + raise ValueError( + f"Destination rank {rank} is absent from the mesh for " + f"{param_info['name']!r}" + ) + return tuple( + _compute_shard_slices( + param_info["global_shape"], + list(dst_mesh.mesh.shape), + mesh_coords, + param_info["dst_placements"], + ) ) @@ -729,7 +761,6 @@ def _weight_update_lifecycle( Native reload initialization invalidates the old runtime layout. Any subsequent exception therefore marks this worker permanently unusable. """ - del transport if self._uses_unquantized_flashinfer_trtllm(): self._validate_weight_update_compatibility() previous_failure = self._nrl_layerwise_reload_failure @@ -745,6 +776,11 @@ def _weight_update_lifecycle( ) model = self.model_runner.model + reload_targets = ( + _unquantized_flashinfer_trtllm_modules(model) + if transport == "nccl_reshard" + else [model] + ) def finalize() -> None: with torch.device(self.device): @@ -755,7 +791,8 @@ def finalize() -> None: try: with set_current_vllm_config(self.model_runner.vllm_config): with torch.device(self.device): - initialize_layerwise_reload(model) + for reload_target in reload_targets: + initialize_layerwise_reload(reload_target) self._nrl_layerwise_reload_active = True yield finalize except Exception as error: @@ -983,14 +1020,76 @@ def post(ctx): return LocalParamSpec(base=vllm_param, pre=pre, post=post) + def _trtllm_grouped_expert_spec(param_info: dict) -> LocalParamSpec: + from torch.distributed._tensor import Shard + + from nemo_rl.weight_sync.nccl_reshard_utils import _STR_TO_DTYPE + + unsupported_shards = [ + placement.dim + for placement in param_info["dst_placements"] + if isinstance(placement, Shard) and placement.dim != 0 + ] + if unsupported_shards: + raise ValueError( + "Unquantized FlashInfer TRTLLM nccl_reshard refit requires " + "expert-parallel destination shards; unsupported tensor shard " + f"dimensions {unsupported_shards} for {param_info['name']!r}" + ) + + pp_stage = param_info.get("pp_stage", 0) + rank = self.pp_comm_groups[pp_stage].rank + local_slices = _local_shard_slices(param_info, rank) + local_shape = tuple( + global_size + if shard_slice.start is None + else shard_slice.stop - shard_slice.start + for global_size, shard_slice in zip( + param_info["global_shape"], local_slices + ) + ) + expert_start = local_slices[0].start or 0 + grouped_proj = param_info["grouped_expert_proj"] + expert_prefix = param_info["name"].rsplit(f".{grouped_proj}.weight", 1)[0] + dtype = _STR_TO_DTYPE[str(param_info["dtype"])] + + def pre(_base: None) -> RefitCtx: + return RefitCtx( + buf=torch.empty(local_shape, dtype=dtype, device=self.device) + ) + + def post(ctx: RefitCtx) -> None: + weights = [ + ( + f"{expert_prefix}.{expert_start + local_idx}." + f"{grouped_proj}.weight", + expert_weight, + ) + for local_idx, expert_weight in enumerate(ctx.buf.unbind(0)) + ] + self._load_full_hf_weights(weights) + + return LocalParamSpec(base=None, pre=pre, post=post) + # Get dict of vllm_param and merged_slice for each hf_name vllm_param_map_and_slices = self._build_hf_to_gen_backend_mapping(refit_info) + param_info_by_name = { + param_info["name"]: param_info + for layer_name in refit_info["layer_names"] + for param_info in refit_info["per_layer_params"][layer_name] + } + use_trtllm_staging = self._uses_unquantized_flashinfer_trtllm() return HFToLocalParamMap( specs={ hf_name: ( - LocalParamSpec(base=vllm_param.data) - if merged_slice is None - else _merged_param_spec(vllm_param, merged_slice) + _trtllm_grouped_expert_spec(param_info_by_name[hf_name]) + if use_trtllm_staging + and param_info_by_name[hf_name].get("grouped_expert_proj") + else ( + LocalParamSpec(base=vllm_param.data) + if merged_slice is None + else _merged_param_spec(vllm_param, merged_slice) + ) ) for hf_name, ( vllm_param, @@ -1145,6 +1244,11 @@ def _to_vllm_name(n: str) -> str: return mapping def nccl_reshard_refit(self) -> bool: + """Receive and finalize one NCCL reshard weight update.""" + with self._weight_update_lifecycle("nccl_reshard") as finalize: + return self._nccl_reshard_refit_impl(finalize) + + def _nccl_reshard_refit_impl(self, finalize: WeightUpdateFinalizer) -> bool: """Receive weights from training workers via xferdtensor. Each HF param's ``LocalParamSpec`` (from ``hf_to_local_param_map``, @@ -1219,26 +1323,18 @@ def _recv_one_param(param_info, group, stream): import time - with self._weight_update_lifecycle("nccl_reshard") as finalize: - misc_t0 = time.perf_counter() - self._receive_and_load_misc_params() - torch.cuda.synchronize() - if torch.distributed.get_rank() == 0: - print( - f"[nccl_reshard_refit] misc recv+load (gen side): " - f"{time.perf_counter() - misc_t0:.2f}s", - flush=True, - ) - torch.cuda.empty_cache() - - # Finalize post-load weight processing: dense Linear + attention/MLA, - # the per-MoE-backend w13 layout (FlashInfer CUTLASS/TRTLLM) that the - # canonical [gate; up] bulk write above defers to here, and the MTP - # drafter's mirror of the same. The FP8 KV-cache per-layer k/v scales - # are finalized by the lifecycle on exit. - finalize() - - torch.cuda.empty_cache() + misc_t0 = time.perf_counter() + self._receive_and_load_misc_params() + torch.cuda.synchronize() + if torch.distributed.get_rank() == 0: + print( + f"[nccl_reshard_refit] misc recv+load (gen side): " + f"{time.perf_counter() - misc_t0:.2f}s", + flush=True, + ) + torch.cuda.empty_cache() + finalize() + torch.cuda.empty_cache() return True def _receive_and_load_misc_params(self) -> None: From 2cadb920cb14a2537c5ee27f2bfd03f40645e37d Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 10:27:20 -0700 Subject: [PATCH 17/68] test(vllm): preserve TRTLLM staging dtype Signed-off-by: seonjinn --- tests/unit/models/generation/test_nccl_reshard_backend.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index bfbbbfe9c88..236c0ba9c1b 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -365,8 +365,12 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): "model.layers.0.mlp.experts.2.gate_proj.weight", "model.layers.0.mlp.experts.3.gate_proj.weight", ] - torch.testing.assert_close(loaded_weights[0][1], torch.full((P, H), 2.0)) - torch.testing.assert_close(loaded_weights[1][1], torch.full((P, H), 3.0)) + torch.testing.assert_close( + loaded_weights[0][1], torch.full((P, H), 2.0, dtype=torch.bfloat16) + ) + torch.testing.assert_close( + loaded_weights[1][1], torch.full((P, H), 3.0, dtype=torch.bfloat16) + ) torch.testing.assert_close(packed_w13, torch.full_like(packed_w13, 7.0)) From 64cbbdc3b65ed66c2cf266f88f57dec83e441507 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 10:32:21 -0700 Subject: [PATCH 18/68] docs(vllm): clarify TRTLLM reshard staging Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/vllm_backend.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 9a77b22cf8e..a44a25ba1c7 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -1005,9 +1005,11 @@ def build_hf_to_local_param_map(self, refit_info: dict) -> HFToLocalParamMap: Wraps the ``(vllm_param, merged_slice)`` resolution from ``_build_hf_to_gen_backend_mapping`` into ``LocalParamSpec``s: - direct (slice ``None``): ``base`` is the live vLLM param; receive in place. - - merged (dense ``gate_up_proj`` / grouped-expert ``w13``): ``pre`` allocs a - recv buffer for this component's ``region`` slice, ``post`` copies it back - (region recomputed each refit to track live storage). + - merged (dense ``gate_up_proj`` / grouped-expert ``w13``): ``pre`` allocates + a receive buffer for this component's ``region`` slice, and ``post`` copies + it back (the region is recomputed each refit to track live storage). + - TRTLLM grouped experts: ``pre`` allocates canonical EP-local BF16 storage, + and ``post`` sends each expert through vLLM's native weight loader. """ def _merged_param_spec(vllm_param, merged_slice): @@ -1020,7 +1022,9 @@ def post(ctx): return LocalParamSpec(base=vllm_param, pre=pre, post=post) - def _trtllm_grouped_expert_spec(param_info: dict) -> LocalParamSpec: + def _trtllm_grouped_expert_spec( + param_info: dict[str, Any], + ) -> LocalParamSpec: from torch.distributed._tensor import Shard from nemo_rl.weight_sync.nccl_reshard_utils import _STR_TO_DTYPE @@ -1045,7 +1049,7 @@ def _trtllm_grouped_expert_spec(param_info: dict) -> LocalParamSpec: if shard_slice.start is None else shard_slice.stop - shard_slice.start for global_size, shard_slice in zip( - param_info["global_shape"], local_slices + param_info["global_shape"], local_slices, strict=True ) ) expert_start = local_slices[0].start or 0 @@ -1256,7 +1260,8 @@ def _nccl_reshard_refit_impl(self, finalize: WeightUpdateFinalizer) -> bool: for a direct param xferdtensor receives straight into the live vLLM param (no hooks); for a merged param (dense gate_up_proj, grouped w13) ``pre`` allocates a temp recv buffer and ``post`` copies the TP-local - slice back into the live merged param. + slice back into the live merged param. TRTLLM grouped experts instead + receive into canonical local tensors and load through vLLM's native path. """ import os from collections import OrderedDict From 9d5a10a58847c8dd2613a9d847e13071410672c9 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 10:43:30 -0700 Subject: [PATCH 19/68] refactor(refit): expose local shard slices Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 23 ++++----------- nemo_rl/weight_sync/nccl_reshard_utils.py | 7 ++--- nemo_rl/weight_sync/xferdtensor.py | 28 +++++++++++++++++++ 3 files changed, 37 insertions(+), 21 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index a44a25ba1c7..1ec005200f9 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -121,25 +121,14 @@ def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: 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 import ( - _compute_shard_slices, - _get_mesh_coords, - ) + from nemo_rl.weight_sync.xferdtensor import get_local_shard_slices dst_mesh = param_info["dst_mesh_info"] - mesh_coords = _get_mesh_coords(dst_mesh, rank) - if mesh_coords is None: - raise ValueError( - f"Destination rank {rank} is absent from the mesh for " - f"{param_info['name']!r}" - ) - return tuple( - _compute_shard_slices( - param_info["global_shape"], - list(dst_mesh.mesh.shape), - mesh_coords, - param_info["dst_placements"], - ) + return get_local_shard_slices( + param_info["global_shape"], + dst_mesh, + param_info["dst_placements"], + rank, ) diff --git a/nemo_rl/weight_sync/nccl_reshard_utils.py b/nemo_rl/weight_sync/nccl_reshard_utils.py index 3f8e638e4cb..7c677fa4c84 100644 --- a/nemo_rl/weight_sync/nccl_reshard_utils.py +++ b/nemo_rl/weight_sync/nccl_reshard_utils.py @@ -91,10 +91,9 @@ class LocalParamSpec: post: ``RefitCtx -> None``; runs after xferdtensor e.g., copy back the received buffer into the merged param. - TODO: A layout that block-permutes the *assembled* param (e.g. FlashInfer - TRTLLM w13) would need a group-level finalize run once after all components - land — a future loop-level addition, not a per-param field. ``pre``/``post`` - covers today's backends (Triton, FlashInfer CUTLASS, Megatron). + Layout-specific backends can receive into canonical storage in ``pre``, load + each logical component in ``post``, and use a transport-level finalizer after + all components land. FlashInfer TRTLLM uses this path for grouped experts. """ base: Any diff --git a/nemo_rl/weight_sync/xferdtensor.py b/nemo_rl/weight_sync/xferdtensor.py index 6f003d58bc4..f0d81e4f7b6 100644 --- a/nemo_rl/weight_sync/xferdtensor.py +++ b/nemo_rl/weight_sync/xferdtensor.py @@ -16,10 +16,13 @@ import logging import os +from collections.abc import Sequence from contextlib import nullcontext +from typing import Any import torch from torch.distributed._tensor import Shard +from torch.distributed.tensor.placement_types import Placement try: from nccl.m2n import ( # pyrefly: ignore[import-error] @@ -255,6 +258,31 @@ def _compute_shard_slices(global_shape, mesh_shape, mesh_coords, placements): return slices +def get_local_shard_slices( + global_shape: list[int] | tuple[int, ...] | torch.Size, + mesh: Any, + placements: Sequence[Placement], + rank: int, +) -> tuple[slice, ...]: + """Return one rank's local slices in a global tensor.""" + mesh_coords = _get_mesh_coords(mesh, rank) + if mesh_coords is None: + raise ValueError(f"Rank {rank} is absent from the destination mesh") + mesh_tensor = getattr(mesh, "mesh", None) + if mesh_tensor is None: + mesh_tensor = getattr(mesh, "_mesh", None) + if mesh_tensor is None: + raise ValueError("DeviceMesh does not expose mesh ranks.") + return tuple( + _compute_shard_slices( + global_shape, + list(mesh_tensor.shape), + mesh_coords, + placements, + ) + ) + + def xferdtensor_golden( src_tensor, src_mesh, From 28445421d502a31cc31975affafbc629be0745a1 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 11:43:32 -0700 Subject: [PATCH 20/68] test(vllm): cover unsupported TRTLLM reshard layout Signed-off-by: seonjinn --- .../generation/test_nccl_reshard_backend.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index 236c0ba9c1b..3517436c48a 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -374,6 +374,38 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): torch.testing.assert_close(packed_w13, torch.full_like(packed_w13, 7.0)) +def test_build_hf_to_local_param_map_rejects_trtllm_tensor_sharding(): + """TRTLLM expert staging supports expert-parallel destination shards only.""" + expert_name = "model.layers.0.mlp.experts.gate_proj.weight" + refit_info = { + "gen_tp_size": 2, + "layer_names": ["model.layers.0"], + "per_layer_params": { + "model.layers.0": [ + { + "name": expert_name, + "global_shape": [4, 32, 16], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "gate_proj", + "dst_mesh_info": MeshInfo(torch.tensor([8, 9])), + "dst_placements": [Shard(1)], + } + ] + }, + } + ext = _make_ext( + { + "model.layers.0.mlp.experts.routed_experts.w13_weight": torch.empty( + 128, 16, 24, 64 + ), + } + ) + ext._uses_unquantized_flashinfer_trtllm = lambda: True + + with pytest.raises(ValueError, match="unsupported tensor shard dimensions"): + ext.build_hf_to_local_param_map(refit_info) + + def test_nccl_reshard_lifecycle_initializes_only_trtllm_moe_modules(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend From ac8e3968a0528cbca2ad1dd4d40e7514b7872da7 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Mon, 17 Aug 2026 14:30:37 -0700 Subject: [PATCH 21/68] fix(vllm): harden TRTLLM NCCL reshard refit Signed-off-by: seonjinn --- docs/design-docs/nccl-reshard-refit.md | 35 ++++++----- .../models/generation/vllm/vllm_backend.py | 35 +++++++++-- nemo_rl/weight_sync/xferdtensor.py | 53 ++++++----------- .../generation/test_nccl_reshard_backend.py | 42 ++++++++++++-- .../models/generation/test_vllm_backend.py | 58 ++++++++++++++++++- .../weight_sync/test_xferdtensor_python.py | 44 ++++++++++++++ 6 files changed, 205 insertions(+), 62 deletions(-) diff --git a/docs/design-docs/nccl-reshard-refit.md b/docs/design-docs/nccl-reshard-refit.md index 8840fa39ca7..36b2ecc7b19 100644 --- a/docs/design-docs/nccl-reshard-refit.md +++ b/docs/design-docs/nccl-reshard-refit.md @@ -36,6 +36,10 @@ single `ValueError` listing every violation. The current requirements are: * **Precision** must match end to end: BF16 train ↔ BF16 gen, or FP8 train (`fp8_param=true` + blockwise recipe) ↔ FP8 gen (`vllm_cfg.precision=fp8`). BF16 train ↔ FP8 gen is not supported yet. +* Unquantized BF16 FlashInfer TRTLLM MoE is supported through vLLM's native + layerwise-reload path. Its grouped expert weights must use expert-parallel + destination sharding with linear expert placement; tensor-sharded expert + destinations and round-robin placement are rejected. * vLLM expert parallelism is supported with the NeMo RL convention `expert_parallel_size == tensor_parallel_size`. * Generation-side, PP > 1 is not supported. @@ -132,10 +136,11 @@ realized **locally**: (sent as-is); grouped MoE experts get a `pre` hook that stacks this rank's per-expert views into a `[num_local_experts, ...]` tensor fresh at each refit. * On the **generation side**, a direct parameter's `base` is the live vLLM parameter - (received into in place); a parameter that is a slice of a fused vLLM tensor (dense - `gate_up_proj`, grouped-expert `w13`/`w2`) gets a `pre` hook that allocates a receive - buffer for its region and a `post` hook that copies the received shard back into the - fused parameter. + (received into in place). Conventional fused parameters use `pre`/`post` hooks to + receive a component and copy it into the appropriate local region. Unquantized + FlashInfer TRTLLM grouped experts instead receive into canonical EP-local staging + tensors; `post` loads each logical expert with its global expert ID through vLLM's + native weight loader. ### Execution Flow: Refit Time @@ -156,7 +161,9 @@ Every training step (with in-flight weight updates, concurrently with generation are distributed across `NRL_REFIT_NUM_STREAMS` CUDA streams so different stages' reshards overlap. For each parameter it runs `pre` (receive-buffer allocation), calls `xferdtensor(None, ..., dst, ..., group, stream)`, then `post` (copy back into the - fused parameter). + fused parameter or load staged TRTLLM experts). After every transfer completes, the + TRTLLM path finalizes vLLM's native layerwise reload once to restore the packed runtime + layout. ### The Misc Path @@ -195,10 +202,11 @@ generation side maps those HF names onto whatever its own storage layout is. `nccl_reshard_refit()` send loop; the misc packed-broadcast producer. * **Generation side** (`vllm_backend.py`): building `hf_to_local_param_map` — mapping HF names onto vLLM's fused parameters (`qkv_proj`, `gate_up_proj`, grouped-expert - `w13_weight`/`w2_weight`) with `pre`/`post` hooks for the slice regions, which is - deliberately **shape-driven** so the same code handles generation TP and generation - EP; the comm bootstrap methods; the `nccl_reshard_refit()` receive loop; the misc - consumer feeding `load_weights`. + `w13_weight`/`w2_weight`) with `pre`/`post` hooks for slice regions or canonical + TRTLLM staging, which is deliberately **shape-driven** so the same code handles + supported generation parallelism; the comm bootstrap methods; the + `nccl_reshard_refit()` receive loop; the misc consumer feeding `load_weights`; and + backend-specific finalization after all weights arrive. **To extend to a new backend**, the only piece with genuinely new logic is `build_hf_to_local_param_map`. Everything else is boilerplate that follows a fixed @@ -207,9 +215,10 @@ contract and can be copied from the existing backend almost verbatim. **The one backend-specific implementation — `build_hf_to_local_param_map`:** resolve each bulk HF name to your local storage as a `LocalParamSpec` — `base` for tensors sent/received as-is, and `pre`/`post` hooks wherever your layout requires staging -(fused/merged tensors, layout conversions, grouped-expert stacking). This is the *only* -place your backend's parameter layout is encoded; all cross-mesh byte movement is -already handled by the shared metadata and `xferdtensor`. +(fused/merged tensors, layout conversions, grouped-expert stacking). Backends that +rebuild runtime storage may also need one transport-level finalizer after all specs have +run. These are the only places the backend's parameter layout is encoded; all cross-mesh +byte movement is already handled by the shared metadata and `xferdtensor`. (A new *training* backend additionally has to produce the HF-named metadata — names, global shapes, dtypes, and the parallelism description the agnostic builder consumes — @@ -275,4 +284,4 @@ Both transports honor the `stream` argument so the transfer is ordered with the | GB200 | DSV3 | BF16 | PP16×EP16 → TP32×DP8 | 97.6% | 1.93s-2.59s | | GB200 | Nemotron Ultra-v3 | BF16 | TP8xEP32xPP2 -> TP8xDP8 | 93.8% | 2.32s | -The feature supports both dense and MoE models. The table above shows the `XferDTensor fraction`, which is the proportion of the refit payload that utilizes the high-performance `bulk` transfer path. As the model size increases, this fraction becomes higher, which is the key to provide a scalable refit time to large models. For FP8 models, the efficiency is currently lower compared to BF16 models. \ No newline at end of file +The feature supports both dense and MoE models. The table above shows the `XferDTensor fraction`, which is the proportion of the refit payload that utilizes the high-performance `bulk` transfer path. As the model size increases, this fraction becomes higher, which is the key to provide a scalable refit time to large models. For FP8 models, the efficiency is currently lower compared to BF16 models. diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 1ec005200f9..4362b9e743b 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -730,12 +730,33 @@ def _uses_unquantized_flashinfer_trtllm(self) -> bool: return _model_uses_unquantized_flashinfer_trtllm(self.model_runner.model) - def _validate_weight_update_compatibility(self) -> None: + def _validate_weight_update_compatibility( + self, transport: WeightUpdateTransport | None = None + ) -> None: """Reject unsupported native layerwise refit combinations.""" - if ( - self._uses_unquantized_flashinfer_trtllm() - and self._mtp_drafter_refit_enabled() - ): + if not self._uses_unquantized_flashinfer_trtllm(): + return + + if transport == "nccl_reshard": + realized_placements = { + getattr( + getattr(module, "expert_map_manager", None), + "placement_strategy", + getattr(module, "expert_placement_strategy", "linear"), + ) + for module in _unquantized_flashinfer_trtllm_modules( + self.model_runner.model + ) + } + unsupported_placements = sorted(realized_placements - {"linear"}) + if unsupported_placements: + raise RuntimeError( + "Unquantized FlashInfer TRTLLM nccl_reshard refit requires " + "linear expert placement; realized " + f"{unsupported_placements!r}" + ) + + if self._mtp_drafter_refit_enabled(): raise RuntimeError( "Unquantized FlashInfer TRTLLM refit does not yet support " "a co-trained MTP drafter" @@ -751,7 +772,7 @@ def _weight_update_lifecycle( subsequent exception therefore marks this worker permanently unusable. """ if self._uses_unquantized_flashinfer_trtllm(): - self._validate_weight_update_compatibility() + self._validate_weight_update_compatibility(transport) previous_failure = self._nrl_layerwise_reload_failure if previous_failure is not None: raise RuntimeError( @@ -976,6 +997,8 @@ def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: Done once ahead of refit; the cached mapping is reused by every ``nccl_reshard_refit`` call. """ + self._validate_weight_update_compatibility("nccl_reshard") + from nemo_rl.weight_sync.nccl_reshard_utils import ( restore_refit_info_placements, ) diff --git a/nemo_rl/weight_sync/xferdtensor.py b/nemo_rl/weight_sync/xferdtensor.py index f0d81e4f7b6..78991395a22 100644 --- a/nemo_rl/weight_sync/xferdtensor.py +++ b/nemo_rl/weight_sync/xferdtensor.py @@ -213,47 +213,32 @@ def _get_mesh_coords(mesh, rank): def _compute_shard_slices(global_shape, mesh_shape, mesh_coords, placements): - """Return the slice of the global tensor this rank owns.""" + """Return the slice of the global tensor this rank owns. + + Match DTensor's sequential ``torch.chunk`` semantics, including uneven + shards, so staging buffers agree with the exact-transfer implementation. + """ slices = [slice(None) for _ in range(len(global_shape))] shard_map = {} for mesh_dim, placement in enumerate(placements): if isinstance(placement, Shard): - shard_map.setdefault(placement.dim, []).append( - (mesh_dim, mesh_shape[mesh_dim], mesh_coords[mesh_dim]) - ) - - for tensor_dim, shard_info in shard_map.items(): - shard_info.sort(key=lambda item: item[0]) - num_chunks = 1 - for _, size, _ in shard_info: - num_chunks *= size - - total_size = int(global_shape[tensor_dim]) - base = total_size // num_chunks - remainder = total_size % num_chunks - sizes = [base + 1 if i < remainder else base for i in range(num_chunks)] - - strides = [] - running = 1 - for _, size, _ in reversed(shard_info): - strides.append(running) - running *= size - strides.reverse() - - # linear_index = this rank's flat chunk number among num_chunks. - # (example: tp coord 2 * stride 1 -> linear_index = 2.) - linear_index = 0 - for (mesh_dim, size, coord), stride in zip(shard_info, strides): + shard_map.setdefault(placement.dim, []).append(mesh_dim) + + for tensor_dim, mesh_dims in shard_map.items(): + start = 0 + local_size = int(global_shape[tensor_dim]) + for mesh_dim in mesh_dims: + size = int(mesh_shape[mesh_dim]) + coord = int(mesh_coords[mesh_dim]) if coord >= size: raise ValueError(f"Invalid mesh coord {coord} for mesh dim {mesh_dim}.") - linear_index += coord * stride - - # This rank owns chunk `linear_index`; its slice starts past every - # earlier chunk. (example: start = 64+64 = 128, end = 192 -> [128:192].) - start = sum(sizes[:linear_index]) - end = start + sizes[linear_index] - slices[tensor_dim] = slice(start, end) + chunk_size = (local_size + size - 1) // size if local_size else 0 + relative_start = min(local_size, coord * chunk_size) + remaining = max(0, local_size - relative_start) + local_size = min(chunk_size, remaining) + start += relative_start + slices[tensor_dim] = slice(start, start + local_size) return slices diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index 3517436c48a..1fa2961b2bb 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -406,7 +406,32 @@ def test_build_hf_to_local_param_map_rejects_trtllm_tensor_sharding(): ext.build_hf_to_local_param_map(refit_info) -def test_nccl_reshard_lifecycle_initializes_only_trtllm_moe_modules(monkeypatch): +def test_prepare_nccl_reshard_refit_info_validates_before_building_map(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext._validate_weight_update_compatibility = MagicMock( + side_effect=RuntimeError("unsupported weight update") + ) + ext.build_hf_to_local_param_map = MagicMock() + restore_refit_info_placements = MagicMock() + monkeypatch.setattr( + "nemo_rl.weight_sync.nccl_reshard_utils.restore_refit_info_placements", + restore_refit_info_placements, + ) + + with pytest.raises(RuntimeError, match="unsupported weight update"): + ext.prepare_nccl_reshard_refit_info({"layer_names": []}) + + ext._validate_weight_update_compatibility.assert_called_once_with("nccl_reshard") + restore_refit_info_placements.assert_not_called() + ext.build_hf_to_local_param_map.assert_not_called() + assert not hasattr(ext, "nccl_reshard_refit_info") + + +def test_nccl_reshard_lifecycle_repeats_for_trtllm_moe_modules(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend model = SimpleNamespace() @@ -421,7 +446,7 @@ def test_nccl_reshard_lifecycle_initializes_only_trtllm_moe_modules(monkeypatch) ext.model_config = model_config ext.device = torch.device("cpu") ext._uses_unquantized_flashinfer_trtllm = lambda: True - ext._validate_weight_update_compatibility = lambda: None + ext._validate_weight_update_compatibility = lambda _transport=None: None ext._maybe_process_mtp_drafter_after_loading = lambda: call_order.append("mtp") monkeypatch.setattr( @@ -445,13 +470,18 @@ def test_nccl_reshard_lifecycle_initializes_only_trtllm_moe_modules(monkeypatch) ), ) - with ext._weight_update_lifecycle("nccl_reshard") as finalize: - call_order.append("transfer") - finalize() + for cycle in range(2): + with ext._weight_update_lifecycle("nccl_reshard") as finalize: + call_order.append(("transfer", cycle)) + finalize() assert call_order == [ ("initialize", trtllm_moe), - "transfer", + ("transfer", 0), + ("finalize", model, model_config), + "mtp", + ("initialize", trtllm_moe), + ("transfer", 1), ("finalize", model, model_config), "mtp", ] diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 7870bad3d68..f1663298c8f 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -121,7 +121,9 @@ def _make_mtp_refit_extension( return ext, drafter_model -def _make_unquantized_moe_model(moe_backend: str) -> SimpleNamespace: +def _make_unquantized_moe_model( + moe_backend: str, expert_placement_strategy: str = "linear" +) -> SimpleNamespace: from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, ) @@ -131,7 +133,12 @@ def _make_unquantized_moe_model(moe_backend: str) -> SimpleNamespace: quant_method = UnquantizedFusedMoEMethod.__new__(UnquantizedFusedMoEMethod) quant_method.unquantized_backend = UnquantizedMoeBackend(moe_backend) - module = SimpleNamespace(quant_method=quant_method) + module = SimpleNamespace( + quant_method=quant_method, + expert_map_manager=SimpleNamespace( + placement_strategy=expert_placement_strategy + ), + ) return SimpleNamespace(modules=lambda: [module]) @@ -237,7 +244,7 @@ def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypa ext.model_config = None ext.device = torch.device("cpu") ext._uses_unquantized_flashinfer_trtllm = lambda: True - ext._validate_weight_update_compatibility = lambda: None + ext._validate_weight_update_compatibility = lambda _transport=None: None ext._maybe_process_mtp_drafter_after_loading = MagicMock() monkeypatch.setattr( @@ -479,6 +486,51 @@ def test_unquantized_reload_rejects_cotrained_mtp_during_prepare(): assert not hasattr(ext, "state_dict_info") +@pytest.mark.vllm +def test_unquantized_reload_rejects_round_robin_expert_placement_for_nccl_only(): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM", "round_robin"), + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), + parallel_config=SimpleNamespace(expert_placement_strategy="round_robin"), + quant_config=None, + ), + ) + ext._mtp_drafter_refit_enabled = lambda: False + + ext.prepare_refit_info({"model.weight": object()}) + + with pytest.raises(RuntimeError, match="linear expert placement"): + ext.prepare_nccl_reshard_refit_info({"layer_names": []}) + + assert hasattr(ext, "state_dict_info") + assert not hasattr(ext, "nccl_reshard_refit_info") + + +@pytest.mark.vllm +def test_unquantized_reload_uses_realized_expert_placement(): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM", "linear"), + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(expert_placement_strategy="round_robin"), + quant_config=None, + ), + ) + ext._mtp_drafter_refit_enabled = lambda: False + + ext._validate_weight_update_compatibility("nccl_reshard") + + @pytest.mark.vllm def test_failed_unquantized_reload_marks_worker_unusable(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend diff --git a/tests/unit/weight_sync/test_xferdtensor_python.py b/tests/unit/weight_sync/test_xferdtensor_python.py index a1f82aea361..0234760bc49 100644 --- a/tests/unit/weight_sync/test_xferdtensor_python.py +++ b/tests/unit/weight_sync/test_xferdtensor_python.py @@ -26,6 +26,7 @@ import torch from torch.distributed._tensor import Replicate, Shard +from nemo_rl.weight_sync.xferdtensor import get_local_shard_slices from nemo_rl.weight_sync import xferdtensor_python as impl @@ -166,6 +167,49 @@ def test_compute_shard_slices_uses_uneven_torch_chunk_semantics(): ] +def test_public_local_shard_slices_matches_exact_transfer_for_uneven_shards(): + mesh = _Mesh([0, 1, 2, 3], (4,)) + + slices = [ + get_local_shard_slices((10,), mesh, (Shard(0),), rank)[0] for rank in range(4) + ] + + assert slices == [ + slice(0, 3), + slice(3, 6), + slice(6, 9), + slice(9, 10), + ] + + +def test_public_local_shard_slices_matches_repeated_shard_dimensions(): + mesh = _Mesh(list(range(6)), (2, 3)) + placements = (Shard(0), Shard(0)) + + public_slices = [ + get_local_shard_slices((10,), mesh, placements, rank)[0] for rank in range(6) + ] + exact_transfer_slices = [ + impl._compute_shard_slices( + global_shape=(10,), + mesh_shape=(2, 3), + coordinates=(rank // 3, rank % 3), + placements=placements, + )[0] + for rank in range(6) + ] + + expected = [ + slice(0, 2), + slice(2, 4), + slice(4, 5), + slice(5, 7), + slice(7, 9), + slice(9, 10), + ] + assert public_slices == exact_transfer_slices == expected + + def test_plan_geometry_tp2_to_tp1_gather(): src_regions, dst_regions, destination_groups, transfers = impl._plan_geometry( _Mesh([0, 1], (2,)), From c119bdb7ab1c7a34c5bea950c44548dae94d7038 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Mon, 17 Aug 2026 21:49:49 -0700 Subject: [PATCH 22/68] style: format TRTLLM reshard test Signed-off-by: seonjinn --- tests/unit/weight_sync/test_xferdtensor_python.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/weight_sync/test_xferdtensor_python.py b/tests/unit/weight_sync/test_xferdtensor_python.py index 0234760bc49..b6b1176e9fd 100644 --- a/tests/unit/weight_sync/test_xferdtensor_python.py +++ b/tests/unit/weight_sync/test_xferdtensor_python.py @@ -26,8 +26,8 @@ import torch from torch.distributed._tensor import Replicate, Shard -from nemo_rl.weight_sync.xferdtensor import get_local_shard_slices from nemo_rl.weight_sync import xferdtensor_python as impl +from nemo_rl.weight_sync.xferdtensor import get_local_shard_slices class _Mesh: From 10e40a5914a30f16b5bcfe9b8b226143a20f77c0 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 19 Aug 2026 17:18:55 -0700 Subject: [PATCH 23/68] fix(vllm): refresh HPC state after layerwise refit Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 11 ++++++ .../models/generation/test_vllm_backend.py | 36 +++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index cafe0e796c5..324c8ba7c80 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -92,6 +92,15 @@ def _detach_pending_layerwise_weights( arguments.arguments["loaded_weight"] = loaded_weight.clone() +def _process_hpc_modules_after_loading(model: torch.nn.Module) -> None: + """Refresh derived HPC state omitted by vLLM's layerwise finalizer.""" + from vllm.model_executor.layers.hpc import HpcModule + + for _, module in model.named_modules(): + if isinstance(module, HpcModule): + module.process_weights_after_loading(model) + + def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: """Return whether a model realized the unquantized TRTLLM MoE backend.""" # Import backend types only when inspecting a constructed vLLM model. @@ -638,6 +647,7 @@ def load_mtp_weights_from_disk(self, model_path: str) -> bool: initialize_layerwise_reload(draft_model) self._load_draft_weights(weights) finalize_layerwise_reload(draft_model, draft_model_config) + _process_hpc_modules_after_loading(draft_model) else: from vllm.model_executor.model_loader.utils import ( process_weights_after_loading, @@ -749,6 +759,7 @@ def _weight_update_lifecycle( def finalize() -> None: with torch.device(self.device): finalize_layerwise_reload(model, self.model_config) + _process_hpc_modules_after_loading(model) self._maybe_process_mtp_drafter_after_loading() torch.accelerator.synchronize() diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 47f58ac68c4..bb326868d09 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -135,6 +135,28 @@ def _make_unquantized_moe_model(moe_backend: str) -> SimpleNamespace: return SimpleNamespace(modules=lambda: [module]) +@pytest.mark.vllm +def test_process_hpc_modules_after_loading(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + class FakeHpcModule: + def __init__(self): + self.process_weights_after_loading = MagicMock() + + hpc_module = FakeHpcModule() + other_module = object() + model = SimpleNamespace( + named_modules=lambda: [("", other_module), ("rope_norm", hpc_module)] + ) + monkeypatch.setattr( + "vllm.model_executor.layers.hpc.HpcModule", FakeHpcModule + ) + + vllm_backend._process_hpc_modules_after_loading(model) + + hpc_module.process_weights_after_loading.assert_called_once_with(model) + + class _DeferredReloadLayer(torch.nn.Module): def __init__(self) -> None: super().__init__() @@ -198,6 +220,11 @@ def set_current_vllm_config(config): ("finalize", reload_model, config) ), ) + monkeypatch.setattr( + vllm_backend, + "_process_hpc_modules_after_loading", + lambda reload_model: call_order.append(("hpc", reload_model)), + ) monkeypatch.setattr( "vllm.model_executor.model_loader.utils.process_weights_after_loading", lambda *_args: pytest.fail( @@ -216,6 +243,7 @@ def set_current_vllm_config(config): ("initialize", model), "load", ("finalize", model, model_config), + ("hpc", model), "mtp", "config_exit", ] @@ -868,6 +896,8 @@ def test_load_mtp_weights_from_disk_uses_layerwise_reload_for_trtllm( tmp_path, monkeypatch ): """TRTLLM draft weights reload into their preserved runtime storage.""" + from nemo_rl.models.generation.vllm import vllm_backend + model_dir = tmp_path / "ckpt" _write_sharded_checkpoint( model_dir, @@ -893,6 +923,11 @@ def test_load_mtp_weights_from_disk_uses_layerwise_reload_for_trtllm( "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", lambda model, config: call_order.append(("finalize", model, config)), ) + monkeypatch.setattr( + vllm_backend, + "_process_hpc_modules_after_loading", + lambda model: call_order.append(("hpc", model)), + ) monkeypatch.setattr( "vllm.model_executor.model_loader.utils.process_weights_after_loading", lambda *_: pytest.fail("TRTLLM draft reload must use the layerwise path"), @@ -908,6 +943,7 @@ def test_load_mtp_weights_from_disk_uses_layerwise_reload_for_trtllm( draft_model, ext.model_runner.vllm_config.speculative_config.draft_model_config, ), + ("hpc", draft_model), ] From 9cc96b0a0dcb486d801b533e4757f5c84a7dcacf Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 26 Aug 2026 16:14:29 -0700 Subject: [PATCH 24/68] Update BF16 TRTLLM refit wording from unquantized to BF16 Signed-off-by: seonjinn --- docs/design-docs/nccl-reshard-refit.md | 4 ++-- nemo_rl/models/generation/vllm/vllm_backend.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/design-docs/nccl-reshard-refit.md b/docs/design-docs/nccl-reshard-refit.md index 4100d987f8f..8603de95e2d 100644 --- a/docs/design-docs/nccl-reshard-refit.md +++ b/docs/design-docs/nccl-reshard-refit.md @@ -39,7 +39,7 @@ single `ValueError` listing every violation. The current requirements are: `vllm_cfg.precision=fp8` and `vllm_cfg.is_mx=true`; the generation ranks quantize each received BF16 shard before installing it. Blockwise-FP8 train → MXFP8 gen is not supported. -* Unquantized BF16 FlashInfer TRTLLM MoE is supported through vLLM's native +* BF16 FlashInfer TRTLLM MoE is supported through vLLM's native layerwise-reload path. Its grouped expert weights must use expert-parallel destination sharding with linear expert placement; tensor-sharded expert destinations and round-robin placement are rejected. @@ -140,7 +140,7 @@ realized **locally**: views into a `[num_local_experts, ...]` tensor fresh at each refit. * On the **generation side**, a direct parameter's `base` is the live vLLM parameter (received into in place). Conventional fused parameters use `pre`/`post` hooks to - receive a component and copy it into the appropriate local region. Unquantized + receive a component and copy it into the appropriate local region. BF16 FlashInfer TRTLLM grouped experts instead receive into canonical EP-local staging tensors; `post` loads each logical expert with its global expert ID through vLLM's native weight loader. diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 83fc52be41e..de849f46730 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -796,7 +796,7 @@ def _validate_native_layerwise_refit( ) if strategy is None: raise RuntimeError( - "Unquantized FlashInfer TRTLLM nccl_reshard refit could " + "BF16 FlashInfer TRTLLM nccl_reshard refit could " "not determine the expert placement strategy of " f"{type(module).__name__}; refusing to assume linear " "placement" @@ -805,7 +805,7 @@ def _validate_native_layerwise_refit( unsupported_placements = sorted(realized_placements - {"linear"}) if unsupported_placements: raise RuntimeError( - "Unquantized FlashInfer TRTLLM nccl_reshard refit requires " + "BF16 FlashInfer TRTLLM nccl_reshard refit requires " "linear expert placement; realized " f"{unsupported_placements!r}" ) @@ -1124,7 +1124,7 @@ def _trtllm_grouped_expert_spec( ] if unsupported_shards: raise ValueError( - "Unquantized FlashInfer TRTLLM nccl_reshard refit requires " + "BF16 FlashInfer TRTLLM nccl_reshard refit requires " "expert-parallel destination shards; unsupported tensor shard " f"dimensions {unsupported_shards} for {param_info['name']!r}" ) @@ -1149,7 +1149,7 @@ def _trtllm_grouped_expert_spec( # instead of silently loading experts onto the wrong ranks. if ep_size > 0 and num_global_experts % ep_size != 0: raise ValueError( - "Unquantized FlashInfer TRTLLM nccl_reshard refit requires " + "BF16 FlashInfer TRTLLM nccl_reshard refit requires " "the global expert count to divide evenly across EP ranks; " f"got {num_global_experts} experts over {ep_size} ranks for " f"{param_info['name']!r}" @@ -1173,7 +1173,7 @@ def _trtllm_grouped_expert_spec( dtype = _STR_TO_DTYPE.get(str(dtype_value)) if dtype is None: raise ValueError( - "Unquantized FlashInfer TRTLLM nccl_reshard refit got an " + "BF16 FlashInfer TRTLLM nccl_reshard refit got an " f"unsupported wire dtype {dtype_value!r} for " f"{param_info['name']!r}" ) @@ -1200,7 +1200,7 @@ def post(ctx: RefitCtx) -> None: missing = [name for name, _ in weights if name not in loaded_names] if missing: raise RuntimeError( - "Unquantized FlashInfer TRTLLM nccl_reshard refit " + "BF16 FlashInfer TRTLLM nccl_reshard refit " f"failed to load staged expert weights {missing!r}" ) From 88ff1ab7bc415cf00275262d618699eddd96cea0 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 26 Aug 2026 19:17:07 -0700 Subject: [PATCH 25/68] Fix nccl_reshard backend unit test fixtures for refit-info validation Signed-off-by: seonjinn --- .../generation/test_nccl_reshard_backend.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index 64f69f3f21e..594dcac7cbb 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -255,25 +255,34 @@ def test_build_hf_to_local_param_map_specs_and_roundtrip(): { "name": "model.layers.0.mlp.gate_proj.weight", "global_shape": [256, H], + "dtype": "torch.float32", + }, + { + "name": "model.layers.0.mlp.up_proj.weight", + "global_shape": [256, H], + "dtype": "torch.float32", }, - {"name": "model.layers.0.mlp.up_proj.weight", "global_shape": [256, H]}, { "name": "model.layers.0.mlp.down_proj.weight", "global_shape": [H, 256], + "dtype": "torch.float32", }, { "name": "model.layers.0.mlp.experts.gate_proj.weight", "global_shape": [E, 128, H], + "dtype": "torch.float32", "grouped_expert_proj": "gate_proj", }, { "name": "model.layers.0.mlp.experts.up_proj.weight", "global_shape": [E, 128, H], + "dtype": "torch.float32", "grouped_expert_proj": "up_proj", }, { "name": "model.layers.0.mlp.experts.down_proj.weight", "global_shape": [E, H, 128], + "dtype": "torch.float32", "grouped_expert_proj": "down_proj", }, ] @@ -357,7 +366,9 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): ext.device = torch.device("cpu") ext.pp_comm_groups = {0: SimpleNamespace(rank=9)} ext._uses_unquantized_flashinfer_trtllm = lambda: True - ext._load_full_hf_weights = MagicMock() + ext._load_full_hf_weights = MagicMock( + side_effect=lambda weights: [name for name, _ in weights] + ) spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) assert spec is not None and spec.pre is not None and spec.post is not None From ddf7495d48715205cb14b1ab0d7854618fe2cbf8 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 2 Sep 2026 01:28:02 -0700 Subject: [PATCH 26/68] fix(vllm): reuse exact shard calculation for refit Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 22 ++++-- nemo_rl/weight_sync/xferdtensor.py | 79 ++++++++----------- .../weight_sync/test_xferdtensor_python.py | 44 ----------- 3 files changed, 47 insertions(+), 98 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 7c810a77057..527ab792cc3 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -147,14 +147,24 @@ def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: 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 import get_local_shard_slices + from nemo_rl.weight_sync.xferdtensor_python import _compute_shard_slices dst_mesh = param_info["dst_mesh_info"] - return get_local_shard_slices( - param_info["global_shape"], - dst_mesh, - param_info["dst_placements"], - rank, + mesh_tensor = getattr(dst_mesh, "mesh", None) + if mesh_tensor is None: + mesh_tensor = getattr(dst_mesh, "_mesh", None) + if mesh_tensor is None: + raise ValueError("Destination DeviceMesh does not expose mesh ranks.") + coordinates = (mesh_tensor == rank).nonzero(as_tuple=False) + if coordinates.numel() == 0: + raise ValueError(f"Rank {rank} is absent from the destination mesh") + return tuple( + _compute_shard_slices( + param_info["global_shape"], + list(mesh_tensor.shape), + coordinates[0].tolist(), + param_info["dst_placements"], + ) ) diff --git a/nemo_rl/weight_sync/xferdtensor.py b/nemo_rl/weight_sync/xferdtensor.py index 78991395a22..3f1841e9c52 100644 --- a/nemo_rl/weight_sync/xferdtensor.py +++ b/nemo_rl/weight_sync/xferdtensor.py @@ -16,13 +16,10 @@ import logging import os -from collections.abc import Sequence from contextlib import nullcontext -from typing import Any import torch from torch.distributed._tensor import Shard -from torch.distributed.tensor.placement_types import Placement try: from nccl.m2n import ( # pyrefly: ignore[import-error] @@ -213,60 +210,46 @@ def _get_mesh_coords(mesh, rank): def _compute_shard_slices(global_shape, mesh_shape, mesh_coords, placements): - """Return the slice of the global tensor this rank owns. - - Match DTensor's sequential ``torch.chunk`` semantics, including uneven - shards, so staging buffers agree with the exact-transfer implementation. - """ + """Return the slice of the global tensor this rank owns.""" slices = [slice(None) for _ in range(len(global_shape))] shard_map = {} for mesh_dim, placement in enumerate(placements): if isinstance(placement, Shard): - shard_map.setdefault(placement.dim, []).append(mesh_dim) - - for tensor_dim, mesh_dims in shard_map.items(): - start = 0 - local_size = int(global_shape[tensor_dim]) - for mesh_dim in mesh_dims: - size = int(mesh_shape[mesh_dim]) - coord = int(mesh_coords[mesh_dim]) + shard_map.setdefault(placement.dim, []).append( + (mesh_dim, mesh_shape[mesh_dim], mesh_coords[mesh_dim]) + ) + + for tensor_dim, shard_info in shard_map.items(): + shard_info.sort(key=lambda item: item[0]) + num_chunks = 1 + for _, size, _ in shard_info: + num_chunks *= size + + total_size = int(global_shape[tensor_dim]) + base = total_size // num_chunks + remainder = total_size % num_chunks + sizes = [base + 1 if i < remainder else base for i in range(num_chunks)] + + strides = [] + running = 1 + for _, size, _ in reversed(shard_info): + strides.append(running) + running *= size + strides.reverse() + + # linear_index = this rank's flat chunk number among num_chunks. + linear_index = 0 + for (mesh_dim, size, coord), stride in zip(shard_info, strides): if coord >= size: raise ValueError(f"Invalid mesh coord {coord} for mesh dim {mesh_dim}.") - chunk_size = (local_size + size - 1) // size if local_size else 0 - relative_start = min(local_size, coord * chunk_size) - remaining = max(0, local_size - relative_start) - local_size = min(chunk_size, remaining) - start += relative_start - slices[tensor_dim] = slice(start, start + local_size) - - return slices + linear_index += coord * stride + start = sum(sizes[:linear_index]) + end = start + sizes[linear_index] + slices[tensor_dim] = slice(start, end) -def get_local_shard_slices( - global_shape: list[int] | tuple[int, ...] | torch.Size, - mesh: Any, - placements: Sequence[Placement], - rank: int, -) -> tuple[slice, ...]: - """Return one rank's local slices in a global tensor.""" - mesh_coords = _get_mesh_coords(mesh, rank) - if mesh_coords is None: - raise ValueError(f"Rank {rank} is absent from the destination mesh") - mesh_tensor = getattr(mesh, "mesh", None) - if mesh_tensor is None: - mesh_tensor = getattr(mesh, "_mesh", None) - if mesh_tensor is None: - raise ValueError("DeviceMesh does not expose mesh ranks.") - return tuple( - _compute_shard_slices( - global_shape, - list(mesh_tensor.shape), - mesh_coords, - placements, - ) - ) - + return slices def xferdtensor_golden( src_tensor, diff --git a/tests/unit/weight_sync/test_xferdtensor_python.py b/tests/unit/weight_sync/test_xferdtensor_python.py index b6b1176e9fd..a1f82aea361 100644 --- a/tests/unit/weight_sync/test_xferdtensor_python.py +++ b/tests/unit/weight_sync/test_xferdtensor_python.py @@ -27,7 +27,6 @@ from torch.distributed._tensor import Replicate, Shard from nemo_rl.weight_sync import xferdtensor_python as impl -from nemo_rl.weight_sync.xferdtensor import get_local_shard_slices class _Mesh: @@ -167,49 +166,6 @@ def test_compute_shard_slices_uses_uneven_torch_chunk_semantics(): ] -def test_public_local_shard_slices_matches_exact_transfer_for_uneven_shards(): - mesh = _Mesh([0, 1, 2, 3], (4,)) - - slices = [ - get_local_shard_slices((10,), mesh, (Shard(0),), rank)[0] for rank in range(4) - ] - - assert slices == [ - slice(0, 3), - slice(3, 6), - slice(6, 9), - slice(9, 10), - ] - - -def test_public_local_shard_slices_matches_repeated_shard_dimensions(): - mesh = _Mesh(list(range(6)), (2, 3)) - placements = (Shard(0), Shard(0)) - - public_slices = [ - get_local_shard_slices((10,), mesh, placements, rank)[0] for rank in range(6) - ] - exact_transfer_slices = [ - impl._compute_shard_slices( - global_shape=(10,), - mesh_shape=(2, 3), - coordinates=(rank // 3, rank % 3), - placements=placements, - )[0] - for rank in range(6) - ] - - expected = [ - slice(0, 2), - slice(2, 4), - slice(4, 5), - slice(5, 7), - slice(7, 9), - slice(9, 10), - ] - assert public_slices == exact_transfer_slices == expected - - def test_plan_geometry_tp2_to_tp1_gather(): src_regions, dst_regions, destination_groups, transfers = impl._plan_geometry( _Mesh([0, 1], (2,)), From d0bcb961cd6a46adfb07089d48da22a3d4bd290f Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 2 Sep 2026 22:36:51 -0700 Subject: [PATCH 27/68] fix(refit): cover Qwen3.5 TRTLLM NCCL reshard Signed-off-by: seonjinn --- docs/guides/models/qwen/qwen3-5.md | 1 + ....5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml | 44 +++++++++ .../models/generation/vllm/vllm_backend.py | 47 ++++++--- tests/test_suites/disabled.txt | 5 + ...n3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh | 39 ++++++++ .../generation/test_nccl_reshard_backend.py | 98 ++++++++++++++++++- .../test_qwen35_bf16_trtllm_recipe.py | 72 ++++++++++++++ .../weight_sync/test_nccl_reshard_utils.py | 39 ++++++++ 8 files changed, 332 insertions(+), 13 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml create mode 100755 tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh create mode 100644 tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py diff --git a/docs/guides/models/qwen/qwen3-5.md b/docs/guides/models/qwen/qwen3-5.md index fb6668dbcdb..8b2896aa34c 100644 --- a/docs/guides/models/qwen/qwen3-5.md +++ b/docs/guides/models/qwen/qwen3-5.md @@ -64,6 +64,7 @@ authoritative settings. | Qwen3.5-9B-Base | LLM | GRPO | Megatron | 1n8g | [`grpo-qwen3.5-9b-1n8g-megatron-fp8.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-9b-1n8g-megatron-fp8.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2-fp8.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2-fp8.yaml) | +| Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 6n4g | [`grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | AutoModel | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml) | | Qwen3.5-35B-A3B-Base | LLM | GRPO | AutoModel | 4n8g | [`grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml) | | Qwen3.5-397B-A17B | LLM | GRPO | Megatron | 32n8g | [`grpo-qwen3.5-397ba17b-32n8g-megatron.v2.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.yaml) | diff --git a/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml new file mode 100644 index 00000000000..674faceff72 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml @@ -0,0 +1,44 @@ +defaults: grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml + +loss_fn: + use_importance_sampling_correction: true + truncated_importance_sampling_type: tis + truncated_importance_sampling_ratio: 2 + +grpo: + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + +checkpointing: + checkpoint_dir: results/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm + +policy: + generation: + refit_transport: nccl_reshard + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 2 + vllm_cfg: + async_engine: true + precision: bfloat16 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + expert_parallel_size: 4 + gpu_memory_utilization: 0.8 + enforce_eager: false + vllm_kwargs: + moe_backend: flashinfer_trtllm + expert_placement_strategy: linear + +logger: + wandb: + name: grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm + +cluster: + gpus_per_node: 4 + num_nodes: 6 + segment_size: 2 diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 527ab792cc3..ed59d82c1d5 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -861,6 +861,13 @@ def _validate_native_layerwise_refit( return if transport == "nccl_reshard": + if self._uses_fp8_kv_cache(): + raise RuntimeError( + "BF16 FlashInfer TRTLLM nccl_reshard refit does not " + "support an FP8 KV cache because its static scales are " + "outside the targeted MoE reload lifecycle" + ) + realized_placements = set() for module in _unquantized_flashinfer_trtllm_modules( self.model_runner.model @@ -1215,6 +1222,7 @@ def post(ctx: RefitCtx) -> None: def _trtllm_grouped_expert_spec( param_info: dict[str, Any], + vllm_param: torch.Tensor, ) -> LocalParamSpec: from torch.distributed._tensor import Shard @@ -1257,7 +1265,18 @@ def _trtllm_grouped_expert_spec( ) pp_stage = param_info.get("pp_stage", 0) - rank = self.pp_comm_groups[pp_stage].rank + pp_comm_groups = self.pp_comm_groups + if pp_comm_groups is None: + raise RuntimeError( + "BF16 FlashInfer TRTLLM nccl_reshard refit mapping was built " + "before the per-PP-stage groups were initialized" + ) + if pp_stage not in pp_comm_groups: + raise RuntimeError( + "BF16 FlashInfer TRTLLM nccl_reshard refit has no " + f"communicator for PP stage {pp_stage}" + ) + rank = pp_comm_groups[pp_stage].rank local_slices = _local_shard_slices(param_info, rank) local_shape = tuple( global_size @@ -1270,12 +1289,16 @@ def _trtllm_grouped_expert_spec( expert_start = 0 if local_slices[0].start is None else local_slices[0].start grouped_proj = param_info["grouped_expert_proj"] expert_prefix = param_info["name"].rsplit(f".{grouped_proj}.weight", 1)[0] - fused_param = ( - "w13_weight" - if grouped_proj in ("gate_proj", "up_proj") - else "w2_weight" - ) - expected_loaded_name = f"{expert_prefix}.{fused_param}" + registered_vllm_name = vllm_names_by_id.get(id(vllm_param)) + if registered_vllm_name is None: + raise ValueError( + "BF16 FlashInfer TRTLLM nccl_reshard refit resolved an " + f"unregistered vLLM parameter for {param_info['name']!r}" + ) + expected_loaded_names = { + registered_vllm_name, + registered_vllm_name.replace(".routed_experts.", "."), + } dtype_value = param_info.get("dtype") dtype = _STR_TO_DTYPE.get(str(dtype_value)) if dtype is None: @@ -1302,13 +1325,13 @@ def post(ctx: RefitCtx) -> None: loaded_names = self._load_full_hf_weights(weights) # AutoWeightsLoader reports the fused destination parameter, # not each per-expert HF source name. - if ( - loaded_names is not None - and expected_loaded_name not in loaded_names + if loaded_names is not None and expected_loaded_names.isdisjoint( + loaded_names ): raise RuntimeError( "BF16 FlashInfer TRTLLM nccl_reshard refit failed to " - f"load fused expert destination {expected_loaded_name!r}; " + "load fused expert destination; expected one of " + f"{sorted(expected_loaded_names)!r}, " f"vLLM reported {sorted(loaded_names)!r}" ) @@ -1360,7 +1383,7 @@ def post(ctx: RefitCtx) -> None: for hf_name, (vllm_param, merged_slice) in vllm_param_map_and_slices.items(): param_info = param_info_by_name[hf_name] if use_trtllm_staging and param_info.get("grouped_expert_proj"): - specs[hf_name] = _trtllm_grouped_expert_spec(param_info) + specs[hf_name] = _trtllm_grouped_expert_spec(param_info, vllm_param) continue wire_dtype_value = param_info.get("dtype") diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index c06a9f253b9..bd5cee279aa 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -43,3 +43,8 @@ tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh # nightly budget is at 3922 of its 3928 GPU-hour cap and this run's ~27 # GPU-hours would exceed it. Move to nightly.txt when the budget has room. tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-ready-first-single-controller.sh + +# Qwen3.5 BF16 FlashInfer TRTLLM + NCCL reshard validation. Keep this +# 6x4-GPU, 20-step run manual until an end-to-end run establishes its +# gen_kl_error/reward bounds and the recurring GB200 suite has budget. +tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh diff --git a/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh new file mode 100755 index 00000000000..41627dd1cf0 --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh @@ -0,0 +1,39 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=6 +GPUS_PER_NODE=4 +STEPS_PER_RUN=20 +MAX_STEPS=20 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) +NUM_MINUTES=240 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" +uv run examples/run_grpo.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir="$LOG_DIR" \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name="$EXP_NAME" \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir="$CKPT_DIR" \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' "$JSON_METRICS") -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py "$JSON_METRICS" \ + 'mean(data["train/gen_kl_error"]) < 0.002' \ + 'max(data["train/reward"]) > 0.5' + + rm -rf "$CKPT_DIR" +fi diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index 71413ac71d1..e6d8f8d9bf6 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -367,7 +367,7 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): ext.pp_comm_groups = {0: SimpleNamespace(rank=9)} ext._uses_unquantized_flashinfer_trtllm = lambda: True ext._load_full_hf_weights = MagicMock( - return_value={"model.layers.0.mlp.experts.w13_weight"} + return_value={"model.layers.0.mlp.experts.routed_experts.w13_weight"} ) spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) @@ -394,6 +394,81 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): torch.testing.assert_close(packed_w13, torch.full_like(packed_w13, 7.0)) +def test_build_hf_to_local_param_map_stages_qwen35_wrapped_experts(): + """Qwen3.5 wrapper prefixes and RoutedExperts names survive staged reload.""" + hidden_size, num_experts, intermediate_size = 16, 4, 32 + hf_prefix = "model.language_model.layers.0.mlp.experts" + runtime_prefix = "language_model.model.layers.0.mlp.experts" + expert_name = f"{hf_prefix}.down_proj.weight" + runtime_name = f"{runtime_prefix}.routed_experts.w2_weight" + refit_info = { + "gen_tp_size": 2, + "layer_names": ["model.language_model.layers.0"], + "per_layer_params": { + "model.language_model.layers.0": [ + { + "name": expert_name, + "global_shape": [ + num_experts, + hidden_size, + intermediate_size, + ], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "down_proj", + "dst_mesh_info": MeshInfo(torch.tensor([8, 9])), + "dst_placements": [Shard(0)], + } + ] + }, + } + packed_w2 = torch.full((128, 16, 24, 64), 7.0) + ext = _make_ext({runtime_name: packed_w2}) + ext.device = torch.device("cpu") + ext.pp_comm_groups = {0: SimpleNamespace(rank=9)} + ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._load_full_hf_weights = MagicMock(return_value={f"{runtime_prefix}.w2_weight"}) + + spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) + assert spec is not None and spec.pre is not None and spec.post is not None + spec.post(spec.pre(spec.base)) + + loaded_weights = ext._load_full_hf_weights.call_args.args[0] + assert [name for name, _ in loaded_weights] == [ + f"{hf_prefix}.2.down_proj.weight", + f"{hf_prefix}.3.down_proj.weight", + ] + + +def test_build_hf_to_local_param_map_requires_pp_groups_for_trtllm_staging(): + """A missing NCCL communicator fails before staged expert placement.""" + expert_name = "model.layers.0.mlp.experts.down_proj.weight" + refit_info = { + "gen_tp_size": 1, + "layer_names": ["model.layers.0"], + "per_layer_params": { + "model.layers.0": [ + { + "name": expert_name, + "global_shape": [2, 16, 32], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "down_proj", + "dst_mesh_info": MeshInfo(torch.tensor([0])), + "dst_placements": [Shard(0)], + } + ] + }, + } + ext = _make_ext( + {"model.layers.0.mlp.experts.routed_experts.w2_weight": torch.empty(2, 16, 32)} + ) + ext.device = torch.device("cpu") + ext.pp_comm_groups = None + ext._uses_unquantized_flashinfer_trtllm = lambda: True + + with pytest.raises(RuntimeError, match="before.*per-PP-stage groups"): + ext.build_hf_to_local_param_map(refit_info) + + def test_build_hf_to_local_param_map_rejects_missing_trtllm_destination(): """The staged load must report the fused destination parameter.""" hidden_size, num_experts, intermediate_size = 16, 4, 32 @@ -496,6 +571,25 @@ def test_prepare_nccl_reshard_refit_info_validates_before_building_map(monkeypat assert not hasattr(ext, "nccl_reshard_refit_info") +def test_nccl_reshard_trtllm_refit_rejects_fp8_kv_cache(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=object()) + ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._uses_fp8_kv_cache = lambda: True + monkeypatch.setattr( + vllm_backend, + "_unquantized_flashinfer_trtllm_modules", + lambda _model: [SimpleNamespace(expert_placement_strategy="linear")], + ) + + with pytest.raises(RuntimeError, match="FP8 KV cache"): + ext._validate_native_layerwise_refit("nccl_reshard") + + def test_legacy_refit_map_is_built_after_comm_groups_exist(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend @@ -524,6 +618,7 @@ def init_nccl_communicator(self, *, device): _FakeGroup, ) monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) ext.init_nccl_reshard_comm_group( @@ -616,6 +711,7 @@ def test_nccl_reshard_refit_runs_transport_lifecycle(monkeypatch): "per_layer_params": {}, "misc_meta": {}, } + ext.pp_comm_groups = {} ext._receive_and_load_misc_params = MagicMock() ext._maybe_process_fp8_kv_cache = MagicMock() finalize = MagicMock() diff --git a/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py b/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py new file mode 100644 index 00000000000..f4483cbbe61 --- /dev/null +++ b/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py @@ -0,0 +1,72 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from typing import Any + +from omegaconf import OmegaConf + +from nemo_rl.utils.config import ( + load_config_with_inheritance, + register_omegaconf_resolvers, +) + + +PROJECT_ROOT = Path(__file__).resolve().parents[4] +RECIPE_NAME = "grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml" + + +def _load_recipe() -> dict[str, Any]: + register_omegaconf_resolvers() + recipe_path = PROJECT_ROOT / "examples/configs/recipes/llm" / RECIPE_NAME + recipe = OmegaConf.to_container( + load_config_with_inheritance(recipe_path), resolve=True + ) + assert isinstance(recipe, dict) + return recipe + + +def test_qwen35_bf16_trtllm_recipe_uses_nccl_reshard() -> None: + recipe = _load_recipe() + generation = recipe["policy"]["generation"] + + assert recipe["loss_fn"]["use_importance_sampling_correction"] is True + assert recipe["cluster"]["gpus_per_node"] == 4 + assert recipe["cluster"]["num_nodes"] == 6 + assert recipe["cluster"]["segment_size"] == 2 + async_grpo = recipe["grpo"]["async_grpo"] + assert async_grpo["enabled"] is True + assert async_grpo["max_trajectory_age_steps"] == 1 + assert async_grpo["in_flight_weight_updates"] is True + assert generation["refit_transport"] == "nccl_reshard" + assert generation["colocated"] == { + "enabled": False, + "resources": {"gpus_per_node": 4, "num_nodes": 2}, + } + + +def test_qwen35_bf16_trtllm_recipe_uses_supported_expert_layout() -> None: + recipe = _load_recipe() + generation = recipe["policy"]["generation"] + vllm_cfg = generation["vllm_cfg"] + + assert vllm_cfg["precision"] == "bfloat16" + assert vllm_cfg["tensor_parallel_size"] == 4 + assert vllm_cfg["expert_parallel_size"] == 4 + assert vllm_cfg["pipeline_parallel_size"] == 1 + assert vllm_cfg["enforce_eager"] is False + assert generation["vllm_kwargs"] == { + "moe_backend": "flashinfer_trtllm", + "expert_placement_strategy": "linear", + } diff --git a/tests/unit/weight_sync/test_nccl_reshard_utils.py b/tests/unit/weight_sync/test_nccl_reshard_utils.py index dfaf985ea4c..daf7f8a3fd8 100644 --- a/tests/unit/weight_sync/test_nccl_reshard_utils.py +++ b/tests/unit/weight_sync/test_nccl_reshard_utils.py @@ -403,6 +403,45 @@ def test_group_expert_params_collapses_to_grouped_hf_entries(): ) +def test_group_expert_params_canonicalizes_qwen35_grouped_slabs(): + base = "model.language_model.layers.0.mlp.experts" + metadata = { + f"{base}.gate_up_proj": { + "shape": [256, 1024, 2048], + "dtype": "torch.bfloat16", + }, + f"{base}.down_proj": { + "shape": [256, 2048, 512], + "dtype": "torch.bfloat16", + }, + "model.visual.proj.weight": { + "shape": [2048, 2048], + "dtype": "torch.bfloat16", + }, + } + + grouped = group_expert_params_in_metadata(metadata) + + assert grouped[f"{base}.gate_proj.weight"] == { + "shape": [256, 512, 2048], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "gate_proj", + } + assert grouped[f"{base}.up_proj.weight"] == { + "shape": [256, 512, 2048], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "up_proj", + } + assert grouped[f"{base}.down_proj.weight"] == { + "shape": [256, 2048, 512], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "down_proj", + } + assert f"{base}.gate_up_proj" not in grouped + assert f"{base}.down_proj" not in grouped + assert grouped["model.visual.proj.weight"] == metadata["model.visual.proj.weight"] + + def test_group_expert_params_no_experts_is_identity(): md = { "model.layers.0.self_attn.q_proj.weight": { From 9a5d850d9f004262262232b33d254bf2f0791de4 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 2 Sep 2026 22:47:32 -0700 Subject: [PATCH 28/68] docs(refit): clarify FP8 KV cache restriction Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/vllm_backend.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index ed59d82c1d5..5d0a630beaa 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -1739,8 +1739,9 @@ def _recv_one_param(param_info, group, stream): # Finalize post-load weight processing: dense Linear + attention/MLA, # the per-MoE-backend w13 layout (FlashInfer CUTLASS/TRTLLM) that the # canonical [gate; up] bulk write above defers to here, and the MTP - # drafter's mirror of the same. The FP8 KV-cache per-layer k/v scales - # are finalized by the lifecycle on exit. + # drafter's mirror of the same. The BF16 TRTLLM nccl_reshard path + # rejects FP8 KV cache above because its static scales are outside this + # targeted MoE lifecycle. finalize() torch.cuda.empty_cache() From 621641de96fe80249a4fc0782cec0a44869490f5 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 2 Sep 2026 22:49:18 -0700 Subject: [PATCH 29/68] refactor(refit): drop unrelated shard comments Signed-off-by: seonjinn --- nemo_rl/weight_sync/xferdtensor.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nemo_rl/weight_sync/xferdtensor.py b/nemo_rl/weight_sync/xferdtensor.py index 3f1841e9c52..6f003d58bc4 100644 --- a/nemo_rl/weight_sync/xferdtensor.py +++ b/nemo_rl/weight_sync/xferdtensor.py @@ -239,18 +239,22 @@ def _compute_shard_slices(global_shape, mesh_shape, mesh_coords, placements): strides.reverse() # linear_index = this rank's flat chunk number among num_chunks. + # (example: tp coord 2 * stride 1 -> linear_index = 2.) linear_index = 0 for (mesh_dim, size, coord), stride in zip(shard_info, strides): if coord >= size: raise ValueError(f"Invalid mesh coord {coord} for mesh dim {mesh_dim}.") linear_index += coord * stride + # This rank owns chunk `linear_index`; its slice starts past every + # earlier chunk. (example: start = 64+64 = 128, end = 192 -> [128:192].) start = sum(sizes[:linear_index]) end = start + sizes[linear_index] slices[tensor_dim] = slice(start, end) return slices + def xferdtensor_golden( src_tensor, src_mesh, From 4e987b5d0e928a36865e11d6701ab1ea6fef8a2e Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 16:55:33 -0700 Subject: [PATCH 30/68] test(refit): cover padded Lightning expert staging Signed-off-by: seonjinn --- .../generation/test_nccl_reshard_backend.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index e6d8f8d9bf6..a4e4c87226e 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -394,6 +394,53 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): torch.testing.assert_close(packed_w13, torch.full_like(packed_w13, 7.0)) +def test_build_hf_to_local_param_map_stages_nemotron_lightning_padded_experts(): + """Receive the logical Nano/Lightning weight instead of its padded runtime form.""" + num_experts, intermediate_size, hidden_size = 128, 928, 2688 + expert_name = "backbone.layers.0.mlp.experts.gate_proj.weight" + runtime_name = "model.layers.0.mlp.experts.routed_experts.w13_weight" + refit_info = { + "gen_tp_size": 1, + "layer_names": ["backbone.layers.0"], + "per_layer_params": { + "backbone.layers.0": [ + { + "name": expert_name, + "global_shape": [ + num_experts, + intermediate_size, + hidden_size, + ], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "gate_proj", + "dst_mesh_info": MeshInfo(torch.tensor([0])), + "dst_placements": [Shard(0)], + } + ] + }, + } + packed_runtime = torch.empty( + num_experts, + hidden_size // 64, + 1024, + 64, + dtype=torch.bfloat16, + device="meta", + ) + ext = _make_ext({runtime_name: packed_runtime}) + ext.device = torch.device("meta") + ext.pp_comm_groups = {0: SimpleNamespace(rank=0)} + ext._uses_unquantized_flashinfer_trtllm = lambda: True + + spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) + assert spec is not None and spec.pre is not None and spec.post is not None + + ctx = spec.pre(spec.base) + assert ctx.buf.shape == (num_experts, intermediate_size, hidden_size) + assert ctx.buf.dtype == torch.bfloat16 + assert ctx.buf.numel() != packed_runtime.numel() + + def test_build_hf_to_local_param_map_stages_qwen35_wrapped_experts(): """Qwen3.5 wrapper prefixes and RoutedExperts names survive staged reload.""" hidden_size, num_experts, intermediate_size = 16, 4, 32 From 7b3f79fa0a56d6f3fded7e3656a92b871d8a5216 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 17:24:39 -0700 Subject: [PATCH 31/68] test(refit): exercise Lightning native loader input Signed-off-by: seonjinn --- .../models/generation/test_nccl_reshard_backend.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index a4e4c87226e..37a3ad183f3 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -397,7 +397,7 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): def test_build_hf_to_local_param_map_stages_nemotron_lightning_padded_experts(): """Receive the logical Nano/Lightning weight instead of its padded runtime form.""" num_experts, intermediate_size, hidden_size = 128, 928, 2688 - expert_name = "backbone.layers.0.mlp.experts.gate_proj.weight" + expert_name = "backbone.layers.0.mlp.experts.up_proj.weight" runtime_name = "model.layers.0.mlp.experts.routed_experts.w13_weight" refit_info = { "gen_tp_size": 1, @@ -412,7 +412,7 @@ def test_build_hf_to_local_param_map_stages_nemotron_lightning_padded_experts(): hidden_size, ], "dtype": "torch.bfloat16", - "grouped_expert_proj": "gate_proj", + "grouped_expert_proj": "up_proj", "dst_mesh_info": MeshInfo(torch.tensor([0])), "dst_placements": [Shard(0)], } @@ -431,6 +431,7 @@ def test_build_hf_to_local_param_map_stages_nemotron_lightning_padded_experts(): ext.device = torch.device("meta") ext.pp_comm_groups = {0: SimpleNamespace(rank=0)} ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._load_full_hf_weights = MagicMock(return_value={runtime_name}) spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) assert spec is not None and spec.pre is not None and spec.post is not None @@ -439,6 +440,13 @@ def test_build_hf_to_local_param_map_stages_nemotron_lightning_padded_experts(): assert ctx.buf.shape == (num_experts, intermediate_size, hidden_size) assert ctx.buf.dtype == torch.bfloat16 assert ctx.buf.numel() != packed_runtime.numel() + spec.post(ctx) + + loaded_weights = ext._load_full_hf_weights.call_args.args[0] + assert len(loaded_weights) == num_experts + assert loaded_weights[0][0] == "backbone.layers.0.mlp.experts.0.up_proj.weight" + assert loaded_weights[-1][0] == "backbone.layers.0.mlp.experts.127.up_proj.weight" + assert loaded_weights[0][1].shape == (intermediate_size, hidden_size) def test_build_hf_to_local_param_map_stages_qwen35_wrapped_experts(): From 748fc71f82586041d102ac4efc1e1311ca095a5f Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 22:26:06 -0700 Subject: [PATCH 32/68] fix(vllm): dispatch mixed TRTLLM refit by module Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 47 ++++- .../generation/test_nccl_reshard_backend.py | 150 +++++++++++++++- .../models/generation/test_vllm_backend.py | 136 ++++++++++++++- .../generation/test_vllm_fp8_quantization.py | 161 ++++++++++++++++++ 4 files changed, 476 insertions(+), 18 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 5d0a630beaa..422f9de507d 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -140,6 +140,23 @@ def _unquantized_flashinfer_trtllm_modules( ] +def _process_mxfp8_modules_after_native_reload(model: torch.nn.Module) -> None: + """Rebuild MXFP8 runtime layouts skipped by the partial native reload.""" + try: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptMxFp8FusedMoE, + ModelOptMxFp8LinearMethod, + ) + except ImportError: + return + + mxfp8_methods = (ModelOptMxFp8FusedMoE, ModelOptMxFp8LinearMethod) + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance(quant_method, mxfp8_methods): + quant_method.process_weights_after_loading(module) + + def _model_uses_unquantized_flashinfer_trtllm(model: torch.nn.Module) -> bool: """Return whether a model realized the unquantized TRTLLM MoE backend.""" return bool(_unquantized_flashinfer_trtllm_modules(model)) @@ -839,13 +856,26 @@ def _uses_unquantized_flashinfer_trtllm(self) -> bool: if not self._supports_unquantized_flashinfer_trtllm_refit(): return False model_runner = getattr(self, "model_runner", None) - vllm_config = getattr(model_runner, "vllm_config", None) - if vllm_config is None: - return False - if getattr(vllm_config, "quant_config", None) is not None: + model = getattr(model_runner, "model", None) + if model is None: return False - return _model_uses_unquantized_flashinfer_trtllm(self.model_runner.model) + return _model_uses_unquantized_flashinfer_trtllm(model) + + def _unquantized_flashinfer_trtllm_param_ids(self) -> set[int]: + """Return parameters owned by realized unquantized TRTLLM MoE modules.""" + if not self._supports_unquantized_flashinfer_trtllm_refit(): + return set() + model_runner = getattr(self, "model_runner", None) + model = getattr(model_runner, "model", None) + if model is None: + return set() + + return { + id(param) + for module in _unquantized_flashinfer_trtllm_modules(model) + for param in module.parameters(recurse=False) + } def _uses_native_layerwise_refit(self, transport: WeightUpdateTransport) -> bool: """Return whether this transport needs vLLM's layerwise lifecycle.""" @@ -949,6 +979,7 @@ def _weight_update_lifecycle( def finalize() -> None: with torch.device(self.device): finalize_layerwise_reload(model, self.model_config) + _process_mxfp8_modules_after_native_reload(model) _refresh_hpc_modules_after_layerwise_reload(model) self._maybe_process_mtp_drafter_after_loading() torch.cuda.synchronize() @@ -1378,11 +1409,13 @@ def post(ctx: RefitCtx) -> None: } vllm_params = dict(self.model_runner.model.named_parameters()) vllm_names_by_id = {id(param): name for name, param in vllm_params.items()} - use_trtllm_staging = self._uses_unquantized_flashinfer_trtllm() + unquantized_trtllm_param_ids = self._unquantized_flashinfer_trtllm_param_ids() specs = {} for hf_name, (vllm_param, merged_slice) in vllm_param_map_and_slices.items(): param_info = param_info_by_name[hf_name] - if use_trtllm_staging and param_info.get("grouped_expert_proj"): + if id(vllm_param) in unquantized_trtllm_param_ids and param_info.get( + "grouped_expert_proj" + ): specs[hf_name] = _trtllm_grouped_expert_spec(param_info, vllm_param) continue diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index 37a3ad183f3..a63e71a690c 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -58,9 +58,18 @@ def _make_ext(vllm_params): named_modules=lambda: [], ) ext.model_runner = SimpleNamespace(model=model) + ext._unquantized_flashinfer_trtllm_param_ids = lambda: set() return ext +def _enable_trtllm_staging(ext): + """Mark every synthetic model parameter as owned by a BF16 TRTLLM module.""" + ext._uses_unquantized_flashinfer_trtllm = lambda: True + ext._unquantized_flashinfer_trtllm_param_ids = lambda: { + id(param) for _, param in ext.model_runner.model.named_parameters() + } + + def _param(*shape): return torch.empty(*shape) @@ -365,7 +374,7 @@ def test_build_hf_to_local_param_map_stages_trtllm_local_experts(): ) ext.device = torch.device("cpu") ext.pp_comm_groups = {0: SimpleNamespace(rank=9)} - ext._uses_unquantized_flashinfer_trtllm = lambda: True + _enable_trtllm_staging(ext) ext._load_full_hf_weights = MagicMock( return_value={"model.layers.0.mlp.experts.routed_experts.w13_weight"} ) @@ -430,7 +439,7 @@ def test_build_hf_to_local_param_map_stages_nemotron_lightning_padded_experts(): ext = _make_ext({runtime_name: packed_runtime}) ext.device = torch.device("meta") ext.pp_comm_groups = {0: SimpleNamespace(rank=0)} - ext._uses_unquantized_flashinfer_trtllm = lambda: True + _enable_trtllm_staging(ext) ext._load_full_hf_weights = MagicMock(return_value={runtime_name}) spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) @@ -449,6 +458,133 @@ def test_build_hf_to_local_param_map_stages_nemotron_lightning_padded_experts(): assert loaded_weights[0][1].shape == (intermediate_size, hidden_size) +def test_build_hf_to_local_param_map_stages_only_bf16_trtllm_experts( + monkeypatch, +): + """Mixed MXFP8 models stage only first/last BF16 expert layers.""" + num_experts, intermediate_size, hidden_size = 2, 64, 32 + first_bf16_name = "backbone.layers.0.mlp.experts.up_proj.weight" + mxfp8_name = "backbone.layers.1.mlp.experts.up_proj.weight" + last_bf16_name = "backbone.layers.2.mlp.experts.up_proj.weight" + first_bf16_runtime_name = "model.layers.0.mlp.experts.routed_experts.w13_weight" + mxfp8_runtime_name = "model.layers.1.mlp.experts.routed_experts.w13_weight" + last_bf16_runtime_name = "model.layers.2.mlp.experts.routed_experts.w13_weight" + refit_info = { + "gen_tp_size": 1, + "layer_names": [ + "backbone.layers.0", + "backbone.layers.1", + "backbone.layers.2", + ], + "per_layer_params": { + "backbone.layers.0": [ + { + "name": first_bf16_name, + "global_shape": [num_experts, intermediate_size, hidden_size], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "up_proj", + "dst_mesh_info": MeshInfo(torch.tensor([0])), + "dst_placements": [Shard(0)], + } + ], + "backbone.layers.1": [ + { + "name": mxfp8_name, + "global_shape": [num_experts, intermediate_size, hidden_size], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "up_proj", + "dst_mesh_info": MeshInfo(torch.tensor([0])), + "dst_placements": [Shard(0)], + } + ], + "backbone.layers.2": [ + { + "name": last_bf16_name, + "global_shape": [num_experts, intermediate_size, hidden_size], + "dtype": "torch.bfloat16", + "grouped_expert_proj": "up_proj", + "dst_mesh_info": MeshInfo(torch.tensor([0])), + "dst_placements": [Shard(0)], + } + ], + }, + } + first_bf16_runtime = torch.empty( + num_experts, hidden_size // 16, intermediate_size, 16, dtype=torch.bfloat16 + ) + last_bf16_runtime = torch.empty( + num_experts, hidden_size // 16, intermediate_size, 16, dtype=torch.bfloat16 + ) + mxfp8_runtime = torch.empty( + num_experts, intermediate_size, hidden_size, dtype=torch.float8_e4m3fn + ) + mxfp8_scale = torch.empty( + num_experts, intermediate_size, hidden_size // 32, dtype=torch.uint8 + ) + ext = _make_ext( + { + first_bf16_runtime_name: first_bf16_runtime, + mxfp8_runtime_name: mxfp8_runtime, + f"{mxfp8_runtime_name}_scale_from_checkpoint": mxfp8_scale, + last_bf16_runtime_name: last_bf16_runtime, + } + ) + ext.device = torch.device("cpu") + ext.pp_comm_groups = {0: SimpleNamespace(rank=0)} + _enable_trtllm_staging(ext) + ext._unquantized_flashinfer_trtllm_param_ids = lambda: { + id(first_bf16_runtime), + id(last_bf16_runtime), + } + ext._load_full_hf_weights = MagicMock(return_value=None) + + def fake_quantize(weight): + return ( + torch.full_like(weight, 3, dtype=torch.float8_e4m3fn), + torch.full( + (*weight.shape[:-1], weight.shape[-1] // 32), + 7, + dtype=torch.uint8, + ), + ) + + monkeypatch.setattr( + "nemo_rl.models.generation.vllm.quantization.fp8.quantize_mxfp8_weight", + fake_quantize, + ) + + specs = ext.build_hf_to_local_param_map(refit_info) + first_bf16_spec = specs.get(first_bf16_name) + mxfp8_spec = specs.get(mxfp8_name) + last_bf16_spec = specs.get(last_bf16_name) + assert first_bf16_spec is not None and first_bf16_spec.pre is not None + assert first_bf16_spec.post is not None and first_bf16_spec.base is None + assert mxfp8_spec is not None and mxfp8_spec.pre is not None + assert mxfp8_spec.post is not None and mxfp8_spec.base is not None + assert last_bf16_spec is not None and last_bf16_spec.pre is not None + assert last_bf16_spec.post is not None and last_bf16_spec.base is None + + first_bf16_spec.post(first_bf16_spec.pre(first_bf16_spec.base)) + mxfp8_spec.post(mxfp8_spec.pre(mxfp8_spec.base)) + last_bf16_spec.post(last_bf16_spec.pre(last_bf16_spec.base)) + + assert ext._load_full_hf_weights.call_count == 2 + loaded_source_names = { + weight_name + for call in ext._load_full_hf_weights.call_args_list + for weight_name, _ in call.args[0] + } + assert loaded_source_names == { + f"backbone.layers.0.mlp.experts.{expert}.up_proj.weight" + for expert in range(num_experts) + } | { + f"backbone.layers.2.mlp.experts.{expert}.up_proj.weight" + for expert in range(num_experts) + } + assert torch.all(mxfp8_runtime.float() == 3) + assert torch.all(mxfp8_scale == 7) + + def test_build_hf_to_local_param_map_stages_qwen35_wrapped_experts(): """Qwen3.5 wrapper prefixes and RoutedExperts names survive staged reload.""" hidden_size, num_experts, intermediate_size = 16, 4, 32 @@ -480,7 +616,7 @@ def test_build_hf_to_local_param_map_stages_qwen35_wrapped_experts(): ext = _make_ext({runtime_name: packed_w2}) ext.device = torch.device("cpu") ext.pp_comm_groups = {0: SimpleNamespace(rank=9)} - ext._uses_unquantized_flashinfer_trtllm = lambda: True + _enable_trtllm_staging(ext) ext._load_full_hf_weights = MagicMock(return_value={f"{runtime_prefix}.w2_weight"}) spec = ext.build_hf_to_local_param_map(refit_info).get(expert_name) @@ -518,7 +654,7 @@ def test_build_hf_to_local_param_map_requires_pp_groups_for_trtllm_staging(): ) ext.device = torch.device("cpu") ext.pp_comm_groups = None - ext._uses_unquantized_flashinfer_trtllm = lambda: True + _enable_trtllm_staging(ext) with pytest.raises(RuntimeError, match="before.*per-PP-stage groups"): ext.build_hf_to_local_param_map(refit_info) @@ -557,7 +693,7 @@ def test_build_hf_to_local_param_map_rejects_missing_trtllm_destination(): ) ext.device = torch.device("cpu") ext.pp_comm_groups = {0: SimpleNamespace(rank=9)} - ext._uses_unquantized_flashinfer_trtllm = lambda: True + _enable_trtllm_staging(ext) ext._load_full_hf_weights = MagicMock( return_value={"model.layers.0.mlp.experts.w13_weight"} ) @@ -595,7 +731,7 @@ def test_build_hf_to_local_param_map_rejects_trtllm_tensor_sharding(): ), } ) - ext._uses_unquantized_flashinfer_trtllm = lambda: True + _enable_trtllm_staging(ext) with pytest.raises(ValueError, match="unsupported tensor shard dimensions"): ext.build_hf_to_local_param_map(refit_info) @@ -697,7 +833,7 @@ def init_nccl_communicator(self, *, device): def test_nccl_reshard_lifecycle_repeats_for_trtllm_moe_modules(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend - model = SimpleNamespace() + model = torch.nn.Module() trtllm_moe = SimpleNamespace() model_config = object() vllm_config = object() diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 2a57b7f57f8..201f2fb25cd 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -334,6 +334,95 @@ def set_current_vllm_config(config): ext._maybe_process_fp8_kv_cache.assert_not_called() +@pytest.mark.vllm +def test_mixed_mxfp8_native_refit_processes_bf16_and_mxfp8_modules(monkeypatch): + """Mixed rollout refits rebuild both runtime expert layouts.""" + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptMxFp8FusedMoE, + ) + + from nemo_rl.models.generation.vllm import vllm_backend + + call_order = [] + model = torch.nn.Module() + first_bf16_moe = torch.nn.Module() + first_bf16_moe.expert_map_manager = SimpleNamespace(placement_strategy="linear") + mxfp8_moe = torch.nn.Module() + mxfp8_moe.quant_method = ModelOptMxFp8FusedMoE.__new__(ModelOptMxFp8FusedMoE) + last_bf16_moe = torch.nn.Module() + last_bf16_moe.expert_map_manager = SimpleNamespace(placement_strategy="linear") + model.add_module("first_bf16_moe", first_bf16_moe) + model.add_module("middle_mxfp8_moe", mxfp8_moe) + model.add_module("last_bf16_moe", last_bf16_moe) + + model_config = object() + vllm_config = SimpleNamespace(quant_config=object()) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=vllm_config) + ext.model_config = model_config + ext.device = torch.device("cpu") + ext._mtp_drafter_refit_enabled = lambda: False + ext._maybe_process_mtp_drafter_after_loading = lambda: call_order.append("mtp") + ext._maybe_process_fp8_kv_cache = MagicMock() + + monkeypatch.setattr( + vllm_backend, + "_unquantized_flashinfer_trtllm_modules", + lambda _model: [first_bf16_moe, last_bf16_moe], + ) + monkeypatch.setattr( + ModelOptMxFp8FusedMoE, + "process_weights_after_loading", + lambda _self, module: call_order.append(("process_mxfp8", module)), + ) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + + @contextlib.contextmanager + def set_current_vllm_config(config): + assert config is vllm_config + call_order.append("config_enter") + try: + yield + finally: + call_order.append("config_exit") + + monkeypatch.setattr("vllm.config.set_current_vllm_config", set_current_vllm_config) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", + lambda module: call_order.append(("initialize", module)), + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", + lambda reload_model, config: call_order.append( + ("finalize", reload_model, config) + ), + ) + monkeypatch.setattr( + vllm_backend, + "_refresh_hpc_modules_after_layerwise_reload", + lambda reload_model: call_order.append(("hpc", reload_model)), + ) + + with ext._weight_update_lifecycle("nccl_reshard") as finalize: + call_order.append("transfer") + finalize() + + assert call_order == [ + "config_enter", + ("initialize", first_bf16_moe), + ("initialize", last_bf16_moe), + "transfer", + ("finalize", model, model_config), + ("process_mxfp8", mxfp8_moe), + ("hpc", model), + "mtp", + "config_exit", + ] + ext._maybe_process_fp8_kv_cache.assert_not_called() + + @pytest.mark.vllm def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypatch): from vllm.model_executor.model_loader.reload import record_metadata_for_reloading @@ -450,7 +539,7 @@ def test_layerwise_reload_propagates_detach_error_after_successful_load(monkeypa def test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend - model = object() + model = SimpleNamespace(modules=lambda: []) model_config = object() vllm_config = SimpleNamespace( kernel_config=SimpleNamespace(moe_backend="flashinfer_trtllm"), @@ -529,7 +618,8 @@ def test_realized_moe_backend_controls_native_refit_lifecycle(): @pytest.mark.vllm -def test_quantized_model_does_not_use_unquantized_refit_lifecycle(): +def test_quantized_model_uses_native_refit_for_realized_bf16_trtllm_modules(): + """A globally quantized model may still contain ignored BF16 MoE layers.""" from nemo_rl.models.generation.vllm import vllm_backend ext = vllm_backend.VllmInternalWorkerExtension.__new__( @@ -540,7 +630,45 @@ def test_quantized_model_does_not_use_unquantized_refit_lifecycle(): vllm_config=SimpleNamespace(quant_config=object()), ) - assert ext._uses_unquantized_flashinfer_trtllm() is False + assert ext._uses_unquantized_flashinfer_trtllm() is True + + +@pytest.mark.vllm +def test_unquantized_trtllm_param_ids_are_scoped_to_realized_modules(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + bf16_experts = torch.nn.Module() + bf16_experts.register_parameter( + "w13_weight", torch.nn.Parameter(torch.empty(2, 4, 8), requires_grad=False) + ) + bf16_experts.register_parameter( + "w2_weight", torch.nn.Parameter(torch.empty(2, 8, 2), requires_grad=False) + ) + mxfp8_experts = torch.nn.Module() + mxfp8_experts.register_parameter( + "w13_weight", + torch.nn.Parameter( + torch.empty(2, 4, 8, dtype=torch.float8_e4m3fn), requires_grad=False + ), + ) + model = torch.nn.Module() + model.add_module("bf16_experts", bf16_experts) + model.add_module("mxfp8_experts", mxfp8_experts) + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model) + monkeypatch.setattr( + vllm_backend, + "_unquantized_flashinfer_trtllm_modules", + lambda _model: [bf16_experts], + ) + + assert ext._unquantized_flashinfer_trtllm_param_ids() == { + id(bf16_experts.w13_weight), + id(bf16_experts.w2_weight), + } @pytest.mark.vllm @@ -549,7 +677,7 @@ def test_quantized_model_does_not_use_unquantized_refit_lifecycle(): [ ("FlashInfer TRTLLM", None, True), ("TRITON", None, False), - ("FlashInfer TRTLLM", object(), False), + ("FlashInfer TRTLLM", object(), True), ], ) def test_weight_update_errors_are_fatal_only_for_native_trtllm_refit( diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index 08d7891c285..f4bbd01b0a2 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -153,6 +153,167 @@ def test_init_fp8_passes_modelopt_ignore_patterns_without_hf_expansion( assert not modelopt_config.is_layer_excluded("model.layers.0.mlp.gate_up_proj") +@pytest.mark.parametrize( + ("num_first_layers_in_bf16", "num_last_layers_in_bf16"), + [ + (0, 0), + (1, 1), + (2, 6), + (7, 3), + (26, 26), + (30, 30), + ], +) +@pytest.mark.parametrize( + "recipe_case", + [ + pytest.param( + types.SimpleNamespace( + model_name="dummy-qwen-model", + hf_layer_prefix="layers", + raw_layer_prefix="model.layers", + mapper_prefixes={}, + hf_target_suffixes=( + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "mlp.experts.gate_proj", + "mlp.experts.up_proj", + "mlp.experts.down_proj", + ), + vllm_target_suffixes=( + "self_attn.qkv_proj", + "self_attn.o_proj", + "mlp.experts", + ), + non_target_suffixes=("mlp.gate", "mlp.shared_experts.up_proj"), + ignore_patterns=( + "*layers.*.mlp.gate", + "*layers.*.mlp.shared_experts.*", + "lm_head", + ), + ), + id="qwen", + ), + pytest.param( + types.SimpleNamespace( + model_name="dummy-nemotron-h-model", + hf_layer_prefix="backbone.layers", + raw_layer_prefix="backbone.layers", + mapper_prefixes={"backbone": "model"}, + hf_target_suffixes=( + "mixer.q_proj", + "mixer.k_proj", + "mixer.v_proj", + "mixer.o_proj", + "mixer.experts.up_proj", + "mixer.experts.down_proj", + ), + vllm_target_suffixes=( + "mixer.qkv_proj", + "mixer.o_proj", + "mixer.experts", + ), + non_target_suffixes=( + "mixer.in_proj", + "mixer.shared_experts.up_proj", + ), + ignore_patterns=( + "*layers.*.mixer.in_proj", + "*layers.*.mixer.shared_experts.*", + "lm_head", + ), + ), + id="nemotron-h", + ), + ], +) +def test_init_fp8_keeps_mixed_recipe_boundary_targets_in_bf16( + fp8_module, + monkeypatch, + num_first_layers_in_bf16, + num_last_layers_in_bf16, + recipe_case, +): + """Keep boundary QKVO and routed experts BF16 across mixed recipes.""" + from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config + from vllm.model_executor.models.utils import WeightsMapper + + fp8 = fp8_module + num_hidden_layers = 52 + param_names = [] + for layer_idx in range(num_hidden_layers): + param_names.extend( + f"{recipe_case.hf_layer_prefix}.{layer_idx}.{suffix}.weight" + for suffix in ( + *recipe_case.hf_target_suffixes, + *recipe_case.non_target_suffixes, + ) + ) + + monkeypatch.setattr( + fp8.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: types.SimpleNamespace( + num_hidden_layers=num_hidden_layers + ), + ) + monkeypatch.setattr( + fp8.AutoModel, + "from_config", + lambda *_args, **_kwargs: types.SimpleNamespace( + named_parameters=lambda: [(name, None) for name in param_names] + ), + ) + monkeypatch.setattr(fp8, "monkey_patch_vllm_ray_executor", lambda _config: None) + + vllm_kwargs = fp8.init_fp8( + { + "precision": "fp8", + "kv_cache_dtype": "auto", + "async_engine": False, + "is_mx": True, + "num_first_layers_in_bf16": num_first_layers_in_bf16, + "num_last_layers_in_bf16": num_last_layers_in_bf16, + "quantization_ignore_patterns": list(recipe_case.ignore_patterns), + }, + recipe_case.model_name, + model_parallel_size=1, + ) + + quant_config = vllm_kwargs["hf_overrides"]["quantization_config"] + modelopt_config = ModelOptMxFp8Config.from_config(quant_config) + mapper = WeightsMapper(orig_to_new_prefix=recipe_case.mapper_prefixes) + modelopt_config.apply_vllm_mapper(mapper.get_unstacked_mapper()) + modelopt_config.packed_modules_mapping.update( + {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + ) + boundary_layers = { + layer_idx + for layer_idx in range(num_hidden_layers) + if layer_idx < num_first_layers_in_bf16 + or layer_idx >= num_hidden_layers - num_last_layers_in_bf16 + } + for layer_idx in range(num_hidden_layers): + if layer_idx in boundary_layers: + for suffix in recipe_case.hf_target_suffixes: + module_name = f"{recipe_case.raw_layer_prefix}.{layer_idx}.{suffix}" + assert module_name in quant_config["ignored_layers"] + + for suffix in recipe_case.vllm_target_suffixes: + module_name = f"model.layers.{layer_idx}.{suffix}" + assert modelopt_config.is_layer_excluded(module_name) + else: + for suffix in recipe_case.vllm_target_suffixes: + module_name = f"model.layers.{layer_idx}.{suffix}" + assert not modelopt_config.is_layer_excluded(module_name) + + for suffix in recipe_case.non_target_suffixes: + module_name = f"model.layers.{layer_idx}.{suffix}" + assert modelopt_config.is_layer_excluded(module_name) + + @pytest.mark.parametrize( "config", [ From feff039f278a8577d030a9aec2298137b245f498 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 23:21:17 -0700 Subject: [PATCH 33/68] test(vllm): cover mixed refit scope boundaries Signed-off-by: seonjinn --- .../models/generation/test_vllm_backend.py | 115 +++++++++++++++--- 1 file changed, 96 insertions(+), 19 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 201f2fb25cd..2b498adf54c 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -203,7 +203,7 @@ def test_init_collective_keeps_generation_ranks_after_the_training_ranks( def _make_unquantized_moe_model( moe_backend: str, expert_placement_strategy: str = "linear" -) -> SimpleNamespace: +) -> torch.nn.Module: from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, ) @@ -213,13 +213,14 @@ def _make_unquantized_moe_model( quant_method = UnquantizedFusedMoEMethod.__new__(UnquantizedFusedMoEMethod) quant_method.unquantized_backend = UnquantizedMoeBackend(moe_backend) - module = SimpleNamespace( - quant_method=quant_method, - expert_map_manager=SimpleNamespace( - placement_strategy=expert_placement_strategy - ), + model = torch.nn.Module() + module = torch.nn.Module() + module.quant_method = quant_method + module.expert_map_manager = SimpleNamespace( + placement_strategy=expert_placement_strategy ) - return SimpleNamespace(modules=lambda: [module]) + model.add_module("moe", module) + return model @pytest.mark.vllm @@ -267,6 +268,7 @@ def test_unquantized_weight_update_uses_layerwise_reload(monkeypatch): call_order = [] model = _make_unquantized_moe_model("FlashInfer TRTLLM") + moe_module = vllm_backend._unquantized_flashinfer_trtllm_modules(model)[0] model_config = object() vllm_config = SimpleNamespace( kernel_config=SimpleNamespace(moe_backend="auto"), quant_config=None @@ -323,7 +325,7 @@ def set_current_vllm_config(config): expected_cycle = [ "config_enter", - ("initialize", model), + ("initialize", moe_module), "load", ("finalize", model, model_config), ("hpc", model), @@ -335,10 +337,12 @@ def set_current_vllm_config(config): @pytest.mark.vllm -def test_mixed_mxfp8_native_refit_processes_bf16_and_mxfp8_modules(monkeypatch): - """Mixed rollout refits rebuild both runtime expert layouts.""" +@pytest.mark.parametrize("transport", ["ipc", "collective", "nccl_reshard"]) +def test_mixed_mxfp8_native_refit_processes_each_module_once(monkeypatch, transport): + """Mixed refits reload BF16 experts and rebuild each MXFP8 layout once.""" from vllm.model_executor.layers.quantization.modelopt import ( ModelOptMxFp8FusedMoE, + ModelOptMxFp8LinearMethod, ) from nemo_rl.models.generation.vllm import vllm_backend @@ -349,10 +353,15 @@ def test_mixed_mxfp8_native_refit_processes_bf16_and_mxfp8_modules(monkeypatch): first_bf16_moe.expert_map_manager = SimpleNamespace(placement_strategy="linear") mxfp8_moe = torch.nn.Module() mxfp8_moe.quant_method = ModelOptMxFp8FusedMoE.__new__(ModelOptMxFp8FusedMoE) + mxfp8_qkv = torch.nn.Module() + mxfp8_qkv.quant_method = ModelOptMxFp8LinearMethod.__new__( + ModelOptMxFp8LinearMethod + ) last_bf16_moe = torch.nn.Module() last_bf16_moe.expert_map_manager = SimpleNamespace(placement_strategy="linear") model.add_module("first_bf16_moe", first_bf16_moe) model.add_module("middle_mxfp8_moe", mxfp8_moe) + model.add_module("middle_mxfp8_qkv", mxfp8_qkv) model.add_module("last_bf16_moe", last_bf16_moe) model_config = object() @@ -372,10 +381,15 @@ def test_mixed_mxfp8_native_refit_processes_bf16_and_mxfp8_modules(monkeypatch): "_unquantized_flashinfer_trtllm_modules", lambda _model: [first_bf16_moe, last_bf16_moe], ) + + def process_mxfp8(_self, module): + call_order.append(("process_mxfp8", module)) + monkeypatch.setattr( - ModelOptMxFp8FusedMoE, - "process_weights_after_loading", - lambda _self, module: call_order.append(("process_mxfp8", module)), + ModelOptMxFp8FusedMoE, "process_weights_after_loading", process_mxfp8 + ) + monkeypatch.setattr( + ModelOptMxFp8LinearMethod, "process_weights_after_loading", process_mxfp8 ) monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) @@ -389,15 +403,28 @@ def set_current_vllm_config(config): call_order.append("config_exit") monkeypatch.setattr("vllm.config.set_current_vllm_config", set_current_vllm_config) + initialized_targets = [] + + def initialize(module): + initialized_targets.append(module) + call_order.append(("initialize", module)) + + def finalize(reload_model, config): + call_order.append(("finalize", reload_model, config)) + for target in initialized_targets: + for module in target.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance( + quant_method, (ModelOptMxFp8FusedMoE, ModelOptMxFp8LinearMethod) + ): + quant_method.process_weights_after_loading(module) + monkeypatch.setattr( "vllm.model_executor.model_loader.reload.initialize_layerwise_reload", - lambda module: call_order.append(("initialize", module)), + initialize, ) monkeypatch.setattr( - "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", - lambda reload_model, config: call_order.append( - ("finalize", reload_model, config) - ), + "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", finalize ) monkeypatch.setattr( vllm_backend, @@ -405,7 +432,7 @@ def set_current_vllm_config(config): lambda reload_model: call_order.append(("hpc", reload_model)), ) - with ext._weight_update_lifecycle("nccl_reshard") as finalize: + with ext._weight_update_lifecycle(transport) as finalize: call_order.append("transfer") finalize() @@ -416,6 +443,7 @@ def set_current_vllm_config(config): "transfer", ("finalize", model, model_config), ("process_mxfp8", mxfp8_moe), + ("process_mxfp8", mxfp8_qkv), ("hpc", model), "mtp", "config_exit", @@ -423,6 +451,55 @@ def set_current_vllm_config(config): ext._maybe_process_fp8_kv_cache.assert_not_called() +@pytest.mark.vllm +@pytest.mark.parametrize("transport", ["ipc", "collective"]) +def test_mixed_native_refit_preserves_post_load_mxfp8_scale(monkeypatch, transport): + """Native reload must not restore MXFP8 linears to pre-load metadata.""" + from vllm.model_executor.model_loader.reload import record_metadata_for_reloading + + from nemo_rl.models.generation.vllm import vllm_backend + + model = torch.nn.Module() + bf16_moe = torch.nn.Module() + bf16_moe.register_parameter( + "weight", torch.nn.Parameter(torch.zeros(2), requires_grad=False) + ) + mxfp8_qkv = torch.nn.Module() + mxfp8_qkv.register_parameter( + "weight", torch.nn.Parameter(torch.zeros(2), requires_grad=False) + ) + model.add_module("first_bf16_moe", bf16_moe) + model.add_module("middle_mxfp8_qkv", mxfp8_qkv) + record_metadata_for_reloading(model) + + mxfp8_qkv.register_parameter( + "weight_scale_from_checkpoint", + torch.nn.Parameter(torch.ones(2), requires_grad=False), + ) + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=model, vllm_config=SimpleNamespace(quant_config=object()) + ) + ext.model_config = object() + ext.device = torch.device("cpu") + ext._maybe_process_fp8_kv_cache = MagicMock() + + monkeypatch.setattr( + vllm_backend, + "_unquantized_flashinfer_trtllm_modules", + lambda _model: [bf16_moe], + ) + monkeypatch.setattr( + "vllm.config.set_current_vllm_config", lambda _config: contextlib.nullcontext() + ) + + with ext._weight_update_lifecycle(transport): + assert hasattr(mxfp8_qkv, "weight_scale_from_checkpoint") + + @pytest.mark.vllm def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypatch): from vllm.model_executor.model_loader.reload import record_metadata_for_reloading From f2214aa62a10a8ea3ae860e32bddf36d00ca0f46 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 23:26:58 -0700 Subject: [PATCH 34/68] test(vllm): model realized MoE backend faithfully Signed-off-by: seonjinn --- tests/unit/models/generation/test_vllm_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 2b498adf54c..b1bb8cfe6e4 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -215,7 +215,7 @@ def _make_unquantized_moe_model( quant_method.unquantized_backend = UnquantizedMoeBackend(moe_backend) model = torch.nn.Module() module = torch.nn.Module() - module.quant_method = quant_method + module.__dict__["quant_method"] = quant_method module.expert_map_manager = SimpleNamespace( placement_strategy=expert_placement_strategy ) From 48301b6a417f05e900b7e1f4679d543e3d6cb501 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 23:29:22 -0700 Subject: [PATCH 35/68] fix(vllm): scope native reload to BF16 TRTLLM modules Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 33 ++++++++++++------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 422f9de507d..6c2b9c1a3ce 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -140,8 +140,17 @@ def _unquantized_flashinfer_trtllm_modules( ] -def _process_mxfp8_modules_after_native_reload(model: torch.nn.Module) -> None: - """Rebuild MXFP8 runtime layouts skipped by the partial native reload.""" +def _reload_target_module_ids( + reload_targets: Sequence[torch.nn.Module], +) -> set[int]: + """Return every module covered by the native reload targets.""" + return {id(module) for target in reload_targets for module in target.modules()} + + +def _process_mxfp8_modules_after_native_reload( + model: torch.nn.Module, reloaded_module_ids: set[int] +) -> None: + """Rebuild MXFP8 layouts outside the partial native reload.""" try: from vllm.model_executor.layers.quantization.modelopt import ( ModelOptMxFp8FusedMoE, @@ -152,6 +161,8 @@ def _process_mxfp8_modules_after_native_reload(model: torch.nn.Module) -> None: mxfp8_methods = (ModelOptMxFp8FusedMoE, ModelOptMxFp8LinearMethod) for module in model.modules(): + if id(module) in reloaded_module_ids: + continue quant_method = getattr(module, "quant_method", None) if isinstance(quant_method, mxfp8_methods): quant_method.process_weights_after_loading(module) @@ -967,19 +978,19 @@ def _weight_update_lifecycle( ) model = self.model_runner.model - # NCCL reshard receives most weights directly into live parameter - # storage; only the TRTLLM MoE modules need the layerwise reload - # lifecycle to rebuild the kernel's private repacked layout. - reload_targets = ( - _unquantized_flashinfer_trtllm_modules(model) - if transport == "nccl_reshard" - else [model] - ) + # Restore only the realized BF16 TRTLLM modules. MXFP8 modules own + # checkpoint-scale parameters created after vLLM recorded reload + # metadata, so restoring the whole mixed model would delete those + # parameters. Their layouts are rebuilt separately after transfer. + reload_targets = _unquantized_flashinfer_trtllm_modules(model) + 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) + _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() From 7c1eb095ac3f6ff4c2efb2014044922e8ddb18c9 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 23:32:21 -0700 Subject: [PATCH 36/68] test(vllm): target deferred reload fixture by module Signed-off-by: seonjinn --- tests/unit/models/generation/test_vllm_backend.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index b1bb8cfe6e4..a524c312ade 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -516,6 +516,11 @@ def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypa ext._uses_unquantized_flashinfer_trtllm = lambda: True ext._validate_native_layerwise_refit = lambda _transport=None: None ext._maybe_process_mtp_drafter_after_loading = MagicMock() + monkeypatch.setattr( + vllm_backend, + "_unquantized_flashinfer_trtllm_modules", + lambda _model: [model.layer], + ) monkeypatch.setattr( "vllm.config.set_current_vllm_config", lambda _: contextlib.nullcontext() From 94a22cbe6b12e1be4244a8b69b4c5ae3de58ce83 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 23:35:20 -0700 Subject: [PATCH 37/68] test(vllm): cover mixed FP8 transport safety Signed-off-by: seonjinn --- .../models/generation/test_vllm_backend.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index a524c312ade..7e8bacabdea 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -500,6 +500,29 @@ def test_mixed_native_refit_preserves_post_load_mxfp8_scale(monkeypatch, transpo assert hasattr(mxfp8_qkv, "weight_scale_from_checkpoint") +@pytest.mark.vllm +def test_fp8_load_uses_buffer_safe_model_loader(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=object(), vllm_config=object()) + weights = [("model.weight", torch.ones(2))] + load_weights = MagicMock() + + monkeypatch.setattr(fp8, "is_fp8_model", lambda _config: True) + monkeypatch.setattr(fp8, "load_weights", load_weights) + + ext._load_hf_weights(weights) + + args = load_weights.call_args.args + assert args[:2] == (weights, ext.model_runner) + assert args[2].__self__ is ext + assert args[2].__func__ is ext._load_full_hf_weights.__func__ + + @pytest.mark.vllm def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypatch): from vllm.model_executor.model_loader.reload import record_metadata_for_reloading @@ -715,6 +738,24 @@ def test_quantized_model_uses_native_refit_for_realized_bf16_trtllm_modules(): assert ext._uses_unquantized_flashinfer_trtllm() is True +@pytest.mark.vllm +@pytest.mark.parametrize("transport", ["ipc", "collective", "nccl_reshard"]) +def test_mixed_native_refit_rejects_fp8_kv_cache(transport): + from nemo_rl.models.generation.vllm import vllm_backend + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace( + model=_make_unquantized_moe_model("FlashInfer TRTLLM") + ) + ext._uses_fp8_kv_cache = lambda: True + ext._mtp_drafter_refit_enabled = lambda: False + + with pytest.raises(RuntimeError, match="FP8 KV cache"): + ext._validate_native_layerwise_refit(transport) + + @pytest.mark.vllm def test_unquantized_trtllm_param_ids_are_scoped_to_realized_modules(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend From 6c63d5d3fc068c4d1bda99f6587f9c1e12875f71 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 23:46:47 -0700 Subject: [PATCH 38/68] test(vllm): exercise refit loaders through module contracts Signed-off-by: seonjinn --- .../generation/test_nccl_reshard_backend.py | 2 +- .../models/generation/test_vllm_backend.py | 7 ++++--- .../generation/test_vllm_fp8_quantization.py | 18 ++++++++++++++++++ 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index a63e71a690c..da46546cd8e 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -834,7 +834,7 @@ def test_nccl_reshard_lifecycle_repeats_for_trtllm_moe_modules(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend model = torch.nn.Module() - trtllm_moe = SimpleNamespace() + trtllm_moe = torch.nn.Module() model_config = object() vllm_config = object() call_order = [] diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 7e8bacabdea..cf835edb9f2 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -518,9 +518,10 @@ def test_fp8_load_uses_buffer_safe_model_loader(monkeypatch): ext._load_hf_weights(weights) args = load_weights.call_args.args - assert args[:2] == (weights, ext.model_runner) - assert args[2].__self__ is ext - assert args[2].__func__ is ext._load_full_hf_weights.__func__ + assert args == (weights, ext.model_runner) + model_load_weights = load_weights.call_args.kwargs["model_load_weights"] + assert model_load_weights.__self__ is ext + assert model_load_weights.__func__ is ext._load_full_hf_weights.__func__ @pytest.mark.vllm diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index f4bbd01b0a2..519dd1e6d8e 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -1288,6 +1288,24 @@ class _MoERunner: ) +def test_load_weights_uses_supplied_model_loader(fp8_module): + fp8 = fp8_module + model = types.SimpleNamespace( + packed_modules_mapping={}, + load_weights=lambda _weights: pytest.fail("must use the supplied loader"), + ) + source = torch.ones(2, dtype=torch.bfloat16) + loaded = [] + + fp8.load_weights( + [("model.norm.weight", source)], + types.SimpleNamespace(model=model), + model_load_weights=lambda weights: loaded.extend(weights), + ) + + assert loaded == [("model.norm.weight", source)] + + @GROUPED_EXPERT_KEY_SHAPES def test_load_weights_passes_grouped_experts_through_for_ignored_bf16_layers( fp8_module, monkeypatch, layers_prefix, wrap_language_model From bde0bf6d4776847f92e9c62da5247b436701ce21 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Thu, 3 Sep 2026 23:52:10 -0700 Subject: [PATCH 39/68] fix(vllm): preserve mixed refit loader ownership Signed-off-by: seonjinn --- .../generation/vllm/quantization/fp8.py | 16 ++++++++++---- .../models/generation/vllm/vllm_backend.py | 22 ++++++++++++------- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index ab171075f30..e2f379d8cd0 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -14,7 +14,7 @@ import os import warnings -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass, field from unittest.mock import patch @@ -538,7 +538,12 @@ def quantize_mxfp8_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Ten return value, scale -def load_weights(weights, model_runner): +def load_weights( + weights, + model_runner, + *, + model_load_weights: Callable[..., object] | None = None, +): global global_fp8_config weights_quantized = [] model = model_runner.model @@ -591,8 +596,11 @@ def load_weights(weights, model_runner): else: weights_quantized.append([k, param_lp]) weights_quantized.append([k + "_scale_inv", param_scale]) - # Finally load the weights into vllm - model.load_weights(weights_quantized) + # Finally load the weights into vllm. Native layerwise reload callers pass + # a wrapper that keeps deferred weight-loader tensors off reusable buffers. + if model_load_weights is None: + model_load_weights = model.load_weights + model_load_weights(weights_quantized) def cast_tensor_to_fp8_blockwise( diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 6c2b9c1a3ce..67420c00b26 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -370,7 +370,11 @@ def _load_hf_weights(self, policy_weights: list[tuple[str, torch.Tensor]]) -> No from nemo_rl.models.generation.vllm.quantization import fp8 if fp8.is_fp8_model(self.model_runner.vllm_config): - fp8.load_weights(policy_weights, self.model_runner) + fp8.load_weights( + policy_weights, + self.model_runner, + model_load_weights=self._load_full_hf_weights, + ) return self._load_full_hf_weights(policy_weights) @@ -901,14 +905,16 @@ def _validate_native_layerwise_refit( if not self._uses_unquantized_flashinfer_trtllm(): return - if transport == "nccl_reshard": - if self._uses_fp8_kv_cache(): - raise RuntimeError( - "BF16 FlashInfer TRTLLM nccl_reshard refit does not " - "support an FP8 KV cache because its static scales are " - "outside the targeted MoE reload lifecycle" - ) + if transport in ("ipc", "collective", "nccl_reshard") and ( + self._uses_fp8_kv_cache() + ): + raise RuntimeError( + "BF16 FlashInfer TRTLLM partial refit does not support an " + "FP8 KV cache because its static scales are outside the " + "targeted reload lifecycle" + ) + if transport == "nccl_reshard": realized_placements = set() for module in _unquantized_flashinfer_trtllm_modules( self.model_runner.model From 7f7b8dde3226db01697eccc5850173eb17e7a739 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 00:19:16 -0700 Subject: [PATCH 40/68] test(vllm): cover nested layer count config Signed-off-by: seonjinn --- .../generation/test_vllm_fp8_quantization.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index 519dd1e6d8e..beef1002137 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -314,6 +314,53 @@ def test_init_fp8_keeps_mixed_recipe_boundary_targets_in_bf16( assert modelopt_config.is_layer_excluded(module_name) +def test_init_fp8_reads_layer_count_from_text_config(fp8_module, monkeypatch): + fp8 = fp8_module + num_hidden_layers = 8 + param_names = [ + f"model.layers.{layer_idx}.mlp.experts.up_proj.weight" + for layer_idx in range(num_hidden_layers) + ] + + monkeypatch.setattr( + fp8.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: types.SimpleNamespace( + text_config=types.SimpleNamespace(num_hidden_layers=num_hidden_layers) + ), + ) + monkeypatch.setattr( + fp8.AutoModel, + "from_config", + lambda *_args, **_kwargs: types.SimpleNamespace( + named_parameters=lambda: [(name, None) for name in param_names] + ), + ) + monkeypatch.setattr(fp8, "monkey_patch_vllm_ray_executor", lambda _config: None) + + vllm_kwargs = fp8.init_fp8( + { + "precision": "fp8", + "kv_cache_dtype": "auto", + "async_engine": False, + "is_mx": True, + "num_first_layers_in_bf16": 2, + "num_last_layers_in_bf16": 2, + }, + "dummy-model-with-text-config", + model_parallel_size=1, + ) + + ignored_layers = vllm_kwargs["hf_overrides"]["quantization_config"][ + "ignored_layers" + ] + assert "model.layers.0.mlp.experts.up_proj.weight" in ignored_layers + assert "model.layers.1.mlp.experts.up_proj.weight" in ignored_layers + assert "model.layers.6.mlp.experts.up_proj.weight" in ignored_layers + assert "model.layers.7.mlp.experts.up_proj.weight" in ignored_layers + assert "model.layers.2.mlp.experts.up_proj.weight" not in ignored_layers + + @pytest.mark.parametrize( "config", [ From 2de2d659c6fe1fac4f5a4c6f520139cbf6db7a1d Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 00:24:29 -0700 Subject: [PATCH 41/68] fix(vllm): resolve nested text layer counts Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/quantization/fp8.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index e2f379d8cd0..b3a335d642a 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -305,6 +305,13 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): with init_empty_weights(): model = AutoModel.from_config(config) param_names = [name for name, _ in model.named_parameters()] + get_text_config = getattr(config, "get_text_config", None) + text_config = ( + get_text_config() + if callable(get_text_config) + else getattr(config, "text_config", config) + ) + num_hidden_layers = text_config.num_hidden_layers bf16_params = [] if num_first_layers_in_bf16 > 0: @@ -315,8 +322,8 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): layers = [ l for l in range( - config.num_hidden_layers - num_last_layers_in_bf16, - config.num_hidden_layers, + num_hidden_layers - num_last_layers_in_bf16, + num_hidden_layers, ) ] bf16_params.extend(_get_params_in_layers(param_names, layers)) From b5d52ed72e1cb6ce45ad95c78d22d9ee66dfd8e7 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 00:28:03 -0700 Subject: [PATCH 42/68] test(vllm): use normalized layer names Signed-off-by: seonjinn --- .../models/generation/test_vllm_fp8_quantization.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index beef1002137..114555aa04c 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -318,7 +318,7 @@ def test_init_fp8_reads_layer_count_from_text_config(fp8_module, monkeypatch): fp8 = fp8_module num_hidden_layers = 8 param_names = [ - f"model.layers.{layer_idx}.mlp.experts.up_proj.weight" + f"layers.{layer_idx}.mlp.experts.up_proj.weight" for layer_idx in range(num_hidden_layers) ] @@ -354,11 +354,11 @@ def test_init_fp8_reads_layer_count_from_text_config(fp8_module, monkeypatch): ignored_layers = vllm_kwargs["hf_overrides"]["quantization_config"][ "ignored_layers" ] - assert "model.layers.0.mlp.experts.up_proj.weight" in ignored_layers - assert "model.layers.1.mlp.experts.up_proj.weight" in ignored_layers - assert "model.layers.6.mlp.experts.up_proj.weight" in ignored_layers - assert "model.layers.7.mlp.experts.up_proj.weight" in ignored_layers - assert "model.layers.2.mlp.experts.up_proj.weight" not in ignored_layers + assert "model.layers.0.mlp.experts.up_proj" in ignored_layers + assert "model.layers.1.mlp.experts.up_proj" in ignored_layers + assert "model.layers.6.mlp.experts.up_proj" in ignored_layers + assert "model.layers.7.mlp.experts.up_proj" in ignored_layers + assert "model.layers.2.mlp.experts.up_proj" not in ignored_layers @pytest.mark.parametrize( From 299a3570bf4657b41c799954d5de065760615eac Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 00:47:48 -0700 Subject: [PATCH 43/68] test(vllm): cover static MTP ownership Signed-off-by: seonjinn --- .../models/generation/test_vllm_backend.py | 53 +++++++++++++++++++ .../generation/test_vllm_fp8_quantization.py | 15 +++--- .../generation/test_vllm_sparse_refit.py | 6 +++ 3 files changed, 68 insertions(+), 6 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index cf835edb9f2..d8c688411a6 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -1321,6 +1321,47 @@ def test_read_mtp_layer_weights_from_checkpoint_filters_and_reads(tmp_path): assert torch.equal(by_name["model.layers.2.shared_head.head.weight"], mtp_head) +@pytest.mark.vllm +def test_read_mtp_layer_weights_from_checkpoint_reads_local_mtp_namespace(tmp_path): + """Read locally numbered MTP layers and their top-level tensors.""" + from nemo_rl.models.generation.vllm.vllm_backend import ( + _read_mtp_layer_weights_from_checkpoint, + ) + + model_dir = tmp_path / "ckpt" + local_layer = torch.randn(4, 4) + projection = torch.randn(4, 4) + norm = torch.randn(4) + base_layer = torch.randn(4, 4) + _write_sharded_checkpoint( + model_dir, + { + "model-00001-of-00002.safetensors": { + "model.mtp.layers.0.mlp.experts.up_proj.weight": local_layer, + "model.mtp.fc.weight": projection, + "model.layers.0.mlp.experts.up_proj.weight": base_layer, + }, + "model-00002-of-00002.safetensors": { + "language_model.mtp.norm.weight": norm, + }, + }, + ) + + weights = _read_mtp_layer_weights_from_checkpoint(str(model_dir), {52}) + + by_name = dict(weights) + assert set(by_name) == { + "model.mtp.layers.0.mlp.experts.up_proj.weight", + "model.mtp.fc.weight", + "language_model.mtp.norm.weight", + } + assert torch.equal( + by_name["model.mtp.layers.0.mlp.experts.up_proj.weight"], local_layer + ) + assert torch.equal(by_name["model.mtp.fc.weight"], projection) + assert torch.equal(by_name["language_model.mtp.norm.weight"], norm) + + @pytest.mark.vllm def test_load_mtp_weights_from_disk_loads_only_mtp_layer(tmp_path, monkeypatch): """Success path: only MTP-layer weights are handed to the drafter, then post-loaded.""" @@ -1505,6 +1546,18 @@ def test_mtp_drafter_refit_enabled(method, from_disk, has_drafter, expected): assert ext._mtp_drafter_refit_enabled() is expected +@pytest.mark.vllm +@pytest.mark.parametrize("weights_from_refit", [False, True]) +def test_configure_mtp_drafter_weight_source(weights_from_refit): + """Checkpoint-loaded MTP stays static for both dummy and auto model loads.""" + ext, _ = _make_mtp_refit_extension(method="mtp", from_disk=False) + + ext.configure_mtp_drafter_weight_source(weights_from_refit) + + assert ext._mtp_drafter_weights_from_refit is weights_from_refit + assert ext._mtp_drafter_refit_enabled() is weights_from_refit + + @pytest.mark.vllm def test_maybe_refit_mtp_drafter_loads_when_enabled(): """A co-trained MTP drafter is fed the (vocab-trimmed) policy weights on refit.""" diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index 114555aa04c..a95e5017c0b 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -329,13 +329,15 @@ def test_init_fp8_reads_layer_count_from_text_config(fp8_module, monkeypatch): text_config=types.SimpleNamespace(num_hidden_layers=num_hidden_layers) ), ) - monkeypatch.setattr( - fp8.AutoModel, - "from_config", - lambda *_args, **_kwargs: types.SimpleNamespace( + from_config_calls = [] + + def from_config(config, **kwargs): + from_config_calls.append((config, kwargs)) + return types.SimpleNamespace( named_parameters=lambda: [(name, None) for name in param_names] - ), - ) + ) + + monkeypatch.setattr(fp8.AutoModel, "from_config", from_config) monkeypatch.setattr(fp8, "monkey_patch_vllm_ray_executor", lambda _config: None) vllm_kwargs = fp8.init_fp8( @@ -359,6 +361,7 @@ def test_init_fp8_reads_layer_count_from_text_config(fp8_module, monkeypatch): assert "model.layers.6.mlp.experts.up_proj" in ignored_layers assert "model.layers.7.mlp.experts.up_proj" in ignored_layers assert "model.layers.2.mlp.experts.up_proj" not in ignored_layers + assert from_config_calls[0][1] == {"trust_remote_code": True} @pytest.mark.parametrize( diff --git a/tests/unit/models/generation/test_vllm_sparse_refit.py b/tests/unit/models/generation/test_vllm_sparse_refit.py index e70c809ca6f..96253439d44 100644 --- a/tests/unit/models/generation/test_vllm_sparse_refit.py +++ b/tests/unit/models/generation/test_vllm_sparse_refit.py @@ -555,6 +555,8 @@ async def test_async_sparse_refit_post_init_records_worker_locality() -> None: worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) worker._sparse_refit_receiver = MagicMock() worker._mtp_load_from_disk = False + worker._mtp_speculative_enabled = True + worker._mtp_weights_from_refit = False worker.report_device_id_async = AsyncMock(return_value=["0"]) worker.llm = MagicMock() worker.llm.collective_rpc = AsyncMock(return_value=["node-0", "node-0"]) @@ -567,6 +569,7 @@ async def test_async_sparse_refit_post_init_records_worker_locality() -> None: ) assert worker.llm.collective_rpc.await_args_list == [ call("bind_numa", args=()), + call("configure_mtp_drafter_weight_source", args=(False,)), call("report_node_hostname", args=()), ] @@ -575,6 +578,8 @@ def test_sync_post_init_binds_numa() -> None: worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl) worker._sparse_refit_receiver = None worker._mtp_load_from_disk = False + worker._mtp_speculative_enabled = True + worker._mtp_weights_from_refit = False worker.report_device_id = MagicMock(return_value=["0"]) worker.llm = MagicMock() @@ -583,6 +588,7 @@ def test_sync_post_init_binds_numa() -> None: assert worker.vllm_device_ids == ["0"] assert worker.llm.collective_rpc.call_args_list == [ call("bind_numa", args=()), + call("configure_mtp_drafter_weight_source", args=(False,)), ] From 52900b8bfa5aab267d2c0e9c75fb20c5ac1bfdf1 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 00:59:04 -0700 Subject: [PATCH 44/68] fix(vllm): preserve static MTP drafters Signed-off-by: seonjinn --- .../generation/vllm/quantization/fp8.py | 4 ++-- .../models/generation/vllm/vllm_backend.py | 23 ++++++++++++------- nemo_rl/models/generation/vllm/vllm_worker.py | 12 ++++++++-- .../generation/vllm/vllm_worker_async.py | 5 ++++ .../models/generation/test_vllm_backend.py | 4 ++-- 5 files changed, 34 insertions(+), 14 deletions(-) diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index b3a335d642a..818d606d004 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -303,7 +303,7 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS) if num_first_layers_in_bf16 > 0 or num_last_layers_in_bf16 > 0: with init_empty_weights(): - model = AutoModel.from_config(config) + model = AutoModel.from_config(config, trust_remote_code=True) param_names = [name for name, _ in model.named_parameters()] get_text_config = getattr(config, "get_text_config", None) text_config = ( @@ -339,7 +339,7 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): ) if quantization_ignored_layer_kws: with init_empty_weights(): - model = AutoModel.from_config(config) + model = AutoModel.from_config(config, trust_remote_code=True) param_names = [ f"model.{name}".removesuffix(".weight").replace( "model.backbone.", "backbone." diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 67420c00b26..899630c2573 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -290,7 +290,11 @@ def _read_mtp_layer_weights_from_checkpoint( shard_to_names: dict[str, list[str]] = {} for name, shard in weight_map.items(): match = layer_re.search(name) - if match is not None and int(match.group(1)) in mtp_layer_indices: + is_mtp_namespace = "mtp" in name.split(".") + is_trailing_mtp_layer = ( + match is not None and int(match.group(1)) in mtp_layer_indices + ) + if is_mtp_namespace or is_trailing_mtp_layer: shard_to_names.setdefault(shard, []).append(name) weights: list[tuple[str, torch.Tensor]] = [] @@ -309,9 +313,9 @@ class VllmInternalWorkerExtension: # ones without probing, matching AbstractPolicyWorker.model_update_group. None and # not {} because a mutable class-level default is shared by every instance. pp_comm_groups: Optional[dict[int, Any]] = None - # True once the MTP drafter has been served by a one-time disk load (see - # load_mtp_weights_from_disk); refit then leaves those static weights alone. - _mtp_drafter_from_disk: bool = False + # False for a checkpoint-loaded static MTP drafter; True only when the + # trainer exports MTP weights in every policy refit stream. + _mtp_drafter_weights_from_refit: bool = True _sparse_delta_applier: Any = None _nrl_named_parameters: dict[str, torch.nn.Parameter] _nrl_layerwise_reload_active: bool = False @@ -665,6 +669,10 @@ def _get_drafter_model(self) -> Any: draft_owner = getattr(self.model_runner, "drafter", None) return getattr(draft_owner, "model", None) if draft_owner else None + def configure_mtp_drafter_weight_source(self, weights_from_refit: bool) -> None: + """Record whether the trainer owns and refreshes the MTP weights.""" + self._mtp_drafter_weights_from_refit = weights_from_refit + def _load_draft_weights( self, draft_weights: list[tuple[str, torch.Tensor]] ) -> None: @@ -693,7 +701,7 @@ def _mtp_drafter_refit_enabled(self) -> bool: does not co-train the MTP layer — to avoid clobbering and re-processing those static weights. """ - if self._mtp_drafter_from_disk: + if not self._mtp_drafter_weights_from_refit: return False spec_config = getattr(self.model_runner.vllm_config, "speculative_config", None) method = getattr(spec_config, "method", None) if spec_config else None @@ -817,9 +825,8 @@ def load_mtp_weights_from_disk(self, model_path: str) -> bool: process_weights_after_loading( draft_model, draft_model_config, self.device ) - # Mark that the MTP drafter is served from a one-time disk load so refit - # does not re-load or re-process these static weights. - self._mtp_drafter_from_disk = True + # This drafter is served by the checkpoint rather than the refit stream. + self._mtp_drafter_weights_from_refit = False logger.info( "[mtp] Loaded MTP draft weights for layers %s from %s", sorted(mtp_layer_indices), diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index ba9ad5fb837..82155d63980 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -562,10 +562,13 @@ def _load_model(self, bundle_indices, seed): # (see VllmInternalWorkerExtension.load_mtp_weights_from_disk). spec_cfg = vllm_kwargs.get("speculative_config") mtp_weights_from_refit = bool(self.cfg.get("_mtp_weights_from_refit")) + self._mtp_speculative_enabled = spec_cfg is not None and spec_cfg.get( + "method" + ) in ("deepseek_mtp", "mtp") + self._mtp_weights_from_refit = mtp_weights_from_refit self._mtp_load_from_disk: bool = ( load_format == "dummy" - and spec_cfg is not None - and spec_cfg.get("method") in ("deepseek_mtp", "mtp") + and self._mtp_speculative_enabled and not mtp_weights_from_refit ) @@ -934,6 +937,11 @@ def post_init(self): if self.llm is not None: self.llm.collective_rpc("bind_numa", args=tuple()) self.vllm_device_ids = self.report_device_id() + if self._mtp_speculative_enabled: + self.llm.collective_rpc( + "configure_mtp_drafter_weight_source", + args=(self._mtp_weights_from_refit,), + ) if self._mtp_load_from_disk: self.llm.collective_rpc( "load_mtp_weights_from_disk", args=(self.model_name,) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 01076d19481..aeefd7e9de9 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -332,6 +332,11 @@ async def post_init_async(self): if self.llm is not None: await self.llm.collective_rpc("bind_numa", args=tuple()) self.vllm_device_ids = await self.report_device_id_async() + if self._mtp_speculative_enabled: + await self.llm.collective_rpc( + "configure_mtp_drafter_weight_source", + args=(self._mtp_weights_from_refit,), + ) if self._mtp_load_from_disk: await self.llm.collective_rpc( "load_mtp_weights_from_disk", args=(self.model_name,) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index d8c688411a6..12c3d61aa7e 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -106,7 +106,7 @@ def _make_mtp_refit_extension( ext = VllmInternalWorkerExtension.__new__(VllmInternalWorkerExtension) ext.device = torch.device("cpu") - ext._mtp_drafter_from_disk = from_disk + ext._mtp_drafter_weights_from_refit = not from_disk spec_config = ( None @@ -1065,7 +1065,7 @@ def process_weights_after_loading(model, model_config, device): ) ext, expected_state_info = _make_collective_update_extension(vllm_backend) if with_mtp: - ext._mtp_drafter_from_disk = False + ext._mtp_drafter_weights_from_refit = True ext.model_runner.drafter = SimpleNamespace(model=draft_model) ext.model_runner.vllm_config = SimpleNamespace( speculative_config=SimpleNamespace( From 347b24d53df72e16ac961b3800f14ee3a06ad38f Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 01:07:36 -0700 Subject: [PATCH 45/68] test(vllm): accept remote-code model construction Signed-off-by: seonjinn --- tests/unit/models/generation/test_vllm_fp8_quantization.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index a95e5017c0b..b0149fde7d9 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -637,7 +637,9 @@ def named_parameters(self): "from_pretrained", lambda *_args, **_kwargs: types.SimpleNamespace(num_hidden_layers=4), ) - monkeypatch.setattr(fp8.AutoModel, "from_config", lambda *_args: FakeModel()) + monkeypatch.setattr( + fp8.AutoModel, "from_config", lambda *_args, **_kwargs: FakeModel() + ) monkeypatch.setattr(fp8, "monkey_patch_vllm_ray_executor", lambda _config: None) with pytest.warns( From 830a6ca6b0237426aca9787b1c92d8101fdb918d Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 05:00:04 -0700 Subject: [PATCH 46/68] fix(mxfp8): validate mixed layer boundaries Signed-off-by: seonjinn --- .../generation/vllm/quantization/fp8.py | 24 ++-- .../generation/test_vllm_fp8_quantization.py | 118 ++++++++++++++++-- 2 files changed, 122 insertions(+), 20 deletions(-) diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index 818d606d004..7c06d2d785e 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -297,6 +297,23 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): # create fp8 kwargs for vllm's LLM(...) num_first_layers_in_bf16 = vllm_cfg.get("num_first_layers_in_bf16", 0) num_last_layers_in_bf16 = vllm_cfg.get("num_last_layers_in_bf16", 0) + get_text_config = getattr(config, "get_text_config", None) + text_config = ( + get_text_config() + if callable(get_text_config) + else getattr(config, "text_config", config) + ) + num_hidden_layers = text_config.num_hidden_layers + for field_name, value in ( + ("num_first_layers_in_bf16", num_first_layers_in_bf16), + ("num_last_layers_in_bf16", num_last_layers_in_bf16), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError(f"{field_name} must be an integer") + if not 0 <= value <= num_hidden_layers: + raise ValueError( + f"{field_name} must be between 0 and {num_hidden_layers}, got {value}" + ) if global_fp8_config.is_mx: fp8_block_quant_kwargs = dict(MXFP8_BLOCK_QUANT_KWARGS) else: @@ -305,13 +322,6 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): with init_empty_weights(): model = AutoModel.from_config(config, trust_remote_code=True) param_names = [name for name, _ in model.named_parameters()] - get_text_config = getattr(config, "get_text_config", None) - text_config = ( - get_text_config() - if callable(get_text_config) - else getattr(config, "text_config", config) - ) - num_hidden_layers = text_config.num_hidden_layers bf16_params = [] if num_first_layers_in_bf16 > 0: diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index b0149fde7d9..baa3acc3bc2 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -157,11 +157,17 @@ def test_init_fp8_passes_modelopt_ignore_patterns_without_hf_expansion( ("num_first_layers_in_bf16", "num_last_layers_in_bf16"), [ (0, 0), + (5, 0), + (0, 4), (1, 1), (2, 6), + (3, 5), + (1, 3), (7, 3), (26, 26), (30, 30), + (40, 0), + (0, 40), ], ) @pytest.mark.parametrize( @@ -170,16 +176,20 @@ def test_init_fp8_passes_modelopt_ignore_patterns_without_hf_expansion( pytest.param( types.SimpleNamespace( model_name="dummy-qwen-model", - hf_layer_prefix="layers", - raw_layer_prefix="model.layers", - mapper_prefixes={}, + num_hidden_layers=40, + nested_text_config=True, + hf_layer_prefix="language_model.layers", + raw_layer_prefix="model.language_model.layers", + vllm_layer_prefix="language_model.model.layers", + mapper_prefixes={"model.language_model.": "language_model.model."}, hf_target_suffixes=( "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", - "mlp.experts.gate_proj", - "mlp.experts.up_proj", + ), + suffixless_hf_target_suffixes=( + "mlp.experts.gate_up_proj", "mlp.experts.down_proj", ), vllm_target_suffixes=( @@ -187,10 +197,15 @@ def test_init_fp8_passes_modelopt_ignore_patterns_without_hf_expansion( "self_attn.o_proj", "mlp.experts", ), - non_target_suffixes=("mlp.gate", "mlp.shared_experts.up_proj"), + non_target_suffixes=( + "mlp.gate", + "mlp.shared_expert.up_proj", + "mlp.shared_expert_gate", + ), ignore_patterns=( "*layers.*.mlp.gate", - "*layers.*.mlp.shared_experts.*", + "*layers.*.mlp.shared_expert.*", + "*layers.*.mlp.shared_expert_gate", "lm_head", ), ), @@ -199,8 +214,11 @@ def test_init_fp8_passes_modelopt_ignore_patterns_without_hf_expansion( pytest.param( types.SimpleNamespace( model_name="dummy-nemotron-h-model", + num_hidden_layers=52, + nested_text_config=False, hf_layer_prefix="backbone.layers", raw_layer_prefix="backbone.layers", + vllm_layer_prefix="model.layers", mapper_prefixes={"backbone": "model"}, hf_target_suffixes=( "mixer.q_proj", @@ -210,6 +228,7 @@ def test_init_fp8_passes_modelopt_ignore_patterns_without_hf_expansion( "mixer.experts.up_proj", "mixer.experts.down_proj", ), + suffixless_hf_target_suffixes=(), vllm_target_suffixes=( "mixer.qkv_proj", "mixer.o_proj", @@ -241,7 +260,7 @@ def test_init_fp8_keeps_mixed_recipe_boundary_targets_in_bf16( from vllm.model_executor.models.utils import WeightsMapper fp8 = fp8_module - num_hidden_layers = 52 + num_hidden_layers = recipe_case.num_hidden_layers param_names = [] for layer_idx in range(num_hidden_layers): param_names.extend( @@ -251,12 +270,20 @@ def test_init_fp8_keeps_mixed_recipe_boundary_targets_in_bf16( *recipe_case.non_target_suffixes, ) ) + param_names.extend( + f"{recipe_case.hf_layer_prefix}.{layer_idx}.{suffix}" + for suffix in recipe_case.suffixless_hf_target_suffixes + ) monkeypatch.setattr( fp8.AutoConfig, "from_pretrained", - lambda *_args, **_kwargs: types.SimpleNamespace( - num_hidden_layers=num_hidden_layers + lambda *_args, **_kwargs: ( + types.SimpleNamespace( + text_config=types.SimpleNamespace(num_hidden_layers=num_hidden_layers) + ) + if recipe_case.nested_text_config + else types.SimpleNamespace(num_hidden_layers=num_hidden_layers) ), ) monkeypatch.setattr( @@ -300,20 +327,85 @@ def test_init_fp8_keeps_mixed_recipe_boundary_targets_in_bf16( for suffix in recipe_case.hf_target_suffixes: module_name = f"{recipe_case.raw_layer_prefix}.{layer_idx}.{suffix}" assert module_name in quant_config["ignored_layers"] + for suffix in recipe_case.suffixless_hf_target_suffixes: + module_name = f"{recipe_case.raw_layer_prefix}.{layer_idx}.{suffix}" + assert module_name in quant_config["ignored_layers"] for suffix in recipe_case.vllm_target_suffixes: - module_name = f"model.layers.{layer_idx}.{suffix}" + module_name = f"{recipe_case.vllm_layer_prefix}.{layer_idx}.{suffix}" assert modelopt_config.is_layer_excluded(module_name) else: for suffix in recipe_case.vllm_target_suffixes: - module_name = f"model.layers.{layer_idx}.{suffix}" + module_name = f"{recipe_case.vllm_layer_prefix}.{layer_idx}.{suffix}" assert not modelopt_config.is_layer_excluded(module_name) for suffix in recipe_case.non_target_suffixes: - module_name = f"model.layers.{layer_idx}.{suffix}" + module_name = f"{recipe_case.vllm_layer_prefix}.{layer_idx}.{suffix}" assert modelopt_config.is_layer_excluded(module_name) +@pytest.mark.parametrize( + "field_name", + ["num_first_layers_in_bf16", "num_last_layers_in_bf16"], +) +@pytest.mark.parametrize( + ("invalid_value", "error_match"), + [ + pytest.param(-1, "must be between 0 and 40", id="negative"), + pytest.param(41, "must be between 0 and 40", id="too-large"), + pytest.param(1.0, "must be an integer", id="float"), + pytest.param("1", "must be an integer", id="string"), + pytest.param(True, "must be an integer", id="bool"), + ], +) +def test_init_fp8_rejects_invalid_bf16_layer_boundaries( + fp8_module, + monkeypatch, + field_name, + invalid_value, + error_match, +): + fp8 = fp8_module + num_hidden_layers = 40 + config = types.SimpleNamespace( + get_text_config=lambda: types.SimpleNamespace( + num_hidden_layers=num_hidden_layers + ) + ) + monkeypatch.setattr( + fp8.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: config, + ) + monkeypatch.setattr( + fp8.AutoModel, + "from_config", + lambda *_args, **_kwargs: types.SimpleNamespace( + named_parameters=lambda: [ + ( + f"language_model.layers.{layer_idx}.mlp.experts.gate_up_proj", + None, + ) + for layer_idx in range(num_hidden_layers) + ] + ), + ) + monkeypatch.setattr(fp8, "monkey_patch_vllm_ray_executor", lambda _config: None) + + vllm_cfg = { + "precision": "fp8", + "kv_cache_dtype": "auto", + "async_engine": False, + "is_mx": True, + "num_first_layers_in_bf16": 0, + "num_last_layers_in_bf16": 0, + } + vllm_cfg[field_name] = invalid_value + + with pytest.raises(ValueError, match=error_match): + fp8.init_fp8(vllm_cfg, "dummy-qwen-model", model_parallel_size=1) + + def test_init_fp8_reads_layer_count_from_text_config(fp8_module, monkeypatch): fp8 = fp8_module num_hidden_layers = 8 From 5b14cad6f45fd0d6332f4d253c3821ab86495429 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 05:10:33 -0700 Subject: [PATCH 47/68] fix(mxfp8): validate scope before patching workers Signed-off-by: seonjinn --- .../generation/vllm/quantization/fp8.py | 44 ++++++++++--------- .../generation/test_vllm_fp8_quantization.py | 10 ++++- 2 files changed, 32 insertions(+), 22 deletions(-) diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index 7c06d2d785e..3fd885bdd25 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -256,9 +256,30 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): quantization_ignore_patterns = [ pattern.strip() for pattern in quantization_ignore_patterns ] + + num_first_layers_in_bf16 = vllm_cfg.get("num_first_layers_in_bf16", 0) + num_last_layers_in_bf16 = vllm_cfg.get("num_last_layers_in_bf16", 0) + get_text_config = getattr(config, "get_text_config", None) + text_config = ( + get_text_config() + if callable(get_text_config) + else getattr(config, "text_config", config) + ) + num_hidden_layers = text_config.num_hidden_layers + for field_name, value in ( + ("num_first_layers_in_bf16", num_first_layers_in_bf16), + ("num_last_layers_in_bf16", num_last_layers_in_bf16), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError(f"{field_name} must be an integer") + if not 0 <= value <= num_hidden_layers: + raise ValueError( + f"{field_name} must be between 0 and {num_hidden_layers}, got {value}" + ) + fp8_config_kwargs = { - "num_first_layers_in_bf16": vllm_cfg.get("num_first_layers_in_bf16", 0), - "num_last_layers_in_bf16": vllm_cfg.get("num_last_layers_in_bf16", 0), + "num_first_layers_in_bf16": num_first_layers_in_bf16, + "num_last_layers_in_bf16": num_last_layers_in_bf16, "model_parallel_size": model_parallel_size, "kv_cache_dtype": kv_cache_dtype, "use_fp8_weights": use_fp8_weights, @@ -295,25 +316,6 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): monkey_patch_vllm_ray_executor(global_fp8_config) # create fp8 kwargs for vllm's LLM(...) - num_first_layers_in_bf16 = vllm_cfg.get("num_first_layers_in_bf16", 0) - num_last_layers_in_bf16 = vllm_cfg.get("num_last_layers_in_bf16", 0) - get_text_config = getattr(config, "get_text_config", None) - text_config = ( - get_text_config() - if callable(get_text_config) - else getattr(config, "text_config", config) - ) - num_hidden_layers = text_config.num_hidden_layers - for field_name, value in ( - ("num_first_layers_in_bf16", num_first_layers_in_bf16), - ("num_last_layers_in_bf16", num_last_layers_in_bf16), - ): - if not isinstance(value, int) or isinstance(value, bool): - raise ValueError(f"{field_name} must be an integer") - if not 0 <= value <= num_hidden_layers: - raise ValueError( - f"{field_name} must be between 0 and {num_hidden_layers}, got {value}" - ) if global_fp8_config.is_mx: fp8_block_quant_kwargs = dict(MXFP8_BLOCK_QUANT_KWARGS) else: diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index baa3acc3bc2..ebf9390826c 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -390,7 +390,12 @@ def test_init_fp8_rejects_invalid_bf16_layer_boundaries( ] ), ) - monkeypatch.setattr(fp8, "monkey_patch_vllm_ray_executor", lambda _config: None) + patch_calls = [] + monkeypatch.setattr( + fp8, + "monkey_patch_vllm_ray_executor", + lambda config: patch_calls.append(config), + ) vllm_cfg = { "precision": "fp8", @@ -405,6 +410,9 @@ def test_init_fp8_rejects_invalid_bf16_layer_boundaries( with pytest.raises(ValueError, match=error_match): fp8.init_fp8(vllm_cfg, "dummy-qwen-model", model_parallel_size=1) + assert patch_calls == [] + assert fp8.global_fp8_config is None + def test_init_fp8_reads_layer_count_from_text_config(fp8_module, monkeypatch): fp8 = fp8_module From 96101df6434b6e5b7f4de684f6679557b99ed601 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 10:55:18 -0700 Subject: [PATCH 48/68] docs(fp8): explain mixed MXFP8 rollout scope Signed-off-by: seonjinn --- docs/fp8.md | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/docs/fp8.md b/docs/fp8.md index dccf6da2932..2e1fe96fab1 100644 --- a/docs/fp8.md +++ b/docs/fp8.md @@ -73,6 +73,43 @@ attention, the router, and the language-model head in BF16: - model.layers.*.mlp.gate ``` +`num_first_layers_in_bf16` and `num_last_layers_in_bf16` keep complete +transformer layers in BF16. NeMo RL reads the model configuration and parameter +names, so these options do not require a model-specific `model.layers` or +`backbone.layers` prefix. The patterns above still apply to every middle layer. + +For example, this Nemotron 3.5 Lightning scope keeps the first two and last six +layers in BF16. In the middle layers it quantizes only the non-shared routed +experts. Attention, Mamba projections, routers, shared experts, latent +projections, and MTP remain in BF16: + +```yaml +policy: + generation: + vllm_cfg: + precision: fp8 + is_mx: true + num_first_layers_in_bf16: 2 + num_last_layers_in_bf16: 6 + quantization_ignore_patterns: + - "*layers.*.mixer.qkv_proj" + - "*layers.*.mixer.o_proj" + - "*layers.*.mixer.in_proj" + - "*layers.*.mixer.out_proj" + - "*layers.*.mixer.up_proj" + - "*layers.*.mixer.down_proj" + - "*layers.*.mixer.gate" + - "*layers.*.mixer.shared_experts.*" + - "*layers.*.mixer.fc1_latent_proj" + - "*layers.*.mixer.fc2_latent_proj" + - "*mtp.*" +``` + +To quantize QKVO as well as the routed experts, remove the `qkv_proj` and +`o_proj` entries. The first two and last six layers still remain entirely in +BF16. Check the logged effective ignore list when adding a new model family; +an ignore pattern that matches no module is usually a naming error. + `lm_head` is always excluded from FP8 and MXFP8 quantization, even when it is not listed in `quantization_ignore_patterns` in the YAML configuration. Models with MTP layers must list their MTP module names explicitly, for example From 8eddc45aa00a1ba48a8e4f978a858f08a077d47f Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 15:40:23 -0700 Subject: [PATCH 49/68] chore: minimize Qwen3.5 TRTLLM recipe Signed-off-by: seonjinn --- .../grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml | 9 --------- .../models/generation/test_qwen35_bf16_trtllm_recipe.py | 1 - 2 files changed, 10 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml index 674faceff72..f97971913cc 100644 --- a/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml @@ -1,19 +1,14 @@ defaults: grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml - loss_fn: use_importance_sampling_correction: true truncated_importance_sampling_type: tis truncated_importance_sampling_ratio: 2 - grpo: async_grpo: enabled: true - max_trajectory_age_steps: 1 in_flight_weight_updates: true - checkpointing: checkpoint_dir: results/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm - policy: generation: refit_transport: nccl_reshard @@ -25,19 +20,15 @@ policy: vllm_cfg: async_engine: true precision: bfloat16 - tensor_parallel_size: 4 - pipeline_parallel_size: 1 expert_parallel_size: 4 gpu_memory_utilization: 0.8 enforce_eager: false vllm_kwargs: moe_backend: flashinfer_trtllm expert_placement_strategy: linear - logger: wandb: name: grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm - cluster: gpus_per_node: 4 num_nodes: 6 diff --git a/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py b/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py index f4483cbbe61..b1f542efc23 100644 --- a/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py +++ b/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py @@ -22,7 +22,6 @@ register_omegaconf_resolvers, ) - PROJECT_ROOT = Path(__file__).resolve().parents[4] RECIPE_NAME = "grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml" From f20709b35feb6e27d0ddda38899ec1ab2901d5af Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 4 Sep 2026 16:53:45 -0700 Subject: [PATCH 50/68] test(vllm): update FP8 refit loader stub Signed-off-by: seonjinn --- tests/unit/models/generation/test_vllm_refit_loader.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/models/generation/test_vllm_refit_loader.py b/tests/unit/models/generation/test_vllm_refit_loader.py index baa769e91a2..a21942baf09 100644 --- a/tests/unit/models/generation/test_vllm_refit_loader.py +++ b/tests/unit/models/generation/test_vllm_refit_loader.py @@ -212,8 +212,9 @@ def test_checkpoint_refit_preserves_nonsharded_fp8_path(monkeypatch): ) monkeypatch.setattr(fp8, "is_fp8_model", lambda _config: True) - def load_fp8_weights(weights, model_runner): + def load_fp8_weights(weights, model_runner, *, model_load_weights): assert model_runner is ext.model_runner + assert model_load_weights == ext._load_full_hf_weights loaded.extend(weights) monkeypatch.setattr(fp8, "load_weights", load_fp8_weights) From d8e8088eaf9f4ed294d275c4f5a5bb89b4720a72 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 5 Sep 2026 18:30:22 -0700 Subject: [PATCH 51/68] test(refit): reproduce mixed TRTLLM init ordering Signed-off-by: seonjinn --- .../models/generation/test_nccl_reshard_backend.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index da46546cd8e..a6101619ff4 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -788,6 +788,8 @@ def test_legacy_refit_map_is_built_after_comm_groups_exist(monkeypatch): vllm_backend.VllmInternalWorkerExtension ) ext.device = torch.device("cpu") + ext.pp_comm_groups = None + ext._uses_unquantized_flashinfer_trtllm = lambda: True ext._validate_native_layerwise_refit = MagicMock() expected_map = HFToLocalParamMap() ext.build_hf_to_local_param_map = MagicMock(return_value=expected_map) @@ -812,6 +814,10 @@ def init_nccl_communicator(self, *, device): monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + ext.prepare_nccl_reshard_refit_info(refit_info) + + ext.build_hf_to_local_param_map.assert_not_called() + ext.init_nccl_reshard_comm_group( rank_prefix=0, pp_ips=["127.0.0.1"], @@ -822,10 +828,6 @@ def init_nccl_communicator(self, *, device): ) assert ext.pp_comm_groups[0].rank == 8 - assert ext.build_hf_to_local_param_map.call_count == 0 - - ext.prepare_nccl_reshard_refit_info(refit_info) - ext.build_hf_to_local_param_map.assert_called_once_with(refit_info) assert ext.hf_to_local_param_map is expected_map From 9326f41a65dd59689df4ca5c0480cfca6630d113 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 5 Sep 2026 18:36:43 -0700 Subject: [PATCH 52/68] fix(refit): build TRTLLM maps after communicator init Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 899630c2573..b4c00e0207a 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -518,6 +518,15 @@ def init_nccl_reshard_comm_group( previous.abort, RELEASE_GRACE_S, "a previous reshard bulk communicator" ) + refit_info = getattr(self, "nccl_reshard_refit_info", None) + if ( + refit_info is not None + and self._uses_unquantized_flashinfer_trtllm() + ): + # TRTLLM expert destinations depend on this worker's rank in each + # per-PP-stage group, so they cannot be mapped during prepare. + self.hf_to_local_param_map = self.build_hf_to_local_param_map(refit_info) + def report_device_id(self) -> str: """Retrieve the UUID of the current CUDA device.""" from nemo_rl.utils.nvml import get_device_uuid @@ -1246,11 +1255,14 @@ def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: self.nccl_reshard_refit_info = ( # pyrefly: ignore[implicitly-defined-attribute] restore_refit_info_placements(refit_info) ) - # Build HFToLocalParamMap after the communicator setup performed by the - # synchronizer, since TRTLLM expert destinations depend on its rank. - self.hf_to_local_param_map = self.build_hf_to_local_param_map( # pyrefly: ignore[implicitly-defined-attribute] - self.nccl_reshard_refit_info - ) + if self._uses_unquantized_flashinfer_trtllm(): + # The TRTLLM expert map needs the per-PP-stage communicator ranks, + # which init_nccl_reshard_comm_group establishes after prepare. + self.hf_to_local_param_map = HFToLocalParamMap() + else: + self.hf_to_local_param_map = self.build_hf_to_local_param_map( + self.nccl_reshard_refit_info + ) def build_hf_to_local_param_map(self, refit_info: dict) -> HFToLocalParamMap: """Build the vLLM-backend ``hf_to_local_param_map`` (HFToLocalParamMap). @@ -1674,6 +1686,13 @@ def nccl_reshard_refit(self, refit_timeout_s: float | None = None) -> bool: Both communicator families are handed to the watchdog -- the per-PP-stage bulk groups and the shared model_update_group -- because the transfer uses them in sequence and a hang can be in either. + Each HF param's ``LocalParamSpec`` (from ``hf_to_local_param_map``, built + during prepare or after PP communicator setup for TRTLLM experts) provides + the dst buffer: + for a direct param xferdtensor receives straight into the live vLLM + param (no hooks); for a merged param (dense gate_up_proj, grouped w13) + ``pre`` allocates a temp recv buffer and ``post`` copies the TP-local + slice back into the live merged param. """ from nemo_rl.distributed.refit_watchdog import ( RefitAborted, From d3f3f06cd641257ea6a2b55009da1918cb267284 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 5 Sep 2026 19:10:30 -0700 Subject: [PATCH 53/68] test(refit): cover both reshard setup orders Signed-off-by: seonjinn --- .../generation/test_nccl_reshard_backend.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/unit/models/generation/test_nccl_reshard_backend.py b/tests/unit/models/generation/test_nccl_reshard_backend.py index a6101619ff4..5bce7c3e5c6 100644 --- a/tests/unit/models/generation/test_nccl_reshard_backend.py +++ b/tests/unit/models/generation/test_nccl_reshard_backend.py @@ -781,7 +781,10 @@ def test_nccl_reshard_trtllm_refit_rejects_fp8_kv_cache(monkeypatch): ext._validate_native_layerwise_refit("nccl_reshard") -def test_legacy_refit_map_is_built_after_comm_groups_exist(monkeypatch): +@pytest.mark.parametrize("prepare_before_comm", [False, True]) +def test_legacy_refit_map_is_built_after_comm_groups_exist( + monkeypatch, prepare_before_comm +): from nemo_rl.models.generation.vllm import vllm_backend ext = vllm_backend.VllmInternalWorkerExtension.__new__( @@ -814,9 +817,9 @@ def init_nccl_communicator(self, *, device): monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) - ext.prepare_nccl_reshard_refit_info(refit_info) - - ext.build_hf_to_local_param_map.assert_not_called() + if prepare_before_comm: + ext.prepare_nccl_reshard_refit_info(refit_info) + ext.build_hf_to_local_param_map.assert_not_called() ext.init_nccl_reshard_comm_group( rank_prefix=0, @@ -828,6 +831,10 @@ def init_nccl_communicator(self, *, device): ) assert ext.pp_comm_groups[0].rank == 8 + if not prepare_before_comm: + ext.build_hf_to_local_param_map.assert_not_called() + ext.prepare_nccl_reshard_refit_info(refit_info) + ext.build_hf_to_local_param_map.assert_called_once_with(refit_info) assert ext.hf_to_local_param_map is expected_map From c2f8c1feceb8a6959e5e43a2c7c712cebca3eb37 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 5 Sep 2026 19:22:43 -0700 Subject: [PATCH 54/68] fix(refit): support both reshard setup orders Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/vllm_backend.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index b4c00e0207a..673ad489593 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -1255,7 +1255,10 @@ def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: self.nccl_reshard_refit_info = ( # pyrefly: ignore[implicitly-defined-attribute] restore_refit_info_placements(refit_info) ) - if self._uses_unquantized_flashinfer_trtllm(): + if ( + self._uses_unquantized_flashinfer_trtllm() + and not self.pp_comm_groups + ): # The TRTLLM expert map needs the per-PP-stage communicator ranks, # which init_nccl_reshard_comm_group establishes after prepare. self.hf_to_local_param_map = HFToLocalParamMap() From ad6d0dc247b9bd61e3f15c8a4302541e1776e600 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sun, 6 Sep 2026 14:00:12 -0700 Subject: [PATCH 55/68] test(vllm): cover FP8 layerwise generator loading Signed-off-by: seonjinn --- .../models/generation/test_vllm_backend.py | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 07255ced1cc..879063b6aa7 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -532,6 +532,50 @@ def test_fp8_load_uses_buffer_safe_model_loader(monkeypatch): assert model_load_weights.__func__ is ext._load_full_hf_weights.__func__ +@pytest.mark.vllm +def test_fp8_layerwise_reload_passes_entire_quantized_generator(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + received_weights = [] + + def model_load_weights(*, weights): + received_weights.extend(weights) + return {name for name, _ in received_weights} + + model = SimpleNamespace(load_weights=model_load_weights) + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + ext.model_runner = SimpleNamespace(model=model, vllm_config=object()) + ext._nrl_layerwise_reload_active = True + source_weights = [("model.weight", torch.ones(2))] + quantized_weights = [ + ("model.weight", torch.ones(2, dtype=torch.float8_e4m3fn)), + ("model.weight_scale", torch.ones(1)), + ] + + def get_quantized_weight_iterator( + weights, model_runner, *, refit_with_reload_api + ): + assert weights is source_weights + assert model_runner is ext.model_runner + assert refit_with_reload_api is False + yield from quantized_weights + + monkeypatch.setattr(fp8, "is_fp8_model", lambda _config: True) + monkeypatch.setattr( + fp8, "get_quantized_weight_iterator", get_quantized_weight_iterator + ) + monkeypatch.setattr( + vllm_backend, "_detach_pending_layerwise_weights", lambda *_args: None + ) + + ext._load_hf_weights(source_weights) + + assert received_weights == quantized_weights + + @pytest.mark.vllm def test_layerwise_reload_preserves_deferred_weight_across_buffer_reuse(monkeypatch): from vllm.model_executor.model_loader.reload import record_metadata_for_reloading From 78d066d821a31a5ab507e25a75f9bd591d9d27a4 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sun, 6 Sep 2026 14:02:48 -0700 Subject: [PATCH 56/68] docs(refit): document BF16 TRTLLM limits Signed-off-by: seonjinn --- docs/design-docs/nccl-reshard-refit.md | 24 ++++++++++++------- ...=> test_vllm_qwen35_bf16_trtllm_recipe.py} | 0 2 files changed, 15 insertions(+), 9 deletions(-) rename tests/unit/models/generation/{test_qwen35_bf16_trtllm_recipe.py => test_vllm_qwen35_bf16_trtllm_recipe.py} (100%) diff --git a/docs/design-docs/nccl-reshard-refit.md b/docs/design-docs/nccl-reshard-refit.md index 8603de95e2d..3f027554b1a 100644 --- a/docs/design-docs/nccl-reshard-refit.md +++ b/docs/design-docs/nccl-reshard-refit.md @@ -42,7 +42,9 @@ single `ValueError` listing every violation. The current requirements are: * BF16 FlashInfer TRTLLM MoE is supported through vLLM's native layerwise-reload path. Its grouped expert weights must use expert-parallel destination sharding with linear expert placement; tensor-sharded expert - destinations and round-robin placement are rejected. + destinations and round-robin placement are rejected. This path does not + support an FP8 KV cache or a co-trained MTP drafter; setup rejects both + combinations. * vLLM expert parallelism is supported with the NeMo RL convention `expert_parallel_size == tensor_parallel_size`. * Generation-side, PP > 1 is not supported. @@ -71,18 +73,22 @@ nccl-reshard-refit implementation: Two FFN-named groups are explicitly excluded and ride the misc path instead: shared-expert weights (`*.shared_expert.*`, which fuse differently on the vLLM side) and co-trained MTP drafter weights (which vLLM keeps in a separate - drafter module updated through `load_weights`). MTP weights are recognized two - ways: bare-`mtp.`-prefix HF names (NemotronH, Qwen3.5) via + drafter module updated through `load_weights`). Co-trained MTP is not supported + with BF16 FlashInfer TRTLLM; this routing applies to other supported backend + combinations. MTP weights are recognized two ways: bare-`mtp.`-prefix HF names + (NemotronH, Qwen3.5) via `is_nccl_reshard_param()`, and DeepSeek-style MTP exported as trailing `model.layers.N` indices via provenance — the Megatron-side name carries an `mtp.` module segment (bare for LM bridges, `language_model.mtp.*` for the VL and EXAONE bridges), so the worker excludes those HF layers when building the metadata (`_collect_mtp_hf_layer_names()`). * **Misc path** — everything else (embeddings, attention projections, layernorms, the - MoE router, `lm_head`, FP8 `_scale_inv` siblings, FP8 KV-cache scales, …). These ride - a packed broadcast (conventional `packed_tensor.py` implementation) over the shared - `model_update_group` and are loaded on the generation side through the backend's - regular `load_weights` machinery. + MoE router, `lm_head`, FP8 `_scale_inv` siblings, FP8 KV-cache scales, …). FP8 + KV-cache scales are supported only by backend combinations that allow an FP8 KV + cache; BF16 FlashInfer TRTLLM rejects that configuration at setup. These tensors + ride a packed broadcast (conventional `packed_tensor.py` implementation) over the + shared `model_update_group` and are loaded on the generation side through the + backend's regular `load_weights` machinery. The feature is integrated into the `nemo_rl/weight_sync/` framework: `create_weight_synchronizer(..., nccl_reshard_refit=True)` returns a @@ -96,8 +102,8 @@ training starts: 1. **`init_collective()`** — creates the `model_update_group`, a NCCL group spanning all training and generation ranks. The bulk path does not use it; it carries the misc - packed-broadcast (and FP8 KV-cache scales), identical to the conventional collective - transport. + packed-broadcast, including FP8 KV-cache scales for backend combinations that support + them, identical to the conventional collective transport. 2. **`init_nccl_reshard_comm_group()`** — creates the bulk-path communicator(s): **one NCCL group per training PP stage**, each spanning that stage's training ranks plus *all* generation ranks (non-PP is simply `pp_size == 1`, a single group over diff --git a/tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py b/tests/unit/models/generation/test_vllm_qwen35_bf16_trtllm_recipe.py similarity index 100% rename from tests/unit/models/generation/test_qwen35_bf16_trtllm_recipe.py rename to tests/unit/models/generation/test_vllm_qwen35_bf16_trtllm_recipe.py From 08d86876b023a9480be208719950383b55cd93b1 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sun, 6 Sep 2026 14:18:48 -0700 Subject: [PATCH 57/68] fix(vllm): preserve FP8 reload iterator contents Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/vllm_backend.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 323f665e3ae..e76b7f97c64 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -334,7 +334,7 @@ def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: return params def _load_full_hf_weights( - self, policy_weights: list[tuple[str, torch.Tensor]] + self, policy_weights: Iterable[tuple[str, torch.Tensor]] ) -> set[str] | None: """Load HF weights and detach any deferred reload tensors from transport storage. @@ -344,12 +344,16 @@ def _load_full_hf_weights( if not getattr(self, "_nrl_layerwise_reload_active", False): return self.model_runner.model.load_weights(weights=policy_weights) - source_storage_ptrs = { - tensor.untyped_storage().data_ptr() for _, tensor in policy_weights - } + source_storage_ptrs = set() + + def track_source_storage() -> Iterator[tuple[str, torch.Tensor]]: + for name, tensor in policy_weights: + source_storage_ptrs.add(tensor.untyped_storage().data_ptr()) + yield name, tensor + load_error: Exception | None = None try: - return self.model_runner.model.load_weights(weights=policy_weights) + return self.model_runner.model.load_weights(weights=track_source_storage()) except Exception as error: load_error = error raise From 38deaed6c7a167ca05aa8db10cd9877e9b0d484f Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sun, 6 Sep 2026 14:23:22 -0700 Subject: [PATCH 58/68] chore(vllm): type tracked reload storage pointers Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/vllm_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index e76b7f97c64..5c9f3a1003c 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -344,7 +344,7 @@ def _load_full_hf_weights( if not getattr(self, "_nrl_layerwise_reload_active", False): return self.model_runner.model.load_weights(weights=policy_weights) - source_storage_ptrs = set() + source_storage_ptrs: set[int] = set() def track_source_storage() -> Iterator[tuple[str, torch.Tensor]]: for name, tensor in policy_weights: From 7ecfc32b5a9a8b4c1c5958b0b1271a13a3839166 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sun, 6 Sep 2026 15:38:10 -0700 Subject: [PATCH 59/68] test(qwen3.5): exempt long model setup from idle reaper Signed-off-by: seonjinn --- .../llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh index 41627dd1cf0..98847fc7a39 100755 --- a/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh +++ b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh @@ -9,6 +9,7 @@ STEPS_PER_RUN=20 MAX_STEPS=20 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) NUM_MINUTES=240 +JOB_REAPER_COMMENT='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"120","reason":"disproportionate_resource_requirement","description":"6-node GRPO with vLLM has long GPU-idle phases during Ray init and model loading"}}' # ===== END CONFIG ===== exit_if_max_steps_reached From eb7529fc99abdd5ea9fceec1b98433a26989d1b4 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sun, 6 Sep 2026 15:40:30 -0700 Subject: [PATCH 60/68] test(qwen3.5): align async launch topology Signed-off-by: seonjinn --- .../llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh index 98847fc7a39..81d7439e8d5 100755 --- a/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh +++ b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh @@ -5,6 +5,7 @@ source "$SCRIPT_DIR/common.env" # ===== BEGIN CONFIG ===== NUM_NODES=6 GPUS_PER_NODE=4 +SEGMENT_SIZE=2 STEPS_PER_RUN=20 MAX_STEPS=20 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) From 90bbc30f45e4236c0ad949f98ae985dbca0a342a Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sun, 6 Sep 2026 17:46:21 -0700 Subject: [PATCH 61/68] fix(vllm): satisfy refit lint checks Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/vllm_backend.py | 11 +++-------- tests/unit/models/generation/test_vllm_backend.py | 4 +--- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 5c9f3a1003c..d9c97655fdd 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -318,6 +318,7 @@ class VllmInternalWorkerExtension: _mtp_drafter_weights_from_refit: bool = True _sparse_delta_applier: Any = None _nrl_named_parameters: dict[str, torch.nn.Parameter] + hf_to_local_param_map: HFToLocalParamMap _nrl_layerwise_reload_active: bool = False # Initialization detaches parameters, so any later failure leaves this # worker unsafe to reuse. Keep the original failure for the worker lifetime. @@ -546,10 +547,7 @@ def init_nccl_reshard_comm_group( ) refit_info = getattr(self, "nccl_reshard_refit_info", None) - if ( - refit_info is not None - and self._uses_unquantized_flashinfer_trtllm() - ): + if refit_info is not None and self._uses_unquantized_flashinfer_trtllm(): # TRTLLM expert destinations depend on this worker's rank in each # per-PP-stage group, so they cannot be mapped during prepare. self.hf_to_local_param_map = self.build_hf_to_local_param_map(refit_info) @@ -1317,10 +1315,7 @@ def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: self.nccl_reshard_refit_info = ( # pyrefly: ignore[implicitly-defined-attribute] restore_refit_info_placements(refit_info) ) - if ( - self._uses_unquantized_flashinfer_trtllm() - and not self.pp_comm_groups - ): + if self._uses_unquantized_flashinfer_trtllm() and not self.pp_comm_groups: # The TRTLLM expert map needs the per-PP-stage communicator ranks, # which init_nccl_reshard_comm_group establishes after prepare. self.hf_to_local_param_map = HFToLocalParamMap() diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 879063b6aa7..a4a74d911d2 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -555,9 +555,7 @@ def model_load_weights(*, weights): ("model.weight_scale", torch.ones(1)), ] - def get_quantized_weight_iterator( - weights, model_runner, *, refit_with_reload_api - ): + def get_quantized_weight_iterator(weights, model_runner, *, refit_with_reload_api): assert weights is source_weights assert model_runner is ext.model_runner assert refit_with_reload_api is False From 7dfc93af31f1a8416bcf15bfd56c52fa2d760749 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 21:24:09 -0700 Subject: [PATCH 62/68] test(vllm): cover batched BF16 TRTLLM layout Signed-off-by: seonjinn --- .../models/generation/test_vllm_backend.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index a4a74d911d2..1d5ff9f5ede 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -249,6 +249,73 @@ def __init__(self): hpc_module.process_weights_after_loading.assert_called_once_with(model) +@pytest.mark.vllm +def test_batched_bf16_trtllm_layout_matches_expertwise_permutation(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + num_experts = 3 + w13_rows = 4 + w2_rows = 3 + cols = 64 + 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]) + calls = [] + + def get_w13_perm(cache, weight, tile_m, *, is_gated_act_gemm): + calls.append(("w13", tuple(weight.shape), tile_m, is_gated_act_gemm)) + return w13_perm + + def get_w2_perm(cache, weight, tile_m): + calls.append(("w2", tuple(weight.shape), 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, + ) + + actual_w13, actual_w2 = ( + vllm_backend._convert_bf16_moe_weights_to_trtllm_block_layout_batched( + {}, w13, w2, is_gated_act_gemm=True + ) + ) + + adjusted_w13_perm = (w13_perm + w13_rows // 2) % w13_rows + expected_w13 = ( + w13.view(torch.uint8) + .view(num_experts, w13_rows, 1, 128) + .permute(0, 2, 1, 3) + .index_select(2, adjusted_w13_perm) + .contiguous() + .view(torch.bfloat16) + ) + expected_w2 = ( + w2.view(torch.uint8) + .view(num_experts, w2_rows, 1, 128) + .permute(0, 2, 1, 3) + .index_select(2, w2_perm) + .contiguous() + .view(torch.bfloat16) + ) + + torch.testing.assert_close(actual_w13, expected_w13) + torch.testing.assert_close(actual_w2, expected_w2) + assert calls == [ + ("w13", (w13_rows, 128), 128, True), + ("w2", (w2_rows, 128), 128), + ] + + class _DeferredReloadLayer(torch.nn.Module): def __init__(self) -> None: super().__init__() From 7c1856a1050617b6a2bb5a4d5e473ed89bb8c71a Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 21:35:18 -0700 Subject: [PATCH 63/68] perf(vllm): batch BF16 TRTLLM expert layout conversion Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 103 ++++++++++++++++-- .../models/generation/test_vllm_backend.py | 63 ++++++++++- 2 files changed, 153 insertions(+), 13 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index d9c97655fdd..d5b936a9c6d 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -173,6 +173,94 @@ 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") + + from flashinfer.fused_moe.core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, + ) + + 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 = _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 = 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 + + original_converter = ( + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + ) + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout = ( + _convert_bf16_moe_weights_to_trtllm_block_layout_batched + ) + try: + yield + finally: + 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 @@ -1033,13 +1121,14 @@ 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: diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 1d5ff9f5ede..897e9a0b9bc 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -257,12 +257,12 @@ def test_batched_bf16_trtllm_layout_matches_expertwise_permutation(monkeypatch): w13_rows = 4 w2_rows = 3 cols = 64 - 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 = 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]) calls = [] @@ -316,6 +316,57 @@ def get_w2_perm(cache, weight, tile_m): ] +@pytest.mark.vllm +def test_batched_bf16_trtllm_layout_is_scoped_to_reload_finalize(monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + from vllm.model_executor.layers.fused_moe.oracle import unquantized + + original_converter = MagicMock() + 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 + is vllm_backend._convert_bf16_moe_weights_to_trtllm_block_layout_batched + ) + + 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: + finalize() + + assert ( + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + is original_converter + ) + + class _DeferredReloadLayer(torch.nn.Module): def __init__(self) -> None: super().__init__() From 1256eaa2a724c30e0dd7da40588482ef4ed31c2b Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 21:46:18 -0700 Subject: [PATCH 64/68] fix(vllm): cover full TRTLLM reload lifecycle Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 34 +++++++------- .../models/generation/test_vllm_backend.py | 44 +++++++------------ 2 files changed, 35 insertions(+), 43 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index d5b936a9c6d..b02c9a65b53 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -15,6 +15,7 @@ import logging import re import socket +import threading from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager from typing import Any, Literal, Optional @@ -43,6 +44,7 @@ ) logger = logging.getLogger(__name__) +_BF16_TRTLLM_LAYOUT_PATCH_LOCK = threading.RLock() try: import vllm # noqa: F401 @@ -247,18 +249,19 @@ 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 - original_converter = ( - unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout - ) - unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout = ( - _convert_bf16_moe_weights_to_trtllm_block_layout_batched - ) - try: - yield - finally: + with _BF16_TRTLLM_LAYOUT_PATCH_LOCK: + original_converter = ( + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + ) unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout = ( - original_converter + _convert_bf16_moe_weights_to_trtllm_block_layout_batched ) + try: + yield + finally: + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout = ( + original_converter + ) def _local_shard_slices(param_info: dict[str, Any], rank: int) -> tuple[slice, ...]: @@ -1133,11 +1136,12 @@ def finalize() -> None: 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 897e9a0b9bc..2b229282531 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -250,13 +250,19 @@ def __init__(self): @pytest.mark.vllm -def test_batched_bf16_trtllm_layout_matches_expertwise_permutation(monkeypatch): +@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 nemo_rl.models.generation.vllm import vllm_backend + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + convert_moe_weights_to_flashinfer_trtllm_block_layout, + ) num_experts = 3 w13_rows = 4 w2_rows = 3 - cols = 64 + cols = 128 w13 = torch.arange(num_experts * w13_rows * cols, dtype=torch.bfloat16).view( num_experts, w13_rows, cols ) @@ -265,14 +271,11 @@ def test_batched_bf16_trtllm_layout_matches_expertwise_permutation(monkeypatch): ) w13_perm = torch.tensor([2, 0, 3, 1]) w2_perm = torch.tensor([1, 2, 0]) - calls = [] def get_w13_perm(cache, weight, tile_m, *, is_gated_act_gemm): - calls.append(("w13", tuple(weight.shape), tile_m, is_gated_act_gemm)) return w13_perm def get_w2_perm(cache, weight, tile_m): - calls.append(("w2", tuple(weight.shape), tile_m)) return w2_perm monkeypatch.setattr( @@ -284,36 +287,17 @@ def get_w2_perm(cache, weight, tile_m): 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=True + {}, w13, w2, is_gated_act_gemm=is_gated_act_gemm ) ) - adjusted_w13_perm = (w13_perm + w13_rows // 2) % w13_rows - expected_w13 = ( - w13.view(torch.uint8) - .view(num_experts, w13_rows, 1, 128) - .permute(0, 2, 1, 3) - .index_select(2, adjusted_w13_perm) - .contiguous() - .view(torch.bfloat16) - ) - expected_w2 = ( - w2.view(torch.uint8) - .view(num_experts, w2_rows, 1, 128) - .permute(0, 2, 1, 3) - .index_select(2, w2_perm) - .contiguous() - .view(torch.bfloat16) - ) - torch.testing.assert_close(actual_w13, expected_w13) torch.testing.assert_close(actual_w2, expected_w2) - assert calls == [ - ("w13", (w13_rows, 128), 128, True), - ("w2", (w2_rows, 128), 128), - ] @pytest.mark.vllm @@ -359,6 +343,10 @@ def finalize_layerwise_reload(_model, _model_config): monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) with ext._weight_update_lifecycle("collective") as finalize: + assert ( + unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout + is vllm_backend._convert_bf16_moe_weights_to_trtllm_block_layout_batched + ) finalize() assert ( From a62e3e068b00048167f8690490f4fd1a34d65bd0 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 21:52:08 -0700 Subject: [PATCH 65/68] fix(vllm): isolate TRTLLM reload converter by context Signed-off-by: seonjinn --- .../models/generation/vllm/vllm_backend.py | 29 ++++++++++++++-- .../models/generation/test_vllm_backend.py | 33 +++++++++++++++---- 2 files changed, 52 insertions(+), 10 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index b02c9a65b53..aa256c19a64 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -18,6 +18,7 @@ import threading from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager +from contextvars import ContextVar from typing import Any, Literal, Optional import torch @@ -45,6 +46,9 @@ 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 @@ -253,12 +257,31 @@ def _use_batched_bf16_trtllm_layout_conversion() -> Iterator[None]: original_converter = ( unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout ) - unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout = ( - _convert_bf16_moe_weights_to_trtllm_block_layout_batched - ) + + 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 ) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 2b229282531..1e4c3faa9ab 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -302,10 +302,20 @@ def get_w2_perm(cache, weight, tile_m): @pytest.mark.vllm def test_batched_bf16_trtllm_layout_is_scoped_to_reload_finalize(monkeypatch): + import threading + from nemo_rl.models.generation.vllm import vllm_backend from vllm.model_executor.layers.fused_moe.oracle import unquantized - original_converter = MagicMock() + 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", @@ -331,10 +341,7 @@ def test_batched_bf16_trtllm_layout_is_scoped_to_reload_finalize(monkeypatch): ) def finalize_layerwise_reload(_model, _model_config): - assert ( - unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout - is vllm_backend._convert_bf16_moe_weights_to_trtllm_block_layout_batched - ) + assert unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout monkeypatch.setattr( "vllm.model_executor.model_loader.reload.finalize_layerwise_reload", @@ -343,16 +350,28 @@ def finalize_layerwise_reload(_model, _model_config): monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) with ext._weight_update_lifecycle("collective") as finalize: - assert ( + active_converter = ( unquantized.convert_moe_weights_to_flashinfer_trtllm_block_layout - is vllm_backend._convert_bf16_moe_weights_to_trtllm_block_layout_batched ) + 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): From 2c2b6cd31f8704dc2cebe678d4ecdd1c0c874d25 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Mon, 17 Aug 2026 21:50:59 -0700 Subject: [PATCH 66/68] style: format TRTLLM refit tests Signed-off-by: seonjinn --- tests/unit/models/generation/test_vllm_backend.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 1e4c3faa9ab..bbe4ec1e0ed 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -254,11 +254,12 @@ def __init__(self): def test_batched_bf16_trtllm_layout_matches_vllm_expertwise_converter( monkeypatch, is_gated_act_gemm ): - from nemo_rl.models.generation.vllm import vllm_backend 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 @@ -304,9 +305,10 @@ def get_w2_perm(cache, weight, tile_m): def test_batched_bf16_trtllm_layout_is_scoped_to_reload_finalize(monkeypatch): import threading - from nemo_rl.models.generation.vllm import vllm_backend 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) From 014a3051eda07c7316e00e2f26a4e56b1d214455 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Mon, 7 Sep 2026 00:49:18 -0700 Subject: [PATCH 67/68] test(qwen3.5): enable BF16 TRTLLM nightly Signed-off-by: seonjinn --- tests/test_suites/disabled.txt | 5 ----- tests/test_suites/nightly_gb200.txt | 3 +++ 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 5a70393152f..fcbfd7e18e1 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -48,11 +48,6 @@ tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh # GPU-hours would exceed it. Move to nightly.txt when the budget has room. tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-ready-first-single-controller.sh -# Qwen3.5 BF16 FlashInfer TRTLLM + NCCL reshard validation. Keep this -# 6x4-GPU, 20-step run manual until an end-to-end run establishes its -# gen_kl_error/reward bounds and the recurring GB200 suite has budget. -tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh - # The Energon SFTv2 smoke test is a 1n2g recipe, but the nightly tiers pin GPUs # per node: nightly.txt requires 8 and nightly_gb200.txt requires 4, so there is # no tier it can join (test_nightly_suites_match_gpus_per_node enforces this). diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index 9ab02ba24ce..419a3b8d309 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -17,6 +17,9 @@ tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.sh tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh +# BF16 FlashInfer TRTLLM + NCCL reshard numerical correctness +tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh + # TRT-LLM generation backend tests/test_suites/llm/grpo-qwen3-1.7b-2n4g-fsdp2-trtllm.sh tests/test_suites/llm/grpo-qwen2.5-0.5b-1n4g-megatron-trtllm-noncolocated-async.sh From f1950fb231b303200d7d03e90a4822057a0de1bb Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 9 Sep 2026 15:27:19 -0700 Subject: [PATCH 68/68] fix(vllm): defer optional FlashInfer import Signed-off-by: seonjinn --- nemo_rl/models/generation/vllm/vllm_backend.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 6dc899acd52..03ecf8d7245 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -19,6 +19,7 @@ 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 @@ -197,16 +198,13 @@ def _convert_bf16_moe_weights_to_trtllm_block_layout_batched( if w13_weight.shape[0] != w2_weight.shape[0]: raise ValueError("W13 and W2 must contain the same number of experts") - from flashinfer.fused_moe.core import ( - _maybe_get_cached_w3_w1_permute_indices, - get_w2_permute_indices_with_cache, - ) + 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 = _maybe_get_cached_w3_w1_permute_indices( + w13_permute_indices = flashinfer_moe_core._maybe_get_cached_w3_w1_permute_indices( cache_permute_indices, w13_expert_uint8, epilogue_tile_m, @@ -215,7 +213,7 @@ def _convert_bf16_moe_weights_to_trtllm_block_layout_batched( if is_gated_act_gemm: rows = w13_expert_uint8.shape[0] w13_permute_indices = (w13_permute_indices + rows // 2) % rows - w2_permute_indices = get_w2_permute_indices_with_cache( + w2_permute_indices = flashinfer_moe_core.get_w2_permute_indices_with_cache( cache_permute_indices, w2_expert_uint8, epilogue_tile_m,