feat: add fastsafetensors loader

Author: lixin010

rtp_llm/BUILD

@@ -97,6 +97,7 @@ requirement([
     "portalocker",
     "concurrent_log_handler",
     "aiter",
+    "fastsafetensors",
 ] + tensorrt)
 
 filegroup(

rtp_llm/distribute/gang_server.py

@@ -13,6 +13,7 @@
 import requests
 import uvicorn
 from fastapi import FastAPI
+import torch.distributed
 
 from rtp_llm.config.py_config_modules import PyEnvConfigs, StaticConfig
 from rtp_llm.config.uvicorn_config import UVICORN_LOGGING_CONFIG
@@ -423,9 +424,16 @@ def start(self):
         master_url = (
             f"tcp://{g_master_info.ip}:{self._gang_info.master.server_port - 1}"
         )
-        logging.info(f"gang worker {g_parallel_info} memory_barrier {master_url}")
+        logging.info(f"gang worker {g_parallel_info} init_process_group {master_url}")
         init_process_timeout = self.py_env_configs.gang_config.dist_barrier_timeout
-        self.memory_barrier(master_url, timeout=init_process_timeout)
+        os.environ["TORCH_DIST_INIT_BARRIER"] = "1"
+        torch.distributed.init_process_group(
+            backend=torch.distributed.Backend.NCCL,
+            init_method=master_url,
+            rank=g_parallel_info.world_rank,
+            world_size=g_parallel_info.world_size,
+            timeout=timedelta(seconds=init_process_timeout),
+        )
 
         logging.info(f"gang worker {g_parallel_info} start_health_check")
         self.start_health_check()

rtp_llm/model_loader/ffn_weight.py

@@ -331,6 +331,58 @@ def _load_raw_tensor(
         after_merge_tensor = self.process_fun(before_merge_tensors).to(convert_type)
         logging.debug("load weight :%s, %s ", self.name, after_merge_tensor.shape)
         return {self.name: after_merge_tensor}
+    
+    def get_tensor_names(
+        self, layer_id: Optional[int], load_config: LoadConfig
+    ) -> set[str]:
+        if self.config.weight_stack:
+            return super().get_tensor_names(layer_id, load_config)
+        names = set[str]()
+        for ckpt_weight in self.weights:
+            selected_experts = load_config.get_selected_experts(
+                layer_id, self.config.expert_num
+            )
+            for expert_id in selected_experts:
+                name = ckpt_weight.name.format(
+                    i=str(layer_id), i_1=str(layer_id + 1), expert_id=str(expert_id)
+                )
+                names.add(name)
+        return names
+
+    def _process_raw_tensors(
+        self,
+        raw_tensors: Dict[str, torch.Tensor],
+        layer_id: Optional[int],
+        device: str,
+        load_config: LoadConfig,
+    ):
+        if self.config.weight_stack:
+            return super()._process_raw_tensors(
+                raw_tensors, layer_id, device, load_config
+            )
+        before_merge_tensors: List[torch.Tensor] = []
+        convert_type = (
+            self.data_type if self.data_type is not None else load_config.compute_dtype
+        )
+        for ckpt_weight in self.weights:
+            selected_experts = load_config.get_selected_experts(
+                layer_id, self.config.expert_num
+            )
+            for expert_id in selected_experts:
+                name = ckpt_weight.name.format(
+                    i=str(layer_id), i_1=str(layer_id + 1), expert_id=str(expert_id)
+                )
+                try:
+                    before_merge_tensors.append(
+                        ckpt_weight.merge_fun([raw_tensors[name]])
+                    )
+                except Exception as e:
+                    logging.error(
+                        f"加载 {self.name}: {name} 失败，完整堆栈:\n{traceback.format_exc()}"
+                    )
+                    raise e
+        after_merge_tensor = self.process_fun(before_merge_tensors).to(convert_type)
+        return {self.name: after_merge_tensor}
 
 
 class MoeWeight(CompositeWeight):

rtp_llm/model_loader/loader.py

@@ -1,8 +1,9 @@
 import gc
 import logging
 import os
+import time
 from collections import OrderedDict
-from typing import Optional
+from typing import Dict, Optional
 
 import safetensors
 import torch
@@ -66,7 +67,8 @@ def load_weights(self, device: str):
         if self._load_config.is_ft_style_weight:
             weights = self._load_from_ft_style(device)
         else:
-            weights = self._load_from_scratch(device)
+            weights = self._load_weight(device)
+            self.force_clean_cuda_memory()
 
         # load dynamic weight
         self._load_dynamic_weights(weights, device)
@@ -203,6 +205,88 @@ def _load_from_ft_style(self, device: str):
         model_weights.global_weights = global_weights
         return model_weights
 
+    def _load_weight(self, device: str):
+        is_safetensor = self._load_config.database.is_safetensor
+        convert_device = self._choose_weight_convert_device(device)
+        if is_safetensor and convert_device != "cpu" and self._is_memory_enough_for_fastsafetensor():
+            return self._load_from_fastsafetensor(device)
+        
+        logging.info(
+            f"database is safetensor: {is_safetensor}, device: {device}, choose devie: {convert_device}"
+        )
+        return self._load_from_scratch(device)
+    
+    def _is_memory_enough_for_fastsafetensor(self):
+        model_size = self._weights_info.config.eval_model_size()
+        device_mem_info = self._load_config.exported_device.get_mem_info()
+        max_file_size = self._load_config.database.get_max_file_size()
+        if device_mem_info is None:
+            return False
+        else:
+            free_mem = device_mem_info.free / (1024.0**2)
+        model_mem = model_size / self._load_config.tp_size / (1024.0**2)
+        max_file_mem = max_file_size / (1024.0**2)
+        logging.debug(f"free mem: {free_mem}, model mem: {model_mem}, max file mem: {max_file_mem}")
+        return (free_mem - model_mem) > (3 * max_file_mem)
+    
+    def _load_from_fastsafetensor(self, device: str):
+        try:
+            all_tensors = self._load_config.database.fastsafetensors_weights_iterator(
+                device, True
+            )
+        except (ModuleNotFoundError, ImportError) as e:
+            logging.warning(f"Failed to import fastsafetensors: {e}")
+            return self._load_from_scratch(device)
+
+        logging.info(f"load weight by device: {device}")
+        model_weights = self._create_model_weights(device)
+        tensor_to_weight_map = self._get_tensor_to_weight_map()
+        direct_io = self._load_config.exported_device.support_dio_load
+        for key, loaded_tensor in all_tensors:
+            if key not in tensor_to_weight_map:
+                continue
+            layer_id, weight = tensor_to_weight_map[key]
+            start = time.time()
+            res, complete = weight.add_tensor(
+                key, layer_id, loaded_tensor, device, self._load_config
+            )
+            logging.debug(
+                f"weight: {type(weight).__name__} add tensor, complete: {complete}, cost {time.time() - start}"
+            )
+            if complete and res is not None:
+                for name, tensor in res.items():
+                    if layer_id is not None and self._load_config.vit_separation != 1:
+                        model_weights.set_layer_weight(layer_id, name, tensor)
+                    else:
+                        model_weights.set_global_weight(name, tensor)
+        for layer_id, name, tensor in self._load_uncomplete_weight_modules(device):
+            if layer_id is not None and self._load_config.vit_separation != 1:
+                model_weights.set_layer_weight(layer_id, name, tensor)
+            else:
+                model_weights.set_global_weight(name, tensor)
+        return model_weights
+    
+    def _load_uncomplete_weight_modules(self, device: str):
+        if self._load_config.vit_separation != 1 and not self._is_attn_model:
+            for layer_id in range(self._load_config.num_layers):
+                layer_weights = self._model_weights_info.layer_weights[layer_id]
+                for weight in layer_weights:
+                    if weight.loaded:
+                        continue
+                    results = weight.load(
+                        self._load_config.database, layer_id, device, self._load_config
+                    )
+                    for name, tensor in results.items():
+                        yield (layer_id, name, tensor)
+        for weight in self._model_weights_info.weights:
+            if self._maybe_skip_weight(weight) or weight.loaded:
+                continue
+            results = weight.load(
+                self._load_config.database, None, device, self._load_config
+            )
+            for name, tensor in results.items():
+                yield (None, name, tensor)
+
     def prepare_weights(self, device: str):
         if self._load_config.vit_separation != 1 and not self._is_attn_model:
             for id in range(self._load_config.num_layers):
@@ -225,6 +309,34 @@ def prepare_weights(self, device: str):
             )
             for name, tensor in weights.items():
                 yield (None, name, tensor)
+    
+    def _get_tensor_to_weight_map(
+        self,
+    ) -> Dict[str, tuple[Optional[int], WeightModule]]:
+        tensor_to_weight_map: Dict[str, tuple[Optional[int], WeightModule]] = {}
+        if self._load_config.vit_separation != 1 and not self._is_attn_model:
+            for layer_id in range(self._load_config.num_layers):
+                layer_weights = self._model_weights_info.layer_weights[layer_id]
+                if isinstance(layer_weights, WeightModule):
+                    names = layer_weights.get_tensor_names(layer_id, self._load_config)
+                    tensor_to_weight_map.update(
+                        {k: (layer_id, layer_weights) for k in names}
+                    )
+                else:
+                    for weight in layer_weights:
+                        names = weight.get_tensor_names(layer_id, self._load_config)
+                        tensor_to_weight_map.update(
+                            {k: (layer_id, weight) for k in names}
+                        )
+        for weight in self._model_weights_info.weights:
+            if self._maybe_skip_weight(weight):
+                continue
+            names = weight.get_tensor_names(None, self._load_config)
+            tensor_to_weight_map.update({k: (None, weight) for k in names})
+        for weight in self._misc_weights_info:
+            names = weight.get_tensor_names(None, self._load_config)
+            tensor_to_weight_map.update({k: (None, weight) for k in names})
+        return tensor_to_weight_map
 
     def _maybe_skip_weight(self, weight: WeightModule):
         if self._task_type == TaskType.LANGUAGE_MODEL:
@@ -270,6 +382,7 @@ def _load_from_scratch(self, device: str):
                 weights.set_layer_weight(layer_id, name, tensor)
             else:
                 weights.set_global_weight(name, tensor)
+            gc.collect()
         return weights
 
     def _load_layer_weights(self, layer_id: int, device: str):

rtp_llm/model_loader/model_weight_info.py

@@ -601,7 +601,6 @@ def __init__(self, num_layers: int, device: str, dtype: torch.dtype):
 
     def set_layer_weight(self, layer_id: int, name: str, tensor: torch.Tensor):
         self.weights[layer_id][name] = tensor
-        gc.collect()
 
     def set_global_weight(self, name: str, tensor: torch.Tensor):
         self.global_weights[name] = tensor

rtp_llm/model_loader/per_block_fp8_quant_weight.py

@@ -774,3 +774,42 @@ def _load_raw_tensor(
             res.update({self.scale.name: scale.contiguous().to(device)})
 
         return res
+    
+    def get_tensor_names(
+        self, layer_id: Optional[int], load_config: LoadConfig
+    ) -> set[str]:
+        return self.kernel.get_tensor_names(layer_id, load_config)
+
+    def _process_raw_tensors(
+        self,
+        raw_tensors: Dict[str, torch.Tensor],
+        layer_id: Optional[int],
+        device: str,
+        load_config: LoadConfig,
+    ):
+        kernel_raw_tensors: Dict[str, torch.Tensor] = {
+            key: raw_tensors[key]
+            for key in self.kernel.get_tensor_names(layer_id, load_config)
+        }
+        kernel = self.kernel._process_raw_tensors(
+            kernel_raw_tensors, layer_id, device, load_config
+        )
+        res = {}
+        scale = None
+        if self.scale:
+            quant_kernel, scale = per_block_cast_to_fp8(
+                kernel.get(self.kernel.name), self.group_size
+            )
+            if quant_kernel.dim() == 2:
+                scale = scale.reshape([scale.shape[0], -1])
+        else:
+            quant_kernel = cast_to_fp8(kernel.get(self.kernel.name))
+        if self.kernel.name == W.moe_w1 or self.kernel.name == W.moe_w2:
+            pass
+        elif quant_kernel.dim() == 2:
+            quant_kernel = quant_kernel.T
+        res = {self.kernel.name: quant_kernel.contiguous().to(device)}
+        if self.scale:
+            scale = scale.T if scale.dim() == 2 else scale
+            res.update({self.scale.name: scale.contiguous().to(device)})
+        return res

rtp_llm/model_loader/static_fp8_quant_weight.py

@@ -706,6 +706,59 @@ def _load_raw_tensor(
             )
             res.update(act_scale_inv)
         return res
+    
+    def get_tensor_names(
+        self, layer_id: Optional[int], load_config: LoadConfig
+    ) -> set[str]:
+        names = self.kernel.get_tensor_names(layer_id, load_config)
+        if self.act_scale:
+            names = names.union(self.act_scale.get_tensor_names(layer_id, load_config))
+        if self.act_scale_inv:
+            names = names.union(
+                self.act_scale_inv.get_tensor_names(layer_id, load_config)
+            )
+        return names
+
+    def _process_raw_tensors(
+        self,
+        raw_tensors: Dict[str, torch.Tensor],
+        layer_id: Optional[int],
+        device: str,
+        load_config: LoadConfig,
+    ):
+        kernel_raw_tensors: Dict[str, torch.Tensor] = {
+            key: raw_tensors[key]
+            for key in self.kernel.get_tensor_names(layer_id, load_config)
+        }
+        kernel = self.kernel._process_raw_tensors(
+            kernel_raw_tensors, layer_id, device, load_config
+        )
+        res = {}
+        quant_kernel, scale = quantize_weight_to_fp8(kernel.get(self.kernel.name))
+        quant_kernel = quant_kernel.T
+        res = {
+            self.kernel.name: quant_kernel.contiguous().to(device),
+            self.scale.name: scale.contiguous().to(device),
+        }
+        if self.act_scale:
+            act_scale_raw_tensors: Dict[str, torch.Tensor] = {
+                key: raw_tensors[key]
+                for key in self.act_scale.get_tensor_names(layer_id, load_config)
+            }
+            act_scale = self.act_scale._process_raw_tensors(
+                act_scale_raw_tensors, layer_id, device, load_config
+            )
+            res.update(act_scale)
+        if self.act_scale_inv:
+            act_scale_inv_raw_tensors: Dict[str, torch.Tensor] = {
+                key: raw_tensors[key]
+                for key in self.act_scale_inv.get_tensor_names(layer_id, load_config)
+            }
+            act_scale_inv = self.act_scale_inv._process_raw_tensors(
+                act_scale_inv_raw_tensors, layer_id, device, load_config
+            )
+            res.update(act_scale_inv)
+        return res
 
 
 class Fp8PerTensorCompressedWeight(CompositeWeight, QuantWeight):