slime：update_weights 完整设计与实现拆解

# slime/backends/megatron_utils/actor.py
update_weight_cls = UpdateWeightFromTensor if self.args.colocate else UpdateWeightFromDistributed
self.weight_updater = update_weight_cls(
    self.args,
    self.model,
    weights_getter=lambda: self.weights_backuper.get("actor"),
    model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name,
    quantization_config=getattr(self.hf_config, "quantization_config", None),
)

# slime/backends/megatron_utils/actor.py
update_weight_cls = UpdateWeightFromTensor if self.args.colocate else UpdateWeightFromDistributed
self.weight_updater = update_weight_cls(
    self.args,
    self.model,
    weights_getter=lambda: self.weights_backuper.get("actor"),
    model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name,
    quantization_config=getattr(self.hf_config, "quantization_config", None),
)

# slime/backends/megatron_utils/actor.py
def update_weights(self) -> None:
    if self.args.debug_train_only or self.args.debug_rollout_only:
        return
 
    if self.args.use_fault_tolerance:
        if dist.get_rank() == 0:
            ray.get(self.rollout_manager.recover_rollout_engines.remote())
        dist.barrier(group=get_gloo_group())
 
    rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
        self.rollout_manager.get_rollout_engines_and_lock.remote()
    )
 
    if self.args.offload_train:
        reload_process_groups()
 
    if num_new_engines > 0:
        self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
        dist.barrier(group=get_gloo_group())
        if dist.get_rank() == 0:
            ray.get(self.rollout_manager.clear_num_new_engines.remote())
 
    with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
        print_memory("before update_weights")
        self.weight_updater.update_weights()
        print_memory("after update_weights")

# slime/backends/megatron_utils/actor.py
def update_weights(self) -> None:
    if self.args.debug_train_only or self.args.debug_rollout_only:
        return
 
    if self.args.use_fault_tolerance:
        if dist.get_rank() == 0:
            ray.get(self.rollout_manager.recover_rollout_engines.remote())
        dist.barrier(group=get_gloo_group())
 
    rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
        self.rollout_manager.get_rollout_engines_and_lock.remote()
    )
 
    if self.args.offload_train:
        reload_process_groups()
 
    if num_new_engines > 0:
        self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
        dist.barrier(group=get_gloo_group())
        if dist.get_rank() == 0:
            ray.get(self.rollout_manager.clear_num_new_engines.remote())
 
    with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
        print_memory("before update_weights")
        self.weight_updater.update_weights()
        print_memory("after update_weights")

# slime/backends/megatron_utils/update_weight/common.py
def _named_params_and_buffers_global(
    args: Namespace, model: Sequence[torch.nn.Module]
) -> Iterator[tuple[str, torch.Tensor]]:
    """
    Yield (global_name, param/buffer) with consistent names across PP/EP. Adjusts indices for
    virtual PP + EP offsets. Handles decoder.layers, mtp.layers (Multi-Token Prediction), expert_bias.
    """
    ep_size = mpu.get_expert_model_parallel_world_size()
    ep_rank = mpu.get_expert_model_parallel_rank()
    if args.num_experts:
        expert_offset = ep_rank * args.num_experts // ep_size
 
    sig = inspect.signature(get_transformer_layer_offset)
    need_vp_stage = "vp_stage" in sig.parameters
 
    for vp_stage, model_module in enumerate(model):
        if need_vp_stage:
            layer_offset = get_transformer_layer_offset(model_module.config, vp_stage)
        else:
            layer_offset = get_transformer_layer_offset(model_module.config)
        for name, param in model_module.named_parameters():
            # for model without ddp wrap
            if not name.startswith("module.module."):
                name = "module." + name
 
            decoder_layers_pattern = r"module\\.module\\.decoder\\.layers\\.(\\d+)\\.(.+)"
            match = re.match(decoder_layers_pattern, name)
            if not match:
                # MTP (Multi-Token Prediction) layers for speculative decoding
                mtp_layers_pattern = r"module\\.module\\.mtp\\.layers\\.(\\d+)\\.(.+)"
                match = re.match(mtp_layers_pattern, name)
                if not match:
                    yield name, param
                    continue
 
                # MTP layer indices start from 0
                layer_idx, rest = match.groups()
                expert_pattern = r"transformer_layer.mlp.experts\\.(.+)\\.weight(\\d+)"
                match = re.match(expert_pattern, rest)
                if not match:
                    yield name, param
                    continue
 
                rest, expert_idx = match.groups()
                expert_idx = int(expert_idx) + expert_offset
                yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.weight{expert_idx}", param
                continue
 
            layer_idx, rest = match.groups()
            layer_idx = int(layer_idx) + layer_offset
 
            # this is hardcoded for te grouped matmul
            expert_pattern = r"mlp.experts\\.(.+)\\.weight(\\d+)"
            match = re.match(expert_pattern, rest)
            if match:
                rest, expert_idx = match.groups()
                expert_idx = int(expert_idx) + expert_offset
                yield f"module.module.decoder.layers.{layer_idx}.mlp.experts.{rest}.weight{expert_idx}", param
            else:
                yield f"module.module.decoder.layers.{layer_idx}.{rest}", param
 
        # treat expert bias as normal parameters
        for name, buffer in model_module.named_buffers():
            # TODO shall we handle (almost) all buffers like Megatron Bridge
            if "expert_bias" not in name:
                continue
            # for model without ddp wrap
            if not name.startswith("module.module."):
                name = "module." + name
 
            decoder_layers_pattern = r"module\\.module\\.decoder\\.layers\\.(\\d+)\\.(.+)"
            match = re.match(decoder_layers_pattern, name)
            if not match:
                yield name, buffer
            else:
                layer_idx, rest = match.groups()
                layer_idx = int(layer_idx) + layer_offset
                yield f"module.module.decoder.layers.{layer_idx}.{rest}", buffer

# slime/backends/megatron_utils/update_weight/common.py
def _named_params_and_buffers_global(
    args: Namespace, model: Sequence[torch.nn.Module]
) -> Iterator[tuple[str, torch.Tensor]]:
    """
    Yield (global_name, param/buffer) with consistent names across PP/EP. Adjusts indices for
    virtual PP + EP offsets. Handles decoder.layers, mtp.layers (Multi-Token Prediction), expert_bias.
    """
    ep_size = mpu.get_expert_model_parallel_world_size()
    ep_rank = mpu.get_expert_model_parallel_rank()
    if args.num_experts:
        expert_offset = ep_rank * args.num_experts // ep_size
 
    sig = inspect.signature(get_transformer_layer_offset)
    need_vp_stage = "vp_stage" in sig.parameters
 
    for vp_stage, model_module in enumerate(model):
        if need_vp_stage:
            layer_offset = get_transformer_layer_offset(model_module.config, vp_stage)
        else:
            layer_offset = get_transformer_layer_offset(model_module.config)
        for name, param in model_module.named_parameters():
            # for model without ddp wrap
            if not name.startswith("module.module."):
                name = "module." + name
 
            decoder_layers_pattern = r"module\\.module\\.decoder\\.layers\\.(\\d+)\\.(.+)"
            match = re.match(decoder_layers_pattern, name)
            if not match:
                # MTP (Multi-Token Prediction) layers for speculative decoding
                mtp_layers_pattern = r"module\\.module\\.mtp\\.layers\\.(\\d+)\\.(.+)"
                match = re.match(mtp_layers_pattern, name)
                if not match:
                    yield name, param
                    continue
 
                # MTP layer indices start from 0
                layer_idx, rest = match.groups()
                expert_pattern = r"transformer_layer.mlp.experts\\.(.+)\\.weight(\\d+)"
                match = re.match(expert_pattern, rest)
                if not match:
                    yield name, param
                    continue
 
                rest, expert_idx = match.groups()
                expert_idx = int(expert_idx) + expert_offset
                yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.weight{expert_idx}", param
                continue
 
            layer_idx, rest = match.groups()
            layer_idx = int(layer_idx) + layer_offset
 
            # this is hardcoded for te grouped matmul
            expert_pattern = r"mlp.experts\\.(.+)\\.weight(\\d+)"
            match = re.match(expert_pattern, rest)
            if match:
                rest, expert_idx = match.groups()
                expert_idx = int(expert_idx) + expert_offset
                yield f"module.module.decoder.layers.{layer_idx}.mlp.experts.{rest}.weight{expert_idx}", param
            else:
                yield f"module.module.decoder.layers.{layer_idx}.{rest}", param
 
        # treat expert bias as normal parameters
        for name, buffer in model_module.named_buffers():
            # TODO shall we handle (almost) all buffers like Megatron Bridge
            if "expert_bias" not in name:
                continue
            # for model without ddp wrap
            if not name.startswith("module.module."):
                name = "module." + name
 
            decoder_layers_pattern = r"module\\.module\\.decoder\\.layers\\.(\\d+)\\.(.+)"
            match = re.match(decoder_layers_pattern, name)
            if not match:
                yield name, buffer
            else:
                layer_idx, rest = match.groups()
                layer_idx = int(layer_idx) + layer_offset
                yield f"module.module.decoder.layers.{layer_idx}.{rest}", buffer

