(self)
| 158 | ) |
| 159 | |
| 160 | def preload(self): |
| 161 | # Initialize shards and load_func |
| 162 | if os.path.isdir(self.model_dir): |
| 163 | shard_files = glob.glob(self.model_dir + "/*." + self.format.value) |
| 164 | elif os.path.isfile(self.model_dir): |
| 165 | shard_files = [self.model_dir] |
| 166 | elif isinstance(self.model_dir, dict): |
| 167 | shard_files = [self.model_dir] |
| 168 | elif isinstance(self.model_dir, PreTrainedModel): |
| 169 | shard_files = [dict(self.model_dir.named_parameters())] |
| 170 | else: |
| 171 | raise NotImplementedError( |
| 172 | "args.model_dir is not a directory, a file or an in-memory module!" |
| 173 | ) |
| 174 | shard_files.sort() |
| 175 | if self.format == ModelWeightsFormat.SAFETENSORS: |
| 176 | self.shards = [ |
| 177 | safe_open(f, framework="pt", device="cpu") for f in shard_files |
| 178 | ] |
| 179 | elif self.format == ModelWeightsFormat.BINARY or self.format == ModelWeightsFormat.PYTORCH: |
| 180 | self.shards = [ |
| 181 | torch.load(f, weights_only=True, map_location="cpu", mmap=True) |
| 182 | for f in shard_files |
| 183 | ] |
| 184 | elif self.format == ModelWeightsFormat.IN_MEMORY: |
| 185 | self.shards = [shard_files[0]] |
| 186 | else: |
| 187 | raise NotImplementedError( |
| 188 | "Only *.safetensors/*.pth/*.bin files are supported.") |
| 189 | for idx, shard in enumerate(self.shards): |
| 190 | self.shard_map.update({k: idx for k in shard.keys()}) |
| 191 | |
| 192 | def load_tensor(self, key, tp_size=1, tp_dim=-1, tp_rank=0): |
| 193 | # Retrieve shard index |
no test coverage detected