586 lines
21 KiB
Python
586 lines
21 KiB
Python
#
|
|
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import time
|
|
from bisect import bisect_left
|
|
from typing import Any
|
|
|
|
import requests
|
|
import torch
|
|
from torch import nn
|
|
from vllm.logger import logger
|
|
from vllm.utils.network_utils import get_ip, get_open_port, join_host_port
|
|
|
|
MAX_TRANSFER_CHUNK_BYTES = 1024**3
|
|
MAX_TRANSFER_CHUNK_WEIGHTS = 512
|
|
|
|
|
|
def _normalize_weight_shape(shape: Any) -> tuple[int, ...] | None:
|
|
if shape is None:
|
|
return None
|
|
if not isinstance(shape, (list, tuple)):
|
|
return None
|
|
if not all(isinstance(dim, int) and dim >= 0 for dim in shape):
|
|
return None
|
|
return tuple(shape)
|
|
|
|
|
|
def _parse_weight_info(weight_info: Any):
|
|
if not isinstance(weight_info, (list, tuple)) or len(weight_info) not in (3, 4):
|
|
return None
|
|
|
|
seed_ptr, seed_len, seed_size = weight_info[:3]
|
|
if not all(isinstance(value, int) for value in (seed_ptr, seed_len, seed_size)):
|
|
return None
|
|
|
|
seed_shape = None
|
|
if len(weight_info) == 4:
|
|
seed_shape = _normalize_weight_shape(weight_info[3])
|
|
if seed_shape is None:
|
|
return None
|
|
|
|
return seed_ptr, seed_len, seed_size, seed_shape
|
|
|
|
|
|
def _reshape_tensor_to_seed_shape(
|
|
name: str,
|
|
tensor: torch.Tensor,
|
|
seed_shape: tuple[int, ...] | None,
|
|
reshape_events: list[tuple[str, tuple[int, ...], tuple[int, ...]]] | None = None,
|
|
) -> bool:
|
|
if seed_shape is None or tuple(tensor.shape) == seed_shape:
|
|
return True
|
|
|
|
if tensor.numel() != _numel_from_shape(seed_shape):
|
|
logger.error(
|
|
"Weight shape mismatch for %s, local shape %s cannot view as seed shape %s",
|
|
name,
|
|
tuple(tensor.shape),
|
|
seed_shape,
|
|
)
|
|
return False
|
|
|
|
local_shape = tuple(tensor.shape)
|
|
try:
|
|
tensor.data = tensor.data.view(seed_shape)
|
|
except Exception as e:
|
|
logger.error(
|
|
"Failed to reshape RFork tensor %s from %s to seed shape %s: %s",
|
|
name,
|
|
local_shape,
|
|
seed_shape,
|
|
e,
|
|
)
|
|
return False
|
|
|
|
if reshape_events is not None:
|
|
reshape_events.append((name, local_shape, seed_shape))
|
|
return True
|
|
|
|
|
|
def _update_registered_weight_shape(
|
|
weight_shape_dict: dict[str, tuple[int, ...]] | None,
|
|
name: str,
|
|
tensor: torch.Tensor,
|
|
) -> None:
|
|
if isinstance(weight_shape_dict, dict):
|
|
weight_shape_dict[name] = tuple(tensor.shape)
|
|
|
|
|
|
def _numel_from_shape(shape: tuple[int, ...]) -> int:
|
|
numel = 1
|
|
for dim in shape:
|
|
numel *= dim
|
|
return numel
|
|
|
|
|
|
def _is_transferable_tensor(tensor: torch.Tensor) -> bool:
|
|
return not tensor.is_meta and tensor.numel() > 0 and _is_tensor_on_transfer_device(tensor)
|
|
|
|
|
|
def _is_tensor_on_transfer_device(tensor: torch.Tensor) -> bool:
|
|
return tensor.device.type == "npu"
|
|
|
|
|
|
def _iter_tensors_in_value(prefix: str, value: Any, visited_object_ids: set[int], scan_objects: bool = False):
|
|
if isinstance(value, torch.Tensor):
|
|
yield prefix, value
|
|
return
|
|
|
|
if isinstance(value, (nn.Module, str, bytes)) or callable(value):
|
|
return
|
|
|
|
if isinstance(value, (list, tuple)):
|
|
for index, item in enumerate(value):
|
|
yield from _iter_tensors_in_value(f"{prefix}.{index}", item, visited_object_ids, scan_objects)
|
|
return
|
|
|
|
if isinstance(value, dict):
|
|
for key, item in value.items():
|
|
yield from _iter_tensors_in_value(f"{prefix}.{key}", item, visited_object_ids, scan_objects)
|
|
return
|
|
|
|
if not scan_objects or not hasattr(value, "__dict__"):
|
|
return
|
|
|
|
value_id = id(value)
|
|
if value_id in visited_object_ids:
|
|
return
|
|
visited_object_ids.add(value_id)
|
|
for attr_name, attr_value in vars(value).items():
|
|
if attr_name.startswith("_"):
|
|
continue
|
|
yield from _iter_tensors_in_value(f"{prefix}.{attr_name}", attr_value, visited_object_ids, scan_objects)
|
|
|
|
|
|
def _try_collect_transferable_tensor(
|
|
name: str,
|
|
tensor: torch.Tensor,
|
|
seen_data_ptrs: set[int],
|
|
collected_tensors: list[tuple[str, torch.Tensor]],
|
|
) -> tuple[bool, bool]:
|
|
if not _is_transferable_tensor(tensor):
|
|
return False, False
|
|
|
|
data_ptr = tensor.data_ptr()
|
|
if data_ptr in seen_data_ptrs:
|
|
return False, True
|
|
|
|
seen_data_ptrs.add(data_ptr)
|
|
collected_tensors.append((name, tensor))
|
|
return True, False
|
|
|
|
|
|
def _collect_transferable_tensors(model: nn.Module) -> list[tuple[str, torch.Tensor]]:
|
|
seen_data_ptrs: set[int] = set()
|
|
collected_tensors: list[tuple[str, torch.Tensor]] = []
|
|
|
|
for name, tensor in model.named_parameters():
|
|
_try_collect_transferable_tensor(
|
|
name,
|
|
tensor,
|
|
seen_data_ptrs,
|
|
collected_tensors,
|
|
)
|
|
|
|
for name, tensor in model.named_buffers():
|
|
_try_collect_transferable_tensor(
|
|
name,
|
|
tensor,
|
|
seen_data_ptrs,
|
|
collected_tensors,
|
|
)
|
|
|
|
# Some Ascend post-load paths replace checkpoint parameters with runtime
|
|
# tensors stored as plain module attributes, e.g. MLA/SFA W_UV and W_UK_T.
|
|
for module_prefix, module in model.named_modules():
|
|
for attr_name, attr_value in vars(module).items():
|
|
if attr_name.startswith("_") or isinstance(attr_value, nn.Module):
|
|
continue
|
|
|
|
scan_objects = attr_name == "impl"
|
|
for tensor_name, tensor in _iter_tensors_in_value(attr_name, attr_value, set(), scan_objects):
|
|
full_name = f"{module_prefix}.{tensor_name}" if module_prefix else tensor_name
|
|
_try_collect_transferable_tensor(
|
|
full_name,
|
|
tensor,
|
|
seen_data_ptrs,
|
|
collected_tensors,
|
|
)
|
|
return collected_tensors
|
|
|
|
|
|
def _iter_transferable_tensors(model: nn.Module):
|
|
yield from _collect_transferable_tensors(model)
|
|
|
|
|
|
def _block_contains_weight_ptr(address: int, size: int, sorted_weight_ptrs: list[int]) -> bool:
|
|
index = bisect_left(sorted_weight_ptrs, address)
|
|
return index < len(sorted_weight_ptrs) and sorted_weight_ptrs[index] < address + size
|
|
|
|
|
|
def _iter_transfer_chunks(
|
|
weight_names: list[str],
|
|
seed_ptr_list: list[int],
|
|
client_ptr_list: list[int],
|
|
client_len_list: list[int],
|
|
):
|
|
chunk_start = 0
|
|
chunk_bytes = 0
|
|
chunk_weights = 0
|
|
|
|
for index, length in enumerate(client_len_list):
|
|
should_flush = chunk_weights > 0 and (
|
|
chunk_bytes + length > MAX_TRANSFER_CHUNK_BYTES or chunk_weights >= MAX_TRANSFER_CHUNK_WEIGHTS
|
|
)
|
|
if should_flush:
|
|
yield (
|
|
weight_names[chunk_start:index],
|
|
seed_ptr_list[chunk_start:index],
|
|
client_ptr_list[chunk_start:index],
|
|
client_len_list[chunk_start:index],
|
|
)
|
|
chunk_start = index
|
|
chunk_bytes = 0
|
|
chunk_weights = 0
|
|
|
|
chunk_bytes += length
|
|
chunk_weights += 1
|
|
|
|
if chunk_weights > 0:
|
|
yield (
|
|
weight_names[chunk_start:],
|
|
seed_ptr_list[chunk_start:],
|
|
client_ptr_list[chunk_start:],
|
|
client_len_list[chunk_start:],
|
|
)
|
|
|
|
|
|
class RForkTransferBackend:
|
|
def __init__(self):
|
|
self.rfork_transfer_engine: Any | None = None
|
|
self.rfork_transfer_engine_session_id = None
|
|
self.rfork_transfer_engine_weights_info_dict = None
|
|
self.rfork_transfer_engine_weights_shape_dict = None
|
|
self.registered_weight_blocks = []
|
|
self._registered_transferable_tensors: list[tuple[str, torch.Tensor]] | None = None
|
|
self._is_initialized = False
|
|
self.init_transfer_engine()
|
|
|
|
def init_transfer_engine(self):
|
|
try:
|
|
from yr.datasystem import TransferEngine # type: ignore[import-not-found]
|
|
except ImportError as e:
|
|
err_msg = (
|
|
"Failed to import TransferEngine from yr.datasystem. "
|
|
"Please install @yuanrong-datasystem/transfer_engine."
|
|
)
|
|
logger.error(err_msg)
|
|
raise ImportError(err_msg) from e
|
|
|
|
transfer_engine = TransferEngine()
|
|
local_hostname = join_host_port(get_ip(), get_open_port())
|
|
ret = transfer_engine.initialize(local_hostname, "ascend", f"npu:{torch.npu.current_device()}")
|
|
if ret.is_error():
|
|
err_msg = (
|
|
f"TransferEngine initialization failed: "
|
|
f"initialize({local_hostname}, 'ascend', "
|
|
f"'npu:{int(torch.npu.current_device())}') -> {ret.to_string()}"
|
|
)
|
|
logger.error(err_msg)
|
|
raise RuntimeError(err_msg)
|
|
|
|
self.rfork_transfer_engine = transfer_engine
|
|
self.rfork_transfer_engine_session_id = local_hostname
|
|
self._is_initialized = True
|
|
|
|
def is_initialized(self) -> bool:
|
|
return self._is_initialized
|
|
|
|
def _get_transfer_engine(self) -> Any:
|
|
if self.rfork_transfer_engine is None:
|
|
raise RuntimeError("TransferEngine is not initialized.")
|
|
return self.rfork_transfer_engine
|
|
|
|
def register_memory_region(self, model):
|
|
transfer_engine = self._get_transfer_engine()
|
|
start_reg_mr_time = time.perf_counter()
|
|
self._registered_transferable_tensors = None
|
|
|
|
weight_mr_dict = {}
|
|
weight_shape_dict = {}
|
|
weight_addr_set = set()
|
|
transferable_tensors = list(_iter_transferable_tensors(model))
|
|
for name, weight in transferable_tensors:
|
|
weight_mr_dict[name] = (
|
|
weight.data_ptr(),
|
|
weight.numel(),
|
|
weight.element_size(),
|
|
)
|
|
weight_shape_dict[name] = tuple(weight.shape)
|
|
weight_addr_set.add(weight.data_ptr())
|
|
|
|
sorted_weight_ptrs = sorted(weight_addr_set)
|
|
|
|
memory_snapshot = torch.npu.memory.memory_snapshot()
|
|
|
|
weight_blocks_for_reg_mr = []
|
|
for segment in memory_snapshot:
|
|
current_weight_block = None
|
|
for block in segment.get("blocks", []):
|
|
address = block.get("address", -1)
|
|
size = block.get("size", -1)
|
|
state = block.get("state", "")
|
|
if address < 0 or size < 0 or state == "":
|
|
continue
|
|
if state == "active_allocated" and _block_contains_weight_ptr(address, size, sorted_weight_ptrs):
|
|
if current_weight_block is None:
|
|
current_weight_block = (address, size)
|
|
elif current_weight_block[0] + current_weight_block[1] == address:
|
|
current_weight_block = (
|
|
current_weight_block[0],
|
|
current_weight_block[1] + size,
|
|
)
|
|
else:
|
|
weight_blocks_for_reg_mr.append(current_weight_block)
|
|
current_weight_block = (address, size)
|
|
if current_weight_block is not None:
|
|
weight_blocks_for_reg_mr.append(current_weight_block)
|
|
|
|
addresses, sizes = zip(*weight_blocks_for_reg_mr) if weight_blocks_for_reg_mr else ((), ())
|
|
ret = transfer_engine.batch_register_memory(addresses, sizes)
|
|
if ret.is_error():
|
|
self._registered_transferable_tensors = None
|
|
logger.error(
|
|
"batch_register_memory failed for %d blocks, ret: %s",
|
|
len(weight_blocks_for_reg_mr),
|
|
ret.to_string(),
|
|
)
|
|
return False
|
|
|
|
self.rfork_transfer_engine_weights_info_dict = weight_mr_dict
|
|
self.rfork_transfer_engine_weights_shape_dict = weight_shape_dict
|
|
self.registered_weight_blocks = weight_blocks_for_reg_mr
|
|
self._registered_transferable_tensors = transferable_tensors
|
|
logger.info(
|
|
"register_memory_region time: %.4fs, weights: %d",
|
|
time.perf_counter() - start_reg_mr_time,
|
|
len(weight_mr_dict),
|
|
)
|
|
return True
|
|
|
|
def unregister_memory_region(self) -> bool:
|
|
transfer_engine = self._get_transfer_engine()
|
|
start_unreg_mr_time = time.perf_counter()
|
|
if not self.registered_weight_blocks:
|
|
self.rfork_transfer_engine_weights_info_dict = None
|
|
self.rfork_transfer_engine_weights_shape_dict = None
|
|
self._registered_transferable_tensors = None
|
|
logger.debug("unregister_memory_region skipped because no blocks are registered.")
|
|
return True
|
|
|
|
ret = transfer_engine.batch_unregister_memory([address for address, _ in self.registered_weight_blocks])
|
|
if ret.is_error():
|
|
logger.error(
|
|
"batch_unregister_memory failed for %d blocks, ret: %s",
|
|
len(self.registered_weight_blocks),
|
|
ret.to_string(),
|
|
)
|
|
return False
|
|
self.rfork_transfer_engine_weights_info_dict = None
|
|
self.rfork_transfer_engine_weights_shape_dict = None
|
|
self.registered_weight_blocks = []
|
|
self._registered_transferable_tensors = None
|
|
logger.info(
|
|
"unregister_memory_region time: %.4fs",
|
|
time.perf_counter() - start_unreg_mr_time,
|
|
)
|
|
return True
|
|
|
|
def recv_from_source(
|
|
self,
|
|
model,
|
|
seed_instance_ip,
|
|
seed_instance_service_port,
|
|
local_seed_key,
|
|
):
|
|
transfer_engine = self._get_transfer_engine()
|
|
seed_url = f"http://{seed_instance_ip}:{seed_instance_service_port}"
|
|
seed_session_id, seed_weight_info, seed_weight_shapes = get_remote_instance_transfer_engine_info(
|
|
seed_url,
|
|
local_seed_key,
|
|
)
|
|
if seed_session_id is None or seed_weight_info is None:
|
|
self._registered_transferable_tensors = None
|
|
logger.error("Cannot get transfer engine session or weight info.")
|
|
return False
|
|
|
|
transferable_tensors = getattr(self, "_registered_transferable_tensors", None)
|
|
if transferable_tensors is None:
|
|
transferable_tensors = list(_iter_transferable_tensors(model))
|
|
|
|
seed_ptr_list = []
|
|
client_ptr_list = []
|
|
client_len_list = []
|
|
weight_names = []
|
|
reshape_events: list[tuple[str, tuple[int, ...], tuple[int, ...]]] = []
|
|
try:
|
|
for name, tensor in transferable_tensors:
|
|
weight_info = seed_weight_info.get(name, None)
|
|
if weight_info is None:
|
|
logger.error("Cannot find weight info for %s.", name)
|
|
return False
|
|
|
|
parsed_weight_info = _parse_weight_info(weight_info)
|
|
if parsed_weight_info is None:
|
|
logger.error("Invalid weight info for %s: %s", name, weight_info)
|
|
return False
|
|
|
|
seed_ptr, seed_len, seed_size, seed_shape = parsed_weight_info
|
|
if seed_shape is None and isinstance(seed_weight_shapes, dict):
|
|
seed_shape = _normalize_weight_shape(seed_weight_shapes.get(name))
|
|
if seed_len != tensor.numel() or seed_size != tensor.element_size():
|
|
logger.error(
|
|
"Weight info mismatch for %s, expected (%s, %s), got (%s, %s)",
|
|
name,
|
|
seed_len,
|
|
seed_size,
|
|
tensor.numel(),
|
|
tensor.element_size(),
|
|
)
|
|
return False
|
|
|
|
if not _reshape_tensor_to_seed_shape(name, tensor, seed_shape, reshape_events):
|
|
return False
|
|
_update_registered_weight_shape(
|
|
self.rfork_transfer_engine_weights_shape_dict,
|
|
name,
|
|
tensor,
|
|
)
|
|
|
|
seed_ptr_list.append(seed_ptr)
|
|
client_ptr_list.append(tensor.data_ptr())
|
|
client_len_list.append(tensor.numel() * tensor.element_size())
|
|
weight_names.append(name)
|
|
finally:
|
|
self._registered_transferable_tensors = None
|
|
transferable_tensors = None
|
|
|
|
if reshape_events:
|
|
sample_events = ", ".join(
|
|
f"{name}: {local_shape}->{seed_shape}" for name, local_shape, seed_shape in reshape_events[:3]
|
|
)
|
|
if len(reshape_events) > 3:
|
|
sample_events += ", ..."
|
|
logger.debug(
|
|
"RFork reshaped %d tensors to match seed shapes: %s",
|
|
len(reshape_events),
|
|
sample_events,
|
|
)
|
|
|
|
transfer_chunks = list(
|
|
_iter_transfer_chunks(
|
|
weight_names,
|
|
seed_ptr_list,
|
|
client_ptr_list,
|
|
client_len_list,
|
|
)
|
|
)
|
|
total_transfer_bytes = sum(client_len_list)
|
|
|
|
transfer_start_time = time.perf_counter()
|
|
logger.info(
|
|
"transfer weights starts, weights: %d, chunks: %d, total bytes: %.2f GiB",
|
|
len(client_len_list),
|
|
len(transfer_chunks),
|
|
total_transfer_bytes / (1024**3),
|
|
)
|
|
for index, (chunk_names, chunk_seed_ptrs, chunk_client_ptrs, chunk_lengths) in enumerate(transfer_chunks, 1):
|
|
chunk_start_time = time.perf_counter()
|
|
logger.debug(
|
|
"transfer weights chunk %d/%d starts, weights: %d, bytes: %.2f GiB, first: %s, last: %s",
|
|
index,
|
|
len(transfer_chunks),
|
|
len(chunk_lengths),
|
|
sum(chunk_lengths) / (1024**3),
|
|
chunk_names[0],
|
|
chunk_names[-1],
|
|
)
|
|
ret = transfer_engine.batch_transfer_sync_read(
|
|
seed_session_id,
|
|
chunk_client_ptrs,
|
|
chunk_seed_ptrs,
|
|
chunk_lengths,
|
|
)
|
|
if ret.is_error():
|
|
logger.error(
|
|
"Failed to transfer weights chunk %d/%d, first: %s, last: %s, ret=%s",
|
|
index,
|
|
len(transfer_chunks),
|
|
chunk_names[0],
|
|
chunk_names[-1],
|
|
ret.to_string(),
|
|
)
|
|
return False
|
|
logger.debug(
|
|
"transfer weights chunk %d/%d done, time: %.4fs",
|
|
index,
|
|
len(transfer_chunks),
|
|
time.perf_counter() - chunk_start_time,
|
|
)
|
|
transfer_time = time.perf_counter() - transfer_start_time
|
|
logger.info("transfer weights time: %.4fs", transfer_time)
|
|
return True
|
|
|
|
|
|
def get_remote_instance_transfer_engine_info(seed_url: str, local_seed_key: str):
|
|
try:
|
|
response = requests.get(
|
|
f"{seed_url}/get_rfork_transfer_engine_info",
|
|
params={"seed_key": local_seed_key},
|
|
)
|
|
if response.status_code != 200:
|
|
logger.error(
|
|
"GET %s/get_rfork_transfer_engine_info failed: %s",
|
|
seed_url,
|
|
response.status_code,
|
|
)
|
|
return None, None, None
|
|
|
|
data = response.json()
|
|
info = data.get("rfork_transfer_engine_info", None)
|
|
if info is not None and isinstance(info, list) and len(info) == 2:
|
|
shape_info = get_remote_instance_weight_shape_info(seed_url, local_seed_key)
|
|
return info[0], info[1], shape_info
|
|
|
|
logger.error(
|
|
"Failed to get rfork_transfer_engine_info in response from %s.",
|
|
seed_url,
|
|
)
|
|
return None, None, None
|
|
except Exception as e:
|
|
logger.error("Exception getting transfer engine info from %s: %s", seed_url, e)
|
|
return None, None, None
|
|
|
|
|
|
def get_remote_instance_weight_shape_info(seed_url: str, local_seed_key: str):
|
|
try:
|
|
response = requests.get(
|
|
f"{seed_url}/get_rfork_transfer_engine_shape_info",
|
|
params={"seed_key": local_seed_key},
|
|
)
|
|
if response.status_code != 200:
|
|
logger.debug(
|
|
"GET %s/get_rfork_transfer_engine_shape_info failed: %s",
|
|
seed_url,
|
|
response.status_code,
|
|
)
|
|
return None
|
|
|
|
data = response.json()
|
|
info = data.get("rfork_transfer_engine_shape_info", None)
|
|
if info is None or isinstance(info, dict):
|
|
return info
|
|
|
|
logger.error(
|
|
"Failed to get rfork_transfer_engine_shape_info in response from %s.",
|
|
seed_url,
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
logger.debug("Exception getting transfer engine shape info from %s: %s", seed_url, e)
|
|
return None
|