# slime/backends/megatron_utils/update_weight/common.py
def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
    if "expert_bias" in name:
        return param
 
    if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
        return param.data
 
    if ".experts." in name:
        tp_size = mpu.get_expert_tensor_parallel_world_size()
        tp_group = mpu.get_expert_tensor_parallel_group()
    else:
        tp_size = mpu.get_tensor_model_parallel_world_size()
        tp_group = mpu.get_tensor_model_parallel_group()
 
    param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
    dist.all_gather(param_partitions, param.data, group=tp_group)
 
    partition_dim = param.partition_dim
    assert param.partition_stride == 1, "partition_stride != 1 is not supported"
    # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
    # TODO: check only GLU is used.
    if "linear_fc1.weight" in name:
        param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
        param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
    # this is bug in megatron's grouped moe.
    if "linear_fc2.weight" in name:
        if partition_dim == 0:
            partition_dim = 1
    param = torch.cat(param_partitions, dim=partition_dim)
    return param

# slime/backends/megatron_utils/update_weight/common.py
def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
    if "expert_bias" in name:
        return param
 
    if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
        return param.data
 
    if ".experts." in name:
        tp_size = mpu.get_expert_tensor_parallel_world_size()
        tp_group = mpu.get_expert_tensor_parallel_group()
    else:
        tp_size = mpu.get_tensor_model_parallel_world_size()
        tp_group = mpu.get_tensor_model_parallel_group()
 
    param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
    dist.all_gather(param_partitions, param.data, group=tp_group)
 
    partition_dim = param.partition_dim
    assert param.partition_stride == 1, "partition_stride != 1 is not supported"
    # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
    # TODO: check only GLU is used.
    if "linear_fc1.weight" in name:
        param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
        param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
    # this is bug in megatron's grouped moe.
    if "linear_fc2.weight" in name:
        if partition_dim == 0:
            partition_dim = 1
    param = torch.cat(param_partitions, dim=partition_dim)
    return param

# slime/backends/megatron_utils/update_weight/common.py
def all_gather_params_async(
    param_infos_and_params: list[tuple[ParamInfo, torch.Tensor]],
) -> list[torch.Tensor]:
    """
    Parallel TP all-gather for multiple params. Loop 1: for each TP param, allocate buffers +
    dist.all_gather(async_op=True) on expert-TP/regular-TP group (skip expert_bias/non-TP/duplicated).
    Loop 2: wait all NCCL handles (enables overlap). Loop 3: concat partitions + apply GLU rechunk/MoE dim fix.
    """
    # Phase 1: Start all async all_gather operations
    gather_tasks = []
    handles = []
 
    for info, param in param_infos_and_params:
        # Prepare async all_gather
        if "expert_bias" in info.name:
            gather_tasks.append((info, param, None, None, None))
            handles.append(None)
        elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
            gather_tasks.append((info, param.data, None, None, None))
            handles.append(None)
        else:
            # Start async all_gather
            if ".experts." in info.name:
                tp_size = mpu.get_expert_tensor_parallel_world_size()
                tp_group = mpu.get_expert_tensor_parallel_group()
            else:
                tp_size = mpu.get_tensor_model_parallel_world_size()
                tp_group = mpu.get_tensor_model_parallel_group()
 
            param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
            handle = dist.all_gather(param_partitions, param.data, group=tp_group, async_op=True)
            gather_tasks.append((info, None, handle, param_partitions, param.partition_dim))
            handles.append(handle)
 
    # Phase 2: Wait for ALL async operations to complete at once
    # This ensures maximum parallelism by not blocking on individual operations
    for handle in handles:
        if handle is not None:
            handle.wait()
 
    # Phase 3: Process all results after all communications are done
    gathered_params = []
    for info, direct_param, handle, param_partitions, partition_dim in gather_tasks:
        if handle is None:
            # No all_gather needed
            param = direct_param
        else:
            # Process the gathered partitions (same logic as original all_gather_param)
            assert partition_dim is not None, "partition_stride != 1 is not supported"
            if "linear_fc1.weight" in info.name:
                param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
                param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
            if "linear_fc2.weight" in info.name:
                if partition_dim == 0:
                    partition_dim = 1
            param = torch.cat(param_partitions, dim=partition_dim)
 
        gathered_params.append(param)
 
    return gathered_params

