From 0af60fc6dcd57679527bf03e87c48131a6f2937d Mon Sep 17 00:00:00 2001 From: brian Date: Fri, 10 Jul 2026 17:43:48 +0800 Subject: [PATCH 1/3] fix(weight-update): post-process bf16 rollout weights --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + .../update_weight_from_distributed.py | 17 +- .../update_weight_from_tensor.py | 15 +- tests/test_bf16_weight_post_process.py | 264 ++++++++++++++++++ 5 files changed, 286 insertions(+), 15 deletions(-) create mode 100644 tests/test_bf16_weight_post_process.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 761496b949..12e257f7d7 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -641,6 +641,10 @@ jobs: "num_gpus": 0, "test_file": "test_empty_colocated_weight_bucket.py" }, + { + "num_gpus": 0, + "test_file": "test_bf16_weight_post_process.py" + }, { "num_gpus": 0, "test_file": "utils/test_hf_checkpoint_saver.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 415586cb84..0b4f8ca45c 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -85,6 +85,7 @@ {'test_file': 'test_placement_group.py', 'num_gpus': 0}, {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0}, {'test_file': 'test_empty_colocated_weight_bucket.py', 'num_gpus': 0}, + {'test_file': 'test_bf16_weight_post_process.py', 'num_gpus': 0}, {'test_file': 'utils/test_hf_checkpoint_saver.py', 'num_gpus': 0}, {'test_file': 'plugin_contracts/test_plugin_rollout_contracts.py', 'num_gpus': 0}, {'test_file': 'plugin_contracts/test_plugin_runtime_hook_contracts.py', 'num_gpus': 0}, diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 1da32420c6..25f7354dab 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -122,13 +122,14 @@ def update_weights(self) -> None: self._send_weights(pbar) if dist.get_rank() == 0: - # int4/fp4 post_process - if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: - post_process_weights( - restore_weights_before_load=False, - post_process_quantization=True, - rollout_engines=self.rollout_engines, - ) + # Re-apply engine-specific post-load transforms. Compressed-tensors + # repacks quantized weights, while BF16 FlashInfer TRT-LLM restores + # its block layout after the canonical weight copy. + post_process_weights( + restore_weights_before_load=False, + post_process_quantization=True, + rollout_engines=self.rollout_engines, + ) ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) @@ -360,7 +361,7 @@ def post_process_weights( rollout_engines: Sequence[ActorHandle], ): """ - Trigger post-process for int4/fp4 quantization on all rollout engines. + Trigger engine-specific post-load processing on all rollout engines. """ ray.get( [ diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index f8efd69c0a..2b95dd224e 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -178,14 +178,15 @@ def update_weights(self) -> None: # IPC handles are now released by the consumers. Clean them up. torch.cuda.ipc_collect() - # int4/fp4 post_process + # Re-apply engine-specific post-load transforms. Compressed-tensors + # repacks quantized weights, while BF16 FlashInfer TRT-LLM restores its + # block layout after the canonical weight copy. if rank == 0: - if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: - post_process_weights( - restore_weights_before_load=False, - post_process_quantization=True, - rollout_engines=self.rollout_engines, - ) + post_process_weights( + restore_weights_before_load=False, + post_process_quantization=True, + rollout_engines=self.rollout_engines, + ) ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) diff --git a/tests/test_bf16_weight_post_process.py b/tests/test_bf16_weight_post_process.py new file mode 100644 index 0000000000..736d8cbde7 --- /dev/null +++ b/tests/test_bf16_weight_post_process.py @@ -0,0 +1,264 @@ +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _package(name, path): + module = types.ModuleType(name) + module.__path__ = [str(path)] + return module + + +def _load_module(monkeypatch, name, path): + sys.modules.pop(name, None) + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _load_update_modules(monkeypatch): + slime_pkg = _package("slime", REPO_ROOT / "slime") + backends_pkg = _package("slime.backends", REPO_ROOT / "slime" / "backends") + megatron_utils_pkg = _package("slime.backends.megatron_utils", REPO_ROOT / "slime" / "backends" / "megatron_utils") + update_weight_pkg = _package( + "slime.backends.megatron_utils.update_weight", + REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight", + ) + slime_utils_pkg = _package("slime.utils", REPO_ROOT / "slime" / "utils") + + ray_mod = types.ModuleType("ray") + ray_mod.get = lambda refs: refs + ray_mod.ObjectRef = object + ray_actor_mod = types.ModuleType("ray.actor") + ray_actor_mod.ActorHandle = object + + megatron_mod = types.ModuleType("megatron") + megatron_core_mod = types.ModuleType("megatron.core") + mpu_mod = types.ModuleType("megatron.core.mpu") + megatron_core_mod.mpu = mpu_mod + + distributed_utils_mod = types.ModuleType("slime.utils.distributed_utils") + distributed_utils_mod.get_gloo_group = lambda: None + distributed_utils_mod.init_process_group = lambda *_args, **_kwargs: None + http_utils_mod = types.ModuleType("slime.utils.http_utils") + http_utils_mod._wrap_ipv6 = lambda host: host + + megatron_to_hf_mod = types.ModuleType("slime.backends.megatron_utils.megatron_to_hf") + megatron_to_hf_mod.convert_to_hf = lambda *_args, **_kwargs: [] + common_mod = types.ModuleType("slime.backends.megatron_utils.update_weight.common") + common_mod.all_gather_param = lambda _name, param: param + common_mod.named_params_and_buffers = lambda *_args, **_kwargs: [] + + sglang_mod = types.ModuleType("slime.backends.megatron_utils.sglang") + sglang_mod.FlattenedTensorBucket = object + sglang_mod.MultiprocessingSerializer = object + iterator_mod = types.ModuleType("slime.backends.megatron_utils.update_weight.hf_weight_iterator_base") + iterator_mod.HfWeightIteratorBase = object + + tqdm_mod = types.ModuleType("tqdm") + tqdm_mod.tqdm = type("tqdm", (), {}) + + modules = { + "slime": slime_pkg, + "slime.backends": backends_pkg, + "slime.backends.megatron_utils": megatron_utils_pkg, + "slime.backends.megatron_utils.update_weight": update_weight_pkg, + "slime.utils": slime_utils_pkg, + "ray": ray_mod, + "ray.actor": ray_actor_mod, + "megatron": megatron_mod, + "megatron.core": megatron_core_mod, + "megatron.core.mpu": mpu_mod, + "slime.utils.distributed_utils": distributed_utils_mod, + "slime.utils.http_utils": http_utils_mod, + "slime.backends.megatron_utils.megatron_to_hf": megatron_to_hf_mod, + "slime.backends.megatron_utils.update_weight.common": common_mod, + "slime.backends.megatron_utils.sglang": sglang_mod, + "slime.backends.megatron_utils.update_weight.hf_weight_iterator_base": iterator_mod, + "tqdm": tqdm_mod, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + update_weight_dir = REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight" + distributed_name = "slime.backends.megatron_utils.update_weight.update_weight_from_distributed" + distributed_module = _load_module( + monkeypatch, + distributed_name, + update_weight_dir / "update_weight_from_distributed.py", + ) + tensor_module = _load_module( + monkeypatch, + "slime.backends.megatron_utils.update_weight.update_weight_from_tensor", + update_weight_dir / "update_weight_from_tensor.py", + ) + return SimpleNamespace(tensor=tensor_module, distributed=distributed_module) + + +@pytest.fixture +def update_modules(monkeypatch): + return _load_update_modules(monkeypatch) + + +class _RemoteMethod: + def __init__(self, event_log, name): + self._event_log = event_log + self._name = name + + def remote(self): + self._event_log.append(self._name) + return self._name + + +class _Engine: + def __init__(self, event_log): + self.pause_generation = _RemoteMethod(event_log, "pause") + self.flush_cache = _RemoteMethod(event_log, "flush") + self.continue_generation = _RemoteMethod(event_log, "continue") + + +def _patch_collectives(monkeypatch, module): + monkeypatch.setattr(module.dist, "get_rank", lambda: 0) + monkeypatch.setattr(module.dist, "barrier", lambda **_kwargs: None) + monkeypatch.setattr(module, "get_gloo_group", lambda: None) + monkeypatch.setattr(module.ray, "get", lambda refs: refs) + + +def _record_post_process(events, calls): + def post_process(**kwargs): + calls.append(kwargs) + events.append("restore" if kwargs["restore_weights_before_load"] else "post_process") + + return post_process + + +def _make_tensor_updater(tensor_module, engine, quantization_config=None): + updater = object.__new__(tensor_module.UpdateWeightFromTensor) + updater.weight_version = 0 + updater.rollout_engines = [engine] + updater.quantization_config = quantization_config + updater.weights_getter = lambda: {} + updater._hf_weight_iterator = SimpleNamespace(get_hf_weight_chunks=lambda _weights: [[]]) + updater._send_hf_params = lambda _chunk: ([], None) + return updater + + +@pytest.mark.unit +def test_tensor_weight_update_post_processes_bf16_engines(monkeypatch, update_modules): + tensor_module = update_modules.tensor + events = [] + calls = [] + engine = _Engine(events) + _patch_collectives(monkeypatch, tensor_module) + monkeypatch.setattr(tensor_module.torch.cuda, "ipc_collect", lambda: None) + monkeypatch.setattr( + tensor_module, + "post_process_weights", + _record_post_process(events, calls), + ) + + updater = _make_tensor_updater(tensor_module, engine) + + updater.update_weights() + + assert events == ["pause", "flush", "post_process", "continue"] + assert calls == [ + { + "restore_weights_before_load": False, + "post_process_quantization": True, + "rollout_engines": [engine], + } + ] + + +@pytest.mark.unit +def test_tensor_weight_update_preserves_compressed_tensor_processing(monkeypatch, update_modules): + tensor_module = update_modules.tensor + events = [] + calls = [] + engine = _Engine(events) + _patch_collectives(monkeypatch, tensor_module) + monkeypatch.setattr(tensor_module.torch.cuda, "ipc_collect", lambda: None) + monkeypatch.setattr( + tensor_module, + "post_process_weights", + _record_post_process(events, calls), + ) + + updater = _make_tensor_updater(tensor_module, engine, {"quant_method": "compressed-tensors"}) + + updater.update_weights() + + assert events == ["pause", "flush", "restore", "post_process", "continue"] + assert [(call["restore_weights_before_load"], call["post_process_quantization"]) for call in calls] == [ + (True, False), + (False, True), + ] + + +@pytest.mark.unit +def test_tensor_weight_update_does_not_resume_after_post_process_failure(monkeypatch, update_modules): + tensor_module = update_modules.tensor + events = [] + engine = _Engine(events) + _patch_collectives(monkeypatch, tensor_module) + monkeypatch.setattr(tensor_module.torch.cuda, "ipc_collect", lambda: None) + + def fail_post_process(**_kwargs): + raise RuntimeError("post-process failed") + + monkeypatch.setattr(tensor_module, "post_process_weights", fail_post_process) + updater = _make_tensor_updater(tensor_module, engine) + + with pytest.raises(RuntimeError, match="post-process failed"): + updater.update_weights() + + assert events == ["pause", "flush"] + + +@pytest.mark.unit +def test_distributed_weight_update_post_processes_bf16_engines(monkeypatch, update_modules): + distributed_module = update_modules.distributed + events = [] + calls = [] + engine = _Engine(events) + _patch_collectives(monkeypatch, distributed_module) + monkeypatch.setattr( + distributed_module, + "post_process_weights", + _record_post_process(events, calls), + ) + + updater = object.__new__(distributed_module.UpdateWeightFromDistributed) + updater.weight_version = 0 + updater.rollout_engines = [engine] + updater.quantization_config = None + updater._group_name = "test" + updater._is_pp_src_rank = False + updater._send_weights = lambda _pbar: events.append("send") + + updater.update_weights() + + assert events == ["pause", "flush", "send", "post_process", "continue"] + assert calls == [ + { + "restore_weights_before_load": False, + "post_process_quantization": True, + "rollout_engines": [engine], + } + ] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) From 8df22e8afa0a015890e3cd85ee0f5848df5c1802 Mon Sep 17 00:00:00 2001 From: brian Date: Thu, 23 Jul 2026 11:31:27 +0800 Subject: [PATCH 2/3] fix(weight-update): scope post-processing by weight format --- .../update_weight_from_distributed.py | 23 ++++++--- .../update_weight_from_tensor.py | 17 ++++--- tests/test_bf16_weight_post_process.py | 49 +++++++++++++++++++ 3 files changed, 73 insertions(+), 16 deletions(-) diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index a673ca0803..5d277cd410 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -123,14 +123,14 @@ def update_weights(self) -> None: self._send_weights(pbar) if dist.get_rank() == 0: - # Re-apply engine-specific post-load transforms. Compressed-tensors - # repacks quantized weights, while BF16 FlashInfer TRT-LLM restores - # its block layout after the canonical weight copy. - post_process_weights( - restore_weights_before_load=False, - post_process_quantization=True, - rollout_engines=self.rollout_engines, - ) + if requires_post_process_after_update(self.quantization_config): + # Compressed-tensors repacks quantized weights, while + # unquantized backends may restore runtime-specific layouts. + post_process_weights( + restore_weights_before_load=False, + post_process_quantization=True, + rollout_engines=self.rollout_engines, + ) ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) @@ -356,6 +356,13 @@ def update_weights_from_distributed( return refs +def requires_post_process_after_update( + quantization_config: Mapping[str, object] | None, +) -> bool: + """Return whether SGLang must finalize weights after a hot update.""" + return not quantization_config or quantization_config.get("quant_method") == "compressed-tensors" + + def post_process_weights( restore_weights_before_load: bool, post_process_quantization: bool, diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index e80835bbde..81be90d1cf 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -22,6 +22,7 @@ connect_rollout_engines_from_distributed, disconnect_rollout_engines_from_distributed, post_process_weights, + requires_post_process_after_update, update_weights_from_distributed, ) @@ -319,15 +320,15 @@ def update_weights(self) -> None: torch.cuda.ipc_collect() torch.cuda.empty_cache() - # Re-apply engine-specific post-load transforms. Compressed-tensors - # repacks quantized weights, while BF16 FlashInfer TRT-LLM restores its - # block layout after the canonical weight copy. if self.rank == 0: - post_process_weights( - restore_weights_before_load=False, - post_process_quantization=True, - rollout_engines=self.rollout_engines, - ) + if requires_post_process_after_update(self.quantization_config): + # Compressed-tensors repacks quantized weights, while + # unquantized backends may restore runtime-specific layouts. + post_process_weights( + restore_weights_before_load=False, + post_process_quantization=True, + rollout_engines=self.rollout_engines, + ) ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) diff --git a/tests/test_bf16_weight_post_process.py b/tests/test_bf16_weight_post_process.py index 6669ceacba..d880d07352 100644 --- a/tests/test_bf16_weight_post_process.py +++ b/tests/test_bf16_weight_post_process.py @@ -217,6 +217,28 @@ def test_tensor_weight_update_preserves_compressed_tensor_processing(monkeypatch ] +@pytest.mark.unit +def test_tensor_weight_update_skips_unrelated_quantization_post_process(monkeypatch, update_modules): + tensor_module = update_modules.tensor + events = [] + calls = [] + engine = _Engine(events) + _patch_collectives(monkeypatch, tensor_module) + monkeypatch.setattr(tensor_module.torch.cuda, "ipc_collect", lambda: None) + monkeypatch.setattr( + tensor_module, + "post_process_weights", + _record_post_process(events, calls), + ) + + updater = _make_tensor_updater(tensor_module, engine, {"quant_method": "fp8"}) + + updater.update_weights() + + assert events == ["pause", "flush", "continue"] + assert calls == [] + + @pytest.mark.unit def test_tensor_weight_update_does_not_resume_after_post_process_failure(monkeypatch, update_modules): tensor_module = update_modules.tensor @@ -270,5 +292,32 @@ def test_distributed_weight_update_post_processes_bf16_engines(monkeypatch, upda ] +@pytest.mark.unit +def test_distributed_weight_update_skips_unrelated_quantization_post_process(monkeypatch, update_modules): + distributed_module = update_modules.distributed + events = [] + calls = [] + engine = _Engine(events) + _patch_collectives(monkeypatch, distributed_module) + monkeypatch.setattr( + distributed_module, + "post_process_weights", + _record_post_process(events, calls), + ) + + updater = object.__new__(distributed_module.UpdateWeightFromDistributed) + updater.weight_version = 0 + updater.rollout_engines = [engine] + updater.quantization_config = {"quant_method": "fp8"} + updater._group_name = "test" + updater._is_pp_src_rank = False + updater._send_weights = lambda _pbar: events.append("send") + + updater.update_weights() + + assert events == ["pause", "flush", "send", "continue"] + assert calls == [] + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) From fb35d5e16e0862dd7ce2bb4daaa7478a22fbef69 Mon Sep 17 00:00:00 2001 From: brian Date: Thu, 23 Jul 2026 11:40:39 +0800 Subject: [PATCH 3/3] test(weight-update): define post-processing scope --- .../update_weight_from_distributed.py | 2 +- tests/test_bf16_weight_post_process.py | 15 +++++++++++++++ tests/test_empty_colocated_weight_bucket.py | 3 +++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 5d277cd410..ddbf947406 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -360,7 +360,7 @@ def requires_post_process_after_update( quantization_config: Mapping[str, object] | None, ) -> bool: """Return whether SGLang must finalize weights after a hot update.""" - return not quantization_config or quantization_config.get("quant_method") == "compressed-tensors" + return quantization_config is None or quantization_config.get("quant_method") == "compressed-tensors" def post_process_weights( diff --git a/tests/test_bf16_weight_post_process.py b/tests/test_bf16_weight_post_process.py index d880d07352..9b64713034 100644 --- a/tests/test_bf16_weight_post_process.py +++ b/tests/test_bf16_weight_post_process.py @@ -117,6 +117,21 @@ def update_modules(monkeypatch): return _load_update_modules(monkeypatch) +@pytest.mark.unit +@pytest.mark.parametrize( + ("quantization_config", "expected"), + [ + (None, True), + ({"quant_method": "compressed-tensors"}, True), + ({"quant_method": "fp8"}, False), + ({}, False), + ({"quant_method": "unknown"}, False), + ], +) +def test_post_process_scope(update_modules, quantization_config, expected): + assert update_modules.distributed.requires_post_process_after_update(quantization_config) is expected + + class _RemoteMethod: def __init__(self, event_log, name): self._event_log = event_log diff --git a/tests/test_empty_colocated_weight_bucket.py b/tests/test_empty_colocated_weight_bucket.py index 7e5db7e0db..7991649bf9 100644 --- a/tests/test_empty_colocated_weight_bucket.py +++ b/tests/test_empty_colocated_weight_bucket.py @@ -126,6 +126,9 @@ def gather_object(obj, object_gather_list, dst, group): update_from_distributed_mod.connect_rollout_engines_from_distributed = lambda *args, **kwargs: None update_from_distributed_mod.disconnect_rollout_engines_from_distributed = lambda *args, **kwargs: None update_from_distributed_mod.post_process_weights = lambda *args, **kwargs: None + update_from_distributed_mod.requires_post_process_after_update = ( + lambda config: config is None or config.get("quant_method") == "compressed-tensors" + ) update_from_distributed_mod.update_weights_from_distributed = lambda *args, **kwargs: [] monkeypatch.setitem(sys.modules, "slime", slime_pkg)