diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index 379dc7a9c59..e003e4ae0ba 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -406,9 +406,20 @@ def collect_named_tensors_for_weight_transfer( yield name, tensor -def begin_weight_update(rollout_engines: Sequence[ActorHandle]): - """Open a weight-update session on all rollout engines (restore packed weights).""" - ray.get([engine.begin_weight_update.remote() for engine in rollout_engines]) +def begin_weight_update(rollout_engines: Sequence[ActorHandle], selector: str = "all"): + """Open a weight-update session on the selected rollout engines (restore packed weights).""" + ray.get([engine.begin_weight_update.remote(selector=selector) for engine in rollout_engines]) + + +def weight_update_selector(args) -> str: + """Exclude the draft only when the trainer provably has no MTP block to send it.""" + if ( + getattr(args, "sglang_speculative_algorithm", None) + and not getattr(args, "mtp_num_layers", None) + and getattr(args, "megatron_to_hf_mode", "raw") != "bridge" + ): + return "target" + return "all" def end_weight_update(rollout_engines: Sequence[ActorHandle]): diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py index 061efdb9df6..774379bbad8 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py @@ -115,6 +115,7 @@ def _update_weight_implementation( self.weight_version, self.rollout_engines, converted_named_tensors, + selector=self._weight_update_selector, ) ray.get(refs) converted_named_tensors.clear() @@ -258,6 +259,7 @@ def update_weights_from_distributed( weight_version: int, rollout_engines: Sequence[ActorHandle], converted_named_tensors: Sequence[tuple[str, torch.Tensor]], + selector: str = "all", ) -> list[ObjectRef]: """ Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). @@ -267,6 +269,7 @@ def update_weights_from_distributed( names=[name for name, _ in converted_named_tensors], dtypes=[param.dtype for _, param in converted_named_tensors], shapes=[param.shape for _, param in converted_named_tensors], + selector=selector, group_name=group_name, weight_version=str(weight_version), ) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py index 1dcf34f4ce2..b722a6164ba 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py @@ -21,6 +21,7 @@ end_weight_update, get_atomic_update_groups, get_named_value_update_units, + weight_update_selector, ) from ..hf_weight_iterator_base import HfWeightIteratorBase @@ -306,13 +307,14 @@ def _send_one_multi_lora_adapter(self, adapter) -> None: def _pause_and_prepare_engines(self) -> None: """Pause rollout engines, flush cache, and open the weight-update session.""" + self._weight_update_selector = weight_update_selector(self.args) if dist.get_rank() == 0: mode = self.args.pause_generation_mode ray.get([engine.pause_generation.remote(mode=mode) for engine in self.rollout_engines]) if mode != "in_place": ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) - begin_weight_update(self.rollout_engines) + begin_weight_update(self.rollout_engines, self._weight_update_selector) def _finalize_and_resume_engines(self) -> None: """Close the weight-update session and resume rollout engines.""" diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 2905c4277fb..755e75c0e95 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -20,7 +20,7 @@ from miles.utils.lora import LORA_ADAPTER_NAME from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer -from .common import _check_weight_sync_results, begin_weight_update, end_weight_update +from .common import _check_weight_sync_results, begin_weight_update, end_weight_update, weight_update_selector from .hf_weight_iterator_base import HfWeightIteratorBase from .update_weight_from_distributed.broadcast import ( connect_rollout_engines_from_distributed, @@ -215,7 +215,7 @@ def update_weights(self) -> None: ray.get([engine.pause_generation.remote(mode=mode) for engine in self.rollout_engines]) ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) if not skip_base_sync: - begin_weight_update(self.rollout_engines) + begin_weight_update(self.rollout_engines, weight_update_selector(self.args)) dist.barrier(group=get_gloo_group()) megatron_local_weights = self.weights_getter() @@ -269,6 +269,7 @@ def _send_base_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: ipc_engine=self._ipc_engine, ipc_gather_src=self._ipc_gather_src, ipc_gather_group=self._ipc_gather_group, + selector=weight_update_selector(self.args), weight_version=self.weight_version, ) if self.use_distribute and self._is_distributed_src_rank: @@ -278,6 +279,7 @@ def _send_base_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: self.weight_version, self.distributed_rollout_engines, hf_named_tensors, + selector=weight_update_selector(self.args), ) if refs_distributed: refs = (refs or []) + refs_distributed @@ -297,6 +299,7 @@ def _send_lora_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: ipc_engine=self._ipc_engine, ipc_gather_src=self._ipc_gather_src, ipc_gather_group=self._ipc_gather_group, + selector=weight_update_selector(self.args), lora_config=self._lora_config, lora_name=LORA_ADAPTER_NAME, lora_loaded=self._lora_loaded, @@ -317,6 +320,7 @@ def _send_to_colocated_engine( lora_name: str | None = None, lora_loaded: bool = False, check_equal: bool = False, + selector: str = "all", ) -> tuple[list[ObjectRef], Any]: # Placeholder ranks (GPU slots reserved but no engine) have no gather group. # gather_object is only collective among group members, so we skip entirely. @@ -394,6 +398,7 @@ def _send_to_colocated_engine( "serialized_named_tensors": [tensors[i] for tensors in serialized_named_tensors], "load_format": "flattened_bucket", "weight_version": str(weight_version), + "selector": selector, } refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs)) diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 095ca2cd677..40ae2866eed 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -319,6 +319,7 @@ def update_weights_from_tensor( load_format: str | None = None, flush_cache: bool = False, weight_version: str | None = None, + selector: str = "all", ): """ Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GPUs. @@ -330,6 +331,7 @@ def update_weights_from_tensor( "serialized_named_tensors": serialized_named_tensors, "load_format": load_format, "flush_cache": flush_cache, + "selector": selector, } if weight_version is not None: payload["weight_version"] = weight_version @@ -584,7 +586,14 @@ def destroy_weights_update_group(self, group_name): pass def update_weights_from_distributed( - self, names, dtypes, shapes, group_name, flush_cache=False, weight_version: str | None = None + self, + names, + dtypes, + shapes, + group_name, + flush_cache=False, + weight_version: str | None = None, + selector: str = "all", ): payload = { "names": names, @@ -592,6 +601,7 @@ def update_weights_from_distributed( "shapes": shapes, "group_name": group_name, "flush_cache": flush_cache, + "selector": selector, } if weight_version is not None: payload["weight_version"] = weight_version @@ -613,9 +623,9 @@ def continue_generation(self): response.raise_for_status() return response - def begin_weight_update(self): + def begin_weight_update(self, selector: str = "all"): """Open a weight-update session on the engine (restores packed weights for loading).""" - return self._make_request("begin_weight_update", {}) + return self._make_request("begin_weight_update", {"selector": selector}) def end_weight_update(self): """Close the weight-update session (post-load + quant post-process on the full model)."""