# slime/backends/megatron_utils/update_weight/common.py
def all_gather_params_async(
    param_infos_and_params: list[tuple[ParamInfo, torch.Tensor]],
) -> list[torch.Tensor]:
    """
    Parallel TP all-gather for multiple params. Loop 1: for each TP param, allocate buffers +
    dist.all_gather(async_op=True) on expert-TP/regular-TP group (skip expert_bias/non-TP/duplicated).
    Loop 2: wait all NCCL handles (enables overlap). Loop 3: concat partitions + apply GLU rechunk/MoE dim fix.
    """
    # Phase 1: Start all async all_gather operations
    gather_tasks = []
    handles = []
 
    for info, param in param_infos_and_params:
        # Prepare async all_gather
        if "expert_bias" in info.name:
            gather_tasks.append((info, param, None, None, None))
            handles.append(None)
        elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
            gather_tasks.append((info, param.data, None, None, None))
            handles.append(None)
        else:
            # Start async all_gather
            if ".experts." in info.name:
                tp_size = mpu.get_expert_tensor_parallel_world_size()
                tp_group = mpu.get_expert_tensor_parallel_group()
            else:
                tp_size = mpu.get_tensor_model_parallel_world_size()
                tp_group = mpu.get_tensor_model_parallel_group()
 
            param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
            handle = dist.all_gather(param_partitions, param.data, group=tp_group, async_op=True)
            gather_tasks.append((info, None, handle, param_partitions, param.partition_dim))
            handles.append(handle)
 
    # Phase 2: Wait for ALL async operations to complete at once
    # This ensures maximum parallelism by not blocking on individual operations
    for handle in handles:
        if handle is not None:
            handle.wait()
 
    # Phase 3: Process all results after all communications are done
    gathered_params = []
    for info, direct_param, handle, param_partitions, partition_dim in gather_tasks:
        if handle is None:
            # No all_gather needed
            param = direct_param
        else:
            # Process the gathered partitions (same logic as original all_gather_param)
            assert partition_dim is not None, "partition_stride != 1 is not supported"
            if "linear_fc1.weight" in info.name:
                param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
                param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
            if "linear_fc2.weight" in info.name:
                if partition_dim == 0:
                    partition_dim = 1
            param = torch.cat(param_partitions, dim=partition_dim)
 
        gathered_params.append(param)
 
    return gathered_params

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py
class HfWeightIteratorBase(ABC):
    @staticmethod
    def create(args, model, **kwargs):
        from .hf_weight_iterator_bridge import HfWeightIteratorBridge
        from .hf_weight_iterator_direct import HfWeightIteratorDirect
 
        c = {
            "raw": HfWeightIteratorDirect,
            "bridge": HfWeightIteratorBridge,
        }[args.megatron_to_hf_mode]
 
        return c(args, model, **kwargs)
 
    def __init__(self, args, model, model_name, quantization_config):
        self.args = args
        self.model = model
        self.model_name = model_name
        self.quantization_config = quantization_config
 
    @abstractmethod
    def get_hf_weight_chunks(self, megatron_local_weights):

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py
class HfWeightIteratorBase(ABC):
    @staticmethod
    def create(args, model, **kwargs):
        from .hf_weight_iterator_bridge import HfWeightIteratorBridge
        from .hf_weight_iterator_direct import HfWeightIteratorDirect
 
        c = {
            "raw": HfWeightIteratorDirect,
            "bridge": HfWeightIteratorBridge,
        }[args.megatron_to_hf_mode]
 
        return c(args, model, **kwargs)
 
    def __init__(self, args, model, model_name, quantization_config):
        self.args = args
        self.model = model
        self.model_name = model_name
        self.quantization_config = quantization_config
 
    @abstractmethod
    def get_hf_weight_chunks(self, megatron_local_weights):

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
class HfWeightIteratorDirect(HfWeightIteratorBase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.megatron_local_param_info_buckets = _get_megatron_local_param_info_buckets(self.args, self.model)
 
    def get_hf_weight_chunks(self, megatron_local_weights):
        rank = dist.get_rank()
 
        for megatron_local_param_infos in tqdm(
            self.megatron_local_param_info_buckets, disable=rank != 0, desc="Update weights"
        ):
            megatron_full_params = _get_megatron_full_params(megatron_local_param_infos, megatron_local_weights)
            hf_named_tensors = self._convert_to_hf_named_tensors(megatron_full_params, megatron_local_param_infos)
            yield hf_named_tensors
            del megatron_full_params
 
    def _convert_to_hf_named_tensors(self, megatron_full_params: Sequence[torch.Tensor], param_infos: list[ParamInfo]):
        hf_named_tensors = []
        for info, param in zip(param_infos, megatron_full_params, strict=False):
            hf_named_tensors.extend(
                convert_to_hf(self.args, self.model_name, info.name, param, self.quantization_config)
            )
        return hf_named_tensors

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
class HfWeightIteratorDirect(HfWeightIteratorBase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.megatron_local_param_info_buckets = _get_megatron_local_param_info_buckets(self.args, self.model)
 
    def get_hf_weight_chunks(self, megatron_local_weights):
        rank = dist.get_rank()
 
        for megatron_local_param_infos in tqdm(
            self.megatron_local_param_info_buckets, disable=rank != 0, desc="Update weights"
        ):
            megatron_full_params = _get_megatron_full_params(megatron_local_param_infos, megatron_local_weights)
            hf_named_tensors = self._convert_to_hf_named_tensors(megatron_full_params, megatron_local_param_infos)
            yield hf_named_tensors
            del megatron_full_params
 
    def _convert_to_hf_named_tensors(self, megatron_full_params: Sequence[torch.Tensor], param_infos: list[ParamInfo]):
        hf_named_tensors = []
        for info, param in zip(param_infos, megatron_full_params, strict=False):
            hf_named_tensors.extend(
                convert_to_hf(self.args, self.model_name, info.name, param, self.quantization_config)
            )
        return hf_named_tensors

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
def _get_megatron_full_params(
    megatron_local_param_infos: Sequence[ParamInfo],
    megatron_local_weights,
) -> Sequence[torch.Tensor]:
    monkey_patch_torch_reductions()
    pp_size = mpu.get_pipeline_model_parallel_world_size()
    ep_size = mpu.get_expert_model_parallel_world_size()
    rank = dist.get_rank()
    # init params:
    params = []
    for info in megatron_local_param_infos:
        if dist.get_rank() == info.src_rank:
            params.append(
                torch.nn.Parameter(
                    megatron_local_weights[info.name].to(device=torch.cuda.current_device(), non_blocking=True),
                    requires_grad=False,
                )
            )
        else:
            params.append(torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device()))
    torch.cuda.synchronize()
 
    # broadcast params across pp ranks
    if pp_size > 1:
        handles = []
        for info, param in zip(megatron_local_param_infos, params, strict=False):
            if info.src_rank in dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()):
                handles.append(
                    torch.distributed.broadcast(
                        param, src=info.src_rank, group=mpu.get_pipeline_model_parallel_group(), async_op=True
                    )
                )
        for handle in handles:
            handle.wait()
 
    # broadcast params across ep ranks
    if ep_size > 1:
        handles = []
        for info, param in zip(megatron_local_param_infos, params, strict=False):
            if ".experts." in info.name:
                src_rank = (
                    info.src_rank
                    if info.src_rank in dist.get_process_group_ranks(mpu.get_expert_model_parallel_group())
                    else rank
                )
                handles.append(
                    torch.distributed.broadcast(
                        param, src=src_rank, group=mpu.get_expert_model_parallel_group(), async_op=True
                    )
                )
        for handle in handles:
            handle.wait()
 
    # Set tp attrs for all params
    for info, param in zip(megatron_local_param_infos, params, strict=False):
        for key, value in info.attrs.items():
            setattr(param, key, value)
 
    # Batch async all_gather for all parameters
    gathered_params = all_gather_params_async(list(zip(megatron_local_param_infos, params, strict=False)))
 
    return gathered_params

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
def _get_megatron_full_params(
    megatron_local_param_infos: Sequence[ParamInfo],
    megatron_local_weights,
) -> Sequence[torch.Tensor]:
    monkey_patch_torch_reductions()
    pp_size = mpu.get_pipeline_model_parallel_world_size()
    ep_size = mpu.get_expert_model_parallel_world_size()
    rank = dist.get_rank()
    # init params:
    params = []
    for info in megatron_local_param_infos:
        if dist.get_rank() == info.src_rank:
            params.append(
                torch.nn.Parameter(
                    megatron_local_weights[info.name].to(device=torch.cuda.current_device(), non_blocking=True),
                    requires_grad=False,
                )
            )
        else:
            params.append(torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device()))
    torch.cuda.synchronize()
 
    # broadcast params across pp ranks
    if pp_size > 1:
        handles = []
        for info, param in zip(megatron_local_param_infos, params, strict=False):
            if info.src_rank in dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()):
                handles.append(
                    torch.distributed.broadcast(
                        param, src=info.src_rank, group=mpu.get_pipeline_model_parallel_group(), async_op=True
                    )
                )
        for handle in handles:
            handle.wait()
 
    # broadcast params across ep ranks
    if ep_size > 1:
        handles = []
        for info, param in zip(megatron_local_param_infos, params, strict=False):
            if ".experts." in info.name:
                src_rank = (
                    info.src_rank
                    if info.src_rank in dist.get_process_group_ranks(mpu.get_expert_model_parallel_group())
                    else rank
                )
                handles.append(
                    torch.distributed.broadcast(
                        param, src=src_rank, group=mpu.get_expert_model_parallel_group(), async_op=True
                    )
                )
        for handle in handles:
            handle.wait()
 
    # Set tp attrs for all params
    for info, param in zip(megatron_local_param_infos, params, strict=False):
        for key, value in info.attrs.items():
            setattr(param, key, value)
 
    # Batch async all_gather for all parameters
    gathered_params = all_gather_params_async(list(zip(megatron_local_param_infos, params, strict=False)))
 
    return gathered_params

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torch.nn.Module]) -> list[list[ParamInfo]]:
    """
    Partition params into buckets ≤ update_weight_buffer_size (with TP replication).
    """
    param_infos = _get_megatron_local_param_infos(args, model)
    param_info_buckets = [[]]  # Start with one empty bucket
    buffer_size = 0  # Track current bucket size in bytes
 
    for info in param_infos:
        # Expert params use expert-TP size, others use regular TP
        if ".experts." in info.name:
            tp_size = mpu.get_expert_tensor_parallel_world_size()
        else:
            tp_size = mpu.get_tensor_model_parallel_world_size()
 
        param_size = info.numel * info.dtype.itemsize * tp_size
 
        # If adding this param would exceed buffer size, start a new bucket
        if buffer_size + param_size > args.update_weight_buffer_size and len(param_info_buckets[-1]) > 0:
            param_info_buckets.append([])
            buffer_size = 0
 
        # Add param to current bucket and update size
        param_info_buckets[-1].append(info)
        buffer_size += param_size
 
    return param_info_buckets

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torch.nn.Module]) -> list[list[ParamInfo]]:
    """
    Partition params into buckets ≤ update_weight_buffer_size (with TP replication).
    """
    param_infos = _get_megatron_local_param_infos(args, model)
    param_info_buckets = [[]]  # Start with one empty bucket
    buffer_size = 0  # Track current bucket size in bytes
 
    for info in param_infos:
        # Expert params use expert-TP size, others use regular TP
        if ".experts." in info.name:
            tp_size = mpu.get_expert_tensor_parallel_world_size()
        else:
            tp_size = mpu.get_tensor_model_parallel_world_size()
 
        param_size = info.numel * info.dtype.itemsize * tp_size
 
        # If adding this param would exceed buffer size, start a new bucket
        if buffer_size + param_size > args.update_weight_buffer_size and len(param_info_buckets[-1]) > 0:
            param_info_buckets.append([])
            buffer_size = 0
 
        # Add param to current bucket and update size
        param_info_buckets[-1].append(info)
        buffer_size += param_size
 
    return param_info_buckets

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
class HfWeightIteratorBridge(HfWeightIteratorBase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
 
        from megatron.bridge import AutoBridge
 
        import slime_plugins.megatron_bridge  # noqa: F401
 
        self._bridge = AutoBridge.from_hf_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
 
    def get_hf_weight_chunks(self, megatron_local_weights):
        # TODO support quantization (e.g. modify megatron-bridge to provide megatron param name)
        renamed_megatron_local_weights = {strip_param_name_prefix(k): v for k, v in megatron_local_weights.items()}
        with megatron_bridge_utils.patch_megatron_model(self.model):
            conversion_tasks = self._bridge.get_conversion_tasks(self.model)
            conversion_tasks = _process_conversion_tasks(conversion_tasks, renamed_megatron_local_weights)
 
            named_weights = self._bridge.export_hf_weights(self.model, cpu=False, conversion_tasks=conversion_tasks)
 
            named_weights = (
                (
                    hf_param_name,
                    postprocess_hf_param(
                        args=self.args,
                        megatron_param_name=megatron_param_name,
                        hf_param_name=hf_param_name,
                        param=weight,
                    ),
                )
                for hf_param_name, weight, megatron_param_name in named_weights
            )
 
            yield from chunk_named_params_by_size(named_weights, chunk_size=self.args.update_weight_buffer_size)

# slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
class HfWeightIteratorBridge(HfWeightIteratorBase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
 
        from megatron.bridge import AutoBridge
 
        import slime_plugins.megatron_bridge  # noqa: F401
 
        self._bridge = AutoBridge.from_hf_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
 
    def get_hf_weight_chunks(self, megatron_local_weights):
        # TODO support quantization (e.g. modify megatron-bridge to provide megatron param name)
        renamed_megatron_local_weights = {strip_param_name_prefix(k): v for k, v in megatron_local_weights.items()}
        with megatron_bridge_utils.patch_megatron_model(self.model):
            conversion_tasks = self._bridge.get_conversion_tasks(self.model)
            conversion_tasks = _process_conversion_tasks(conversion_tasks, renamed_megatron_local_weights)
 
            named_weights = self._bridge.export_hf_weights(self.model, cpu=False, conversion_tasks=conversion_tasks)
 
            named_weights = (
                (
                    hf_param_name,
                    postprocess_hf_param(
                        args=self.args,
                        megatron_param_name=megatron_param_name,
                        hf_param_name=hf_param_name,
                        param=weight,
                    ),
                )
                for hf_param_name, weight, megatron_param_name in named_weights
            )
 
            yield from chunk_named_params_by_size(named_weights, chunk_size=self.args.update_weight_buffer_size)

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
class UpdateWeightFromTensor:
    """
    Update rollout engines from tensor dict:
    load(dict→GPU) → broadcast PP/EP(GPU NCCL) → gather TP(GPU NCCL) → convert HF(GPU) → send.
    Colocated: GPU→CPU serialize → gather_object(Gloo CPU, collects from rollout_num_gpus_per_engine ranks) → Ray IPC to engine.
    Distributed: GPU NCCL broadcast to remote engines.
    """

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
class UpdateWeightFromTensor:
    """
    Update rollout engines from tensor dict:
    load(dict→GPU) → broadcast PP/EP(GPU NCCL) → gather TP(GPU NCCL) → convert HF(GPU) → send.
    Colocated: GPU→CPU serialize → gather_object(Gloo CPU, collects from rollout_num_gpus_per_engine ranks) → Ray IPC to engine.
    Distributed: GPU NCCL broadcast to remote engines.
    """

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
def connect_rollout_engines(
    self, rollout_engines: Sequence[ActorHandle], rollout_engine_lock: ActorHandle
) -> None:
    """
    Split colocated/distributed engines. Global source rank (DP=TP=PP=0) creates NCCL
    for distributed. Map ranks to colocated IPC engines.
    """
    self.rollout_engines = rollout_engines
    colocate_engine_nums = (
        self.args.actor_num_nodes * self.args.actor_num_gpus_per_node // self.args.rollout_num_gpus_per_engine
    )
    self.use_distribute = len(rollout_engines) > colocate_engine_nums
 
    if self.use_distribute:
        self.rollout_engines = rollout_engines[:colocate_engine_nums]
        self.distributed_rollout_engines = rollout_engines[colocate_engine_nums:]
        self._is_distributed_src_rank = (
            mpu.get_data_parallel_rank(with_context_parallel=True) == 0
            and mpu.get_tensor_model_parallel_rank() == 0
            and mpu.get_pipeline_model_parallel_rank() == 0
        )
        self._group_name = "slime"
        if self._is_distributed_src_rank:
            if self._model_update_groups is not None:
                disconnect_rollout_engines_from_distributed(
                    self.args, self._group_name, self._model_update_groups, self.distributed_rollout_engines
                )
 
            self._model_update_groups = connect_rollout_engines_from_distributed(
                self.args, self._group_name, self.distributed_rollout_engines
            )
 
    # Here we assume the gpu id of rollout engines and train actors are the same.
    for i, engine in enumerate(self.rollout_engines):
        start_rank = i * self.args.rollout_num_gpus_per_engine
        end_rank = (i + 1) * self.args.rollout_num_gpus_per_engine
        group_ranks = list(range(start_rank, end_rank))
        if dist.get_rank() in group_ranks:
            self._ipc_engine = engine

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
def connect_rollout_engines(
    self, rollout_engines: Sequence[ActorHandle], rollout_engine_lock: ActorHandle
) -> None:
    """
    Split colocated/distributed engines. Global source rank (DP=TP=PP=0) creates NCCL
    for distributed. Map ranks to colocated IPC engines.
    """
    self.rollout_engines = rollout_engines
    colocate_engine_nums = (
        self.args.actor_num_nodes * self.args.actor_num_gpus_per_node // self.args.rollout_num_gpus_per_engine
    )
    self.use_distribute = len(rollout_engines) > colocate_engine_nums
 
    if self.use_distribute:
        self.rollout_engines = rollout_engines[:colocate_engine_nums]
        self.distributed_rollout_engines = rollout_engines[colocate_engine_nums:]
        self._is_distributed_src_rank = (
            mpu.get_data_parallel_rank(with_context_parallel=True) == 0
            and mpu.get_tensor_model_parallel_rank() == 0
            and mpu.get_pipeline_model_parallel_rank() == 0
        )
        self._group_name = "slime"
        if self._is_distributed_src_rank:
            if self._model_update_groups is not None:
                disconnect_rollout_engines_from_distributed(
                    self.args, self._group_name, self._model_update_groups, self.distributed_rollout_engines
                )
 
            self._model_update_groups = connect_rollout_engines_from_distributed(
                self.args, self._group_name, self.distributed_rollout_engines
            )
 
    # Here we assume the gpu id of rollout engines and train actors are the same.
    for i, engine in enumerate(self.rollout_engines):
        start_rank = i * self.args.rollout_num_gpus_per_engine
        end_rank = (i + 1) * self.args.rollout_num_gpus_per_engine
        group_ranks = list(range(start_rank, end_rank))
        if dist.get_rank() in group_ranks:
            self._ipc_engine = engine

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
@torch.no_grad()
def update_weights(self) -> None:
    """
    version++, flush caches, process buckets. Progress on rank 0.
    """
    self.weight_version += 1
 
    rank = dist.get_rank()
    if rank == 0:
        ray.get([engine.pause_generation.remote() for engine in self.rollout_engines])
        ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
        if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]:
            post_process_weights(
                restore_weights_before_load=True,
                post_process_quantization=False,
                rollout_engines=self.rollout_engines,
            )
    dist.barrier(group=get_gloo_group())
 
    megatron_local_weights = self.weights_getter()
 
    for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
        refs, long_lived_tensors = self._send_hf_params(hf_named_tensors)
        ray.get(refs)
        del long_lived_tensors
 
    dist.barrier(group=get_gloo_group())
 
    # int4/fp4 post_process
    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,
            )
        ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
    dist.barrier(group=get_gloo_group())

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
@torch.no_grad()
def update_weights(self) -> None:
    """
    version++, flush caches, process buckets. Progress on rank 0.
    """
    self.weight_version += 1
 
    rank = dist.get_rank()
    if rank == 0:
        ray.get([engine.pause_generation.remote() for engine in self.rollout_engines])
        ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
        if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]:
            post_process_weights(
                restore_weights_before_load=True,
                post_process_quantization=False,
                rollout_engines=self.rollout_engines,
            )
    dist.barrier(group=get_gloo_group())
 
    megatron_local_weights = self.weights_getter()
 
    for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
        refs, long_lived_tensors = self._send_hf_params(hf_named_tensors)
        ray.get(refs)
        del long_lived_tensors
 
    dist.barrier(group=get_gloo_group())
 
    # int4/fp4 post_process
    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,
            )
        ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
    dist.barrier(group=get_gloo_group())

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
def _send_to_colocated_engine(
    hf_named_tensors: list[tuple[str, torch.Tensor]],
    *,
    ipc_engine,
    ipc_gather_src,
    ipc_gather_group,
    weight_version,
) -> tuple[list[ObjectRef], Any]:
    # TODO improve
    long_live_tensors = []
 
    if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
        converted_named_tensors_by_dtypes = {"dtype": hf_named_tensors}
    else:
        converted_named_tensors_by_dtypes = {}
        for name, tensor in hf_named_tensors:
            dtype = tensor.dtype
            if dtype not in converted_named_tensors_by_dtypes:
                converted_named_tensors_by_dtypes[dtype] = []
            converted_named_tensors_by_dtypes[dtype].append((name, tensor))
 
    serialized_tensors = []
    for _dtype, named_tensors in converted_named_tensors_by_dtypes.items():
        flattened_tensor_bucket = FlattenedTensorBucket(named_tensors=named_tensors)
        metadata = flattened_tensor_bucket.get_metadata()
        flattened_tensor_data = {
            "flattened_tensor": flattened_tensor_bucket.get_flattened_tensor(),
            "metadata": metadata,
        }
        long_live_tensors.append(flattened_tensor_data)
        serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True))
 
    serialized_named_tensors = (
        [None] * dist.get_world_size(ipc_gather_group) if ipc_gather_src == dist.get_rank() else None
    )
    dist.gather_object(
        serialized_tensors,
        object_gather_list=serialized_named_tensors,
        dst=ipc_gather_src,
        group=ipc_gather_group,
    )
 
    refs = []
    if dist.get_rank() == ipc_gather_src:
        # TODO: here we assume all ranks have the same number of dtypes, not sure if that is correct.
        num_dtypes = len(serialized_named_tensors[0])
        for i in range(num_dtypes):
            kwargs = {
                "serialized_named_tensors": [tensors[i] for tensors in serialized_named_tensors],
                "load_format": "flattened_bucket",
                "weight_version": str(weight_version),
            }
            refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs))
 
    return refs, long_lived_tensors

# slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
def _send_to_colocated_engine(
    hf_named_tensors: list[tuple[str, torch.Tensor]],
    *,
    ipc_engine,
    ipc_gather_src,
    ipc_gather_group,
    weight_version,
) -> tuple[list[ObjectRef], Any]:
    # TODO improve
    long_live_tensors = []
 
    if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
        converted_named_tensors_by_dtypes = {"dtype": hf_named_tensors}
    else:
        converted_named_tensors_by_dtypes = {}
        for name, tensor in hf_named_tensors:
            dtype = tensor.dtype
            if dtype not in converted_named_tensors_by_dtypes:
                converted_named_tensors_by_dtypes[dtype] = []
            converted_named_tensors_by_dtypes[dtype].append((name, tensor))
 
    serialized_tensors = []
    for _dtype, named_tensors in converted_named_tensors_by_dtypes.items():
        flattened_tensor_bucket = FlattenedTensorBucket(named_tensors=named_tensors)
        metadata = flattened_tensor_bucket.get_metadata()
        flattened_tensor_data = {
            "flattened_tensor": flattened_tensor_bucket.get_flattened_tensor(),
            "metadata": metadata,
        }
        long_live_tensors.append(flattened_tensor_data)
        serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True))
 
    serialized_named_tensors = (
        [None] * dist.get_world_size(ipc_gather_group) if ipc_gather_src == dist.get_rank() else None
    )
    dist.gather_object(
        serialized_tensors,
        object_gather_list=serialized_named_tensors,
        dst=ipc_gather_src,
        group=ipc_gather_group,
    )
 
    refs = []
    if dist.get_rank() == ipc_gather_src:
        # TODO: here we assume all ranks have the same number of dtypes, not sure if that is correct.
        num_dtypes = len(serialized_named_tensors[0])
        for i in range(num_dtypes):
            kwargs = {
                "serialized_named_tensors": [tensors[i] for tensors in serialized_named_tensors],
                "load_format": "flattened_bucket",
                "weight_version": str(weight_version),
            }
            refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs))
 
    return refs, long_lived_tensors

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
class UpdateWeightFromDistributed:
    """
    Training rank 0 broadcasts HF-converted tensors to rollout engines via NCCL custom process group.
    """

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
class UpdateWeightFromDistributed:
    """
    Training rank 0 broadcasts HF-converted tensors to rollout engines via NCCL custom process group.
    """

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def connect_rollout_engines_from_distributed(
    args: Namespace, group_name: str, rollout_engines: Sequence[ActorHandle]
) -> dist.ProcessGroup:
    """
    Create NCCL group: training rank 0 + all engine GPUs. Blocks until joined.
    """
    master_address = ray._private.services.get_node_ip_address()
    with socket.socket() as sock:
        sock.bind(("", 0))
        master_port = sock.getsockname()[1]
    world_size = len(rollout_engines) * args.rollout_num_gpus_per_engine + 1
 
    refs = [
        engine.init_weights_update_group.remote(
            master_address,
            master_port,
            i * args.rollout_num_gpus_per_engine + 1,
            world_size,
            group_name,
            backend="nccl",
        )
        for i, engine in enumerate(rollout_engines)
    ]
    model_update_groups = init_process_group(
        backend="nccl",
        init_method=f"tcp://{master_address}:{master_port}",
        world_size=world_size,
        rank=0,
        group_name=group_name,
    )
    ray.get(refs)
    return model_update_groups

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def connect_rollout_engines_from_distributed(
    args: Namespace, group_name: str, rollout_engines: Sequence[ActorHandle]
) -> dist.ProcessGroup:
    """
    Create NCCL group: training rank 0 + all engine GPUs. Blocks until joined.
    """
    master_address = ray._private.services.get_node_ip_address()
    with socket.socket() as sock:
        sock.bind(("", 0))
        master_port = sock.getsockname()[1]
    world_size = len(rollout_engines) * args.rollout_num_gpus_per_engine + 1
 
    refs = [
        engine.init_weights_update_group.remote(
            master_address,
            master_port,
            i * args.rollout_num_gpus_per_engine + 1,
            world_size,
            group_name,
            backend="nccl",
        )
        for i, engine in enumerate(rollout_engines)
    ]
    model_update_groups = init_process_group(
        backend="nccl",
        init_method=f"tcp://{master_address}:{master_port}",
        world_size=world_size,
        rank=0,
        group_name=group_name,
    )
    ray.get(refs)
    return model_update_groups

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def update_weights_from_distributed(
    group_name: str,
    group: dist.ProcessGroup,
    weight_version: int,
    rollout_engines: Sequence[ActorHandle],
    converted_named_tensors: Sequence[tuple[str, torch.Tensor]],
) -> list[ObjectRef]:
    """
    Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines).
    """
    refs = [
        engine.update_weights_from_distributed.remote(
            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],
            group_name=group_name,
            weight_version=str(weight_version),
        )
        for engine in rollout_engines
    ]
 
    handles = []
    for _, param in converted_named_tensors:
        handles.append(dist.broadcast(param.data, 0, group=group, async_op=True))
    for handle in handles:
        handle.wait()
 
    return refs

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def update_weights_from_distributed(
    group_name: str,
    group: dist.ProcessGroup,
    weight_version: int,
    rollout_engines: Sequence[ActorHandle],
    converted_named_tensors: Sequence[tuple[str, torch.Tensor]],
) -> list[ObjectRef]:
    """
    Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines).
    """
    refs = [
        engine.update_weights_from_distributed.remote(
            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],
            group_name=group_name,
            weight_version=str(weight_version),
        )
        for engine in rollout_engines
    ]
 
    handles = []
    for _, param in converted_named_tensors:
        handles.append(dist.broadcast(param.data, 0, group=group, async_op=True))
    for handle in handles:
        handle.wait()
 
    return refs

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def _update_bucket_weights_from_distributed(
    self, converted_named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None
) -> None:
    """
    Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock.
    """
    # lock the rollout engines to prevent dead lock on broadcast.
    while not ray.get(self.rollout_engine_lock.acquire.remote()):
        time.sleep(0.1)
 
    refs = update_weights_from_distributed(
        self._group_name,
        self._model_update_groups,
        self.weight_version,
        self.rollout_engines,
        converted_named_tensors,
    )
 
    ray.get(refs)
    converted_named_tensors.clear()
    ray.get(self.rollout_engine_lock.release.remote())
    pbar.update(1)

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def _update_bucket_weights_from_distributed(
    self, converted_named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None
) -> None:
    """
    Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock.
    """
    # lock the rollout engines to prevent dead lock on broadcast.
    while not ray.get(self.rollout_engine_lock.acquire.remote()):
        time.sleep(0.1)
 
    refs = update_weights_from_distributed(
        self._group_name,
        self._model_update_groups,
        self.weight_version,
        self.rollout_engines,
        converted_named_tensors,
    )
 
    ray.get(refs)
    converted_named_tensors.clear()
    ray.get(self.rollout_engine_lock.release.remote())
    pbar.update(1)

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def _update_weight_from_distributed(
    self,
    name: str,
    param: torch.nn.Parameter,
    converted_named_tensors: list[tuple[str, torch.Tensor]],
    buffer_size: int,
    pbar: tqdm | None = None,
) -> int | None:
    """
    Non-expert: gather TP → rm pad → HF → buffer (flush if full). All gather, PP source buffers.
    Returns updated bytes on source, None on non-source.
    """
    param = all_gather_param(name, param)
    if not self._is_pp_src_rank:
        return
 
    param_size = param.numel() * param.element_size()
    if buffer_size + param_size > self.args.update_weight_buffer_size:
        self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar)
        buffer_size = 0
    converted_named_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config)
    buffer_size += param_size
    return buffer_size

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def _update_weight_from_distributed(
    self,
    name: str,
    param: torch.nn.Parameter,
    converted_named_tensors: list[tuple[str, torch.Tensor]],
    buffer_size: int,
    pbar: tqdm | None = None,
) -> int | None:
    """
    Non-expert: gather TP → rm pad → HF → buffer (flush if full). All gather, PP source buffers.
    Returns updated bytes on source, None on non-source.
    """
    param = all_gather_param(name, param)
    if not self._is_pp_src_rank:
        return
 
    param_size = param.numel() * param.element_size()
    if buffer_size + param_size > self.args.update_weight_buffer_size:
        self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar)
        buffer_size = 0
    converted_named_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config)
    buffer_size += param_size
    return buffer_size

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def _update_expert_weight_from_distributed(
    self,
    name: str,
    param: torch.nn.Parameter,
    named_tensors: list[tuple[str, torch.Tensor]],
    buffer_size: int,
    pbar: tqdm | None = None,
) -> int:
    """
    Expert: gather TP → rm pad → buffer. EP gather + HF deferred. Threshold × EP size.
    """
    param = all_gather_param(name, param)
 
    param_size = param.numel() * param.element_size()
    if (
        buffer_size + param_size
    ) * mpu.get_expert_model_parallel_world_size() > self.args.update_weight_buffer_size:
        self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar)
        buffer_size = 0
 
    named_tensors.append((name, param))
    buffer_size += param_size
    return buffer_size
 
