[BUGFIX] Invoke "post_process_quantization" by default after weight updating

miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py

@@ -150,19 +150,18 @@ def _pause_and_prepare_engines(self) -> None:
                     post_process_quantization=False,
                 )
 
-    def _finalize_and_resume_engines(self) -> None:
+    def _finalize_and_resume_engines(self, post_load_weights: bool = False) -> None:
         """Run post-process if needed and resume rollout engines."""
         if dist.get_rank() == 0:
-            # int4/fp4 post_process, mxfp8 post-process (swizzle MoE scales).
-            if self.quantization_config and self.quantization_config["quant_method"] in [
-                "compressed-tensors",
-                "mxfp8",
-            ]:
-                post_process_weights(
-                    rollout_engines=self.rollout_engines,
-                    restore_weights_before_load=False,
-                    post_process_quantization=True,
-                )
+            # post_process_quantization is related to the process_weights_after_loading
+            # in the sglang rollout side, which should always be invoked after weight
+            # updating.
+            post_process_weights(
+                rollout_engines=self.rollout_engines,
+                restore_weights_before_load=False,
+                post_process_quantization=True,
+                post_load_weights=post_load_weights,
+            )
             ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
 
     @torch.no_grad()

miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p.py

@@ -18,7 +18,6 @@
 
 from miles.utils.distributed_utils import get_gloo_group
 
-from ..common import post_process_weights
 from .mixin import DistBucketedWeightUpdateMixin
 from .p2p_transfer_utils import (
     P2PTransferManager,
@@ -125,11 +124,7 @@ def _finalize_and_resume_engines(self):
                     for engine in self.rollout_engines
                 ]
             )
-            post_process_weights(
-                rollout_engines=self.rollout_engines,
-                post_load_weights=True,
-            )
-        super()._finalize_and_resume_engines()
+        super()._finalize_and_resume_engines(post_load_weights=True)
 
     def _update_weight_implementation(
         self, converted_named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None

miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py

@@ -205,17 +205,15 @@ def update_weights(self) -> None:
 
         dist.barrier(group=get_gloo_group())
 
-        # int4/fp4 post_process, mxfp8 post-process (swizzle MoE scales).
         if rank == 0:
-            if self.quantization_config and self.quantization_config["quant_method"] in [
-                "compressed-tensors",
-                "mxfp8",
-            ]:
-                post_process_weights(
-                    rollout_engines=self.rollout_engines,
-                    restore_weights_before_load=False,
-                    post_process_quantization=True,
-                )
+            # `post_process_quantization` is related to the `process_weights_after_loading`
+            # in the sglang rollout side, which should always be invoked after weight
+            # updating.
+            post_process_weights(
+                rollout_engines=self.rollout_engines,
+                restore_weights_before_load=False,
+                post_process_quantization=True,
+            )
             ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
         dist.barrier(group=get_gloo_group())
 

tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py

@@ -210,11 +210,14 @@ def test_raises_on_zero_lora_chunks(self, mock_iter_base, mock_dist, mock_ray, m
         with pytest.raises(RuntimeError, match="zero chunks"):
             updater.update_weights()
 
+    @patch("miles.backends.megatron_utils.update_weight.common.ray")
     @patch(f"{_UW_MODULE}.get_gloo_group", return_value=MagicMock())
     @patch(f"{_UW_MODULE}.ray")
     @patch(f"{_UW_MODULE}.dist")
     @patch(f"{_UW_MODULE}.HfWeightIteratorBase")
-    def test_no_raise_for_base_model_zero_chunks(self, mock_iter_base, mock_dist, mock_ray, mock_gloo):
+    def test_no_raise_for_base_model_zero_chunks(
+        self, mock_iter_base, mock_dist, mock_ray, mock_gloo, mock_common_ray
+    ):
         """Base model weight sync with zero chunks is valid (e.g. empty model state)."""
         from miles.backends.megatron_utils.update_weight.update_weight_from_tensor import UpdateWeightFromTensor