def _update_expert_bucket_weights_from_distributed(
    self, named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None
) -> None:
    """
    Gather EP → HF → broadcast. Clears buffer.
    """
    names = [name for name, _ in named_tensors]
    all_names = [None] * mpu.get_expert_model_parallel_world_size()
    dist.all_gather_object(all_names, names, group=mpu.get_expert_model_parallel_group())
 
    for names in all_names:
        assert len(named_tensors) == len(names), f"mismatch names length: {len(named_tensors)} != {len(names)}"
 
    all_gathered_params = [[] for _ in range(mpu.get_expert_model_parallel_world_size())]
    handles = []
    for i, (_name, param) in enumerate(named_tensors):
        params = [
            torch.empty_like(param.data, device=torch.cuda.current_device())
            for _ in range(mpu.get_expert_model_parallel_world_size())
        ]
        handle = dist.all_gather(params, param.data, group=mpu.get_expert_model_parallel_group(), async_op=True)
        handles.append(handle)
        for ep_rank, names in enumerate(all_names):
            all_gathered_params[ep_rank].append((names[i], params[ep_rank]))
    for handle in handles:
        handle.wait()
 
    named_tensors.clear()
    if not self._is_pp_src_rank:
        return
 
    all_gathered_params = sum(all_gathered_params, [])
    converted_hf_tensors = []
    for name, param in all_gathered_params:
        converted_hf_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config)
 
    self._update_bucket_weights_from_distributed(converted_hf_tensors, pbar)

# slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
def _update_expert_weight_from_distributed(
    self,
    name: str,
    param: torch.nn.Parameter,
    named_tensors: list[tuple[str, torch.Tensor]],
    buffer_size: int,
    pbar: tqdm | None = None,
) -> int:
    """
    Expert: gather TP → rm pad → buffer. EP gather + HF deferred. Threshold × EP size.
    """
    param = all_gather_param(name, param)
 
    param_size = param.numel() * param.element_size()
    if (
        buffer_size + param_size
    ) * mpu.get_expert_model_parallel_world_size() > self.args.update_weight_buffer_size:
        self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar)
        buffer_size = 0
 
    named_tensors.append((name, param))
    buffer_size += param_size
    return buffer_size
 
def _update_expert_bucket_weights_from_distributed(
    self, named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None
) -> None:
    """
    Gather EP → HF → broadcast. Clears buffer.
    """
    names = [name for name, _ in named_tensors]
    all_names = [None] * mpu.get_expert_model_parallel_world_size()
    dist.all_gather_object(all_names, names, group=mpu.get_expert_model_parallel_group())
 
    for names in all_names:
        assert len(named_tensors) == len(names), f"mismatch names length: {len(named_tensors)} != {len(names)}"
 
    all_gathered_params = [[] for _ in range(mpu.get_expert_model_parallel_world_size())]
    handles = []
    for i, (_name, param) in enumerate(named_tensors):
        params = [
            torch.empty_like(param.data, device=torch.cuda.current_device())
            for _ in range(mpu.get_expert_model_parallel_world_size())
        ]
        handle = dist.all_gather(params, param.data, group=mpu.get_expert_model_parallel_group(), async_op=True)
        handles.append(handle)
        for ep_rank, names in enumerate(all_names):
            all_gathered_params[ep_rank].append((names[i], params[ep_rank]))
    for handle in handles:
        handle.wait()
 
    named_tensors.clear()
    if not self._is_pp_src_rank:
        return
 
    all_gathered_params = sum(all_gathered_params, [])
    converted_hf_tensors = []
    for name, param in all_gathered_params:
        converted_hf_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config)
 
    self._update_bucket_weights_from_distributed(converted_hf_tensors, pbar)

# python/sglang/srt/entrypoints/engine.py
def init_weights_update_group(
    self,
    master_address: str,
    master_port: int,
    rank_offset: int,
    world_size: int,
    group_name: str,
    backend: str = "nccl",
):
    """Initialize parameter update group."""
    obj = InitWeightsUpdateGroupReqInput(
        master_address=master_address,
        master_port=master_port,
        rank_offset=rank_offset,
        world_size=world_size,
        group_name=group_name,
        backend=backend,
    )
    return self.loop.run_until_complete(
        self.tokenizer_manager.init_weights_update_group(obj, None)
    )
 
def update_weights_from_distributed(
    self,
    names: list[str],
    dtypes: list[str],
    shapes: list[list[int]],
    group_name: str = "weight_update_group",
    flush_cache: bool = True,
    load_format: Optional[str] = None,
):
    """Update weights from distributed source."""
    obj = UpdateWeightsFromDistributedReqInput(
        names=names,
        dtypes=dtypes,
        shapes=shapes,
        group_name=group_name,
        flush_cache=flush_cache,
        load_format=load_format,
    )
    return self.loop.run_until_complete(
        self.tokenizer_manager.update_weights_from_distributed(obj, None)
    )
 
def update_weights_from_tensor(
    self,
    named_tensors: List[Tuple[str, torch.Tensor]],
    load_format: Optional[str] = None,
    flush_cache: bool = True,
):
    """Update weights from distributed source. If there are going to be more updates, set `flush_cache` to be false
    to avoid duplicated cache cleaning operation."""
    if load_format == "flattened_bucket":
        serialized_named_tensors = named_tensors
    else:
        serialized_named_tensors = [
            MultiprocessingSerializer.serialize(named_tensors)
            for _ in range(self.server_args.tp_size)
        ]
    obj = UpdateWeightsFromTensorReqInput(
        serialized_named_tensors=serialized_named_tensors,
        load_format=load_format,
        flush_cache=flush_cache,
    )
    return self.loop.run_until_complete(
        self.tokenizer_manager.update_weights_from_tensor(obj, None)
    )

# python/sglang/srt/entrypoints/engine.py
def init_weights_update_group(
    self,
    master_address: str,
    master_port: int,
    rank_offset: int,
    world_size: int,
    group_name: str,
    backend: str = "nccl",
):
    """Initialize parameter update group."""
    obj = InitWeightsUpdateGroupReqInput(
        master_address=master_address,
        master_port=master_port,
        rank_offset=rank_offset,
        world_size=world_size,
        group_name=group_name,
        backend=backend,
    )
    return self.loop.run_until_complete(
        self.tokenizer_manager.init_weights_update_group(obj, None)
    )
 
def update_weights_from_distributed(
    self,
    names: list[str],
    dtypes: list[str],
    shapes: list[list[int]],
    group_name: str = "weight_update_group",
    flush_cache: bool = True,
    load_format: Optional[str] = None,
):
    """Update weights from distributed source."""
    obj = UpdateWeightsFromDistributedReqInput(
        names=names,
        dtypes=dtypes,
        shapes=shapes,
        group_name=group_name,
        flush_cache=flush_cache,
        load_format=load_format,
    )
    return self.loop.run_until_complete(
        self.tokenizer_manager.update_weights_from_distributed(obj, None)
    )
 
def update_weights_from_tensor(
    self,
    named_tensors: List[Tuple[str, torch.Tensor]],
    load_format: Optional[str] = None,
    flush_cache: bool = True,
):
    """Update weights from distributed source. If there are going to be more updates, set `flush_cache` to be false
    to avoid duplicated cache cleaning operation."""
    if load_format == "flattened_bucket":
        serialized_named_tensors = named_tensors
    else:
        serialized_named_tensors = [
            MultiprocessingSerializer.serialize(named_tensors)
            for _ in range(self.server_args.tp_size)
        ]
    obj = UpdateWeightsFromTensorReqInput(
        serialized_named_tensors=serialized_named_tensors,
        load_format=load_format,
        flush_cache=flush_cache,
    )
    return self.loop.run_until_complete(
        self.tokenizer_manager.update_weights_from_tensor(obj, None)
    )

# slime/backends/sglang_utils/sglang_engine.py
def update_weights_from_tensor(
    self,
    serialized_named_tensors: list[str],
    load_format: str | None = None,
    flush_cache: bool = False,
    weight_version: str | None = None,
):
    payload = {
        "serialized_named_tensors": serialized_named_tensors,
        "load_format": load_format,
        "flush_cache": flush_cache,
    }
    if weight_version is not None:
        payload["weight_version"] = weight_version
    return self._make_request(
        "update_weights_from_tensor",
        payload,
    )
 
def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend):
    return self._make_request(
        "init_weights_update_group",
        {
            "master_address": master_address,
            "master_port": master_port,
            "rank_offset": rank_offset,
            "world_size": world_size,
            "group_name": group_name,
            "backend": backend,
        },
    )
 
def update_weights_from_distributed(
    self, names, dtypes, shapes, group_name, flush_cache=False, weight_version: str | None = None
):
    payload = {
        "names": names,
        "dtypes": [str(dtype).replace("torch.", "") for dtype in dtypes],
        "shapes": shapes,
        "group_name": group_name,
        "flush_cache": flush_cache,
    }
    if weight_version is not None:
        payload["weight_version"] = weight_version
    return self._make_request(
        "update_weights_from_distributed",
        payload,
    )
 
def pause_generation(self):
    response = requests.post(f"http://{self.server_host}:{self.server_port}/pause_generation", json={})
    response.raise_for_status()
    return response
 
def continue_generation(self):
    response = requests.post(f"http://{self.server_host}:{self.server_port}/continue_generation", json={})
    response.raise_for_status()
    return response
 
def post_process_weights(
    self,
    restore_weights_before_load: bool = False,
    post_process_quantization: bool = False,
):
    return self._make_request(
        "post_process_weights",
        {
            "restore_weights_before_load": restore_weights_before_load,
            "post_process_quantization": post_process_quantization,
        },
    )

# slime/backends/sglang_utils/sglang_engine.py
def update_weights_from_tensor(
    self,
    serialized_named_tensors: list[str],
    load_format: str | None = None,
    flush_cache: bool = False,
    weight_version: str | None = None,
):
    payload = {
        "serialized_named_tensors": serialized_named_tensors,
        "load_format": load_format,
        "flush_cache": flush_cache,
    }
    if weight_version is not None:
        payload["weight_version"] = weight_version
    return self._make_request(
        "update_weights_from_tensor",
        payload,
    )
 
def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend):
    return self._make_request(
        "init_weights_update_group",
        {
            "master_address": master_address,
            "master_port": master_port,
            "rank_offset": rank_offset,
            "world_size": world_size,
            "group_name": group_name,
            "backend": backend,
        },
    )
 
def update_weights_from_distributed(
    self, names, dtypes, shapes, group_name, flush_cache=False, weight_version: str | None = None
):
    payload = {
        "names": names,
        "dtypes": [str(dtype).replace("torch.", "") for dtype in dtypes],
        "shapes": shapes,
        "group_name": group_name,
        "flush_cache": flush_cache,
    }
    if weight_version is not None:
        payload["weight_version"] = weight_version
    return self._make_request(
        "update_weights_from_distributed",
        payload,
    )
 
def pause_generation(self):
    response = requests.post(f"http://{self.server_host}:{self.server_port}/pause_generation", json={})
    response.raise_for_status()
    return response
 
def continue_generation(self):
    response = requests.post(f"http://{self.server_host}:{self.server_port}/continue_generation", json={})
    response.raise_for_status()
    return response
 
def post_process_weights(
    self,
    restore_weights_before_load: bool = False,
    post_process_quantization: bool = False,
):
    return self._make_request(
        "post_process_weights",
        {
            "restore_weights_before_load": restore_weights_before_load,
            "post_process_quantization": post_process_quantization,
        },
    )

# python/sglang/srt/managers/tokenizer_manager.py
async def pause_generation(self, obj: PauseGenerationReqInput):
    async with self.is_pause_cond:
        self.is_pause = True
        if obj.mode != "abort":
            await self.send_to_scheduler.send_pyobj(obj)
        else:
            # we are using the model_update_lock to check if there is still on-going requests.
            while True:
                # TODO: maybe make it async instead of fire-and-forget
                self.abort_request(abort_all=True)
                is_locked = await self.model_update_lock.is_locked()
                if not is_locked:
                    break
                await asyncio.sleep(1.0)
 
async def continue_generation(self, obj: ContinueGenerationReqInput):
    async with self.is_pause_cond:
        self.is_pause = False
        await self.send_to_scheduler.send_pyobj(obj)
        self.is_pause_cond.notify_all()

# python/sglang/srt/managers/tokenizer_manager.py
async def pause_generation(self, obj: PauseGenerationReqInput):
    async with self.is_pause_cond:
        self.is_pause = True
        if obj.mode != "abort":
            await self.send_to_scheduler.send_pyobj(obj)
        else:
            # we are using the model_update_lock to check if there is still on-going requests.
            while True:
                # TODO: maybe make it async instead of fire-and-forget
                self.abort_request(abort_all=True)
                is_locked = await self.model_update_lock.is_locked()
                if not is_locked:
                    break
                await asyncio.sleep(1.0)
 
async def continue_generation(self, obj: ContinueGenerationReqInput):
    async with self.is_pause_cond:
        self.is_pause = False
        await self.send_to_scheduler.send_pyobj(obj)
        self.is_pause_cond.notify_all()

# python/sglang/srt/managers/tokenizer_communicator_mixin.py
async def update_weights_from_distributed(
    self: TokenizerManager,
    obj: UpdateWeightsFromDistributedReqInput,
    request: Optional[fastapi.Request] = None,
) -> Tuple[bool, str]:
    self.auto_create_handle_loop()
    assert (
        self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
    ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
 
    if obj.abort_all_requests:
        self.abort_request(abort_all=True)
 
    # Immediately update the weights if the engine is in paused state
    async with self.is_pause_cond:
        is_paused = self.is_pause
 
    lock_context = (
        self.model_update_lock.writer_lock if not is_paused else nullcontext()
    )
    async with lock_context:
        results = await self.update_weights_from_distributed_communicator(obj)
 
    success, message = _Communicator.merge_results(results)
    if success and obj.weight_version is not None:
        self._update_weight_version_if_provided(obj.weight_version)
        message += f" Weight version updated to {obj.weight_version}."
 
    return success, message
 
async def update_weights_from_tensor(
    self: TokenizerManager,
    obj: UpdateWeightsFromTensorReqInput,
    request: Optional[fastapi.Request] = None,
) -> Tuple[bool, str]:
    self.auto_create_handle_loop()
    assert (
        self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
    ), "dp_size must be 1 or dp attention must be enabled for update weights from tensor"
 
    if obj.abort_all_requests:
        self.abort_request(abort_all=True)
 
    # Immediately update the weights if the engine is in paused state
    async with self.is_pause_cond:
        is_paused = self.is_pause
 
    lock_context = (
        self.model_update_lock.writer_lock if not is_paused else nullcontext()
    )
    async with lock_context:
        results = await self.update_weights_from_tensor_communicator(obj)
 
    success, message = _Communicator.merge_results(results)
    if success and obj.weight_version is not None:
        self._update_weight_version_if_provided(obj.weight_version)
        message += f" Weight version updated to {obj.weight_version}."
 
    return success, message

# python/sglang/srt/managers/tokenizer_communicator_mixin.py
async def update_weights_from_distributed(
    self: TokenizerManager,
    obj: UpdateWeightsFromDistributedReqInput,
    request: Optional[fastapi.Request] = None,
) -> Tuple[bool, str]:
    self.auto_create_handle_loop()
    assert (
        self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
    ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
 
    if obj.abort_all_requests:
        self.abort_request(abort_all=True)
 
    # Immediately update the weights if the engine is in paused state
    async with self.is_pause_cond:
        is_paused = self.is_pause
 
    lock_context = (
        self.model_update_lock.writer_lock if not is_paused else nullcontext()
    )
    async with lock_context:
        results = await self.update_weights_from_distributed_communicator(obj)
 
    success, message = _Communicator.merge_results(results)
    if success and obj.weight_version is not None:
        self._update_weight_version_if_provided(obj.weight_version)
        message += f" Weight version updated to {obj.weight_version}."
 
    return success, message
 
async def update_weights_from_tensor(
    self: TokenizerManager,
    obj: UpdateWeightsFromTensorReqInput,
    request: Optional[fastapi.Request] = None,
) -> Tuple[bool, str]:
    self.auto_create_handle_loop()
    assert (
        self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
    ), "dp_size must be 1 or dp attention must be enabled for update weights from tensor"
 
    if obj.abort_all_requests:
        self.abort_request(abort_all=True)
 
    # Immediately update the weights if the engine is in paused state
    async with self.is_pause_cond:
        is_paused = self.is_pause
 
    lock_context = (
        self.model_update_lock.writer_lock if not is_paused else nullcontext()
    )
    async with lock_context:
        results = await self.update_weights_from_tensor_communicator(obj)
 
    success, message = _Communicator.merge_results(results)
    if success and obj.weight_version is not None:
        self._update_weight_version_if_provided(obj.weight_version)
        message += f" Weight version updated to {obj.weight_version}."
 
    return success, message

# python/sglang/srt/managers/tokenizer_manager.py
# Build meta_info and return value
meta_info = {
    "id": rid,
    "finish_reason": recv_obj.finished_reasons[i],
    "prompt_tokens": recv_obj.prompt_tokens[i],
    "weight_version": self.server_args.weight_version,
    "total_retractions": recv_obj.retraction_counts[i],
}

# python/sglang/srt/managers/tokenizer_manager.py
# Build meta_info and return value
meta_info = {
    "id": rid,
    "finish_reason": recv_obj.finished_reasons[i],
    "prompt_tokens": recv_obj.prompt_tokens[i],
    "weight_version": self.server_args.weight_version,
    "total_retractions": recv_obj.retraction_counts[i],
}

# python/sglang/srt/managers/scheduler_update_weights_mixin.py
def update_weights_from_distributed(
    self,
    recv_req: UpdateWeightsFromDistributedReqInput,
) -> Tuple[bool, str]:
    """Update the online model parameter."""
    success, message = self.tp_worker.update_weights_from_distributed(recv_req)
    if success:
        if recv_req.flush_cache:
            flush_cache_success = self.flush_cache()
            assert flush_cache_success, "Cache flush failed after updating weights"
    else:
        logger.error(message)
    return UpdateWeightsFromDistributedReqOutput(success, message)
 
def update_weights_from_tensor(
    self: Scheduler, recv_req: UpdateWeightsFromTensorReqInput
):
    """Update the online model parameter from tensors."""
    worker = self.draft_worker or self.tp_worker
    success, message = worker.update_weights_from_tensor(recv_req)
    # TODO extract common code b/t update_weights_from_distributed and update_weights_from_tensor later
    if success:
        if recv_req.flush_cache:
            flush_cache_success = self.flush_cache()
            assert flush_cache_success, "Cache flush failed after updating weights"
    else:
        logger.error(message)
    torch.distributed.barrier(group=self.tp_cpu_group)
    return UpdateWeightsFromTensorReqOutput(success, message)