weights.py 14.9 KB
Newer Older
1
2
import torch

3
from abc import ABC, abstractmethod
4
from contextlib import contextmanager
5
from pathlib import Path
6
from typing import Dict, List, Optional, Union, Type
7
from safetensors import safe_open
8
9
from dataclasses import dataclass

10
from text_generation_server.utils.import_utils import SYSTEM
11
12
13
14
15
16
17
18
19
20
21
22
23


class WeightsLoader(ABC):
    """
    Instances of this type implement higher-level weight loading.

    At a low-level, every weight is stored in the Safetensors format.
    The interpretation of weights may be different however, for instance
    could be packed, quantized weights. Loaders are responsible for
    interpreting the raw tensors, sharding tensors in a manner compatible
    with the format, etc.
    """

24
25
26
27
28
29
30
    @abstractmethod
    def get_weights(self, weights: "Weights", prefix: str):
        """
        Get weights at the given prefix and apply without tensor paralllism.
        """
        ...

31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
    @abstractmethod
    def get_weights_col_packed(
        self,
        weights: "Weights",
        prefix: str,
        block_sizes: Union[int, List[int]],
    ):
        """
        Get the packed weights at the given prefix with column-splitting for
        tensor parallelism. This method should be used when multiple different
        weights are packed into a tensor, for instance, query/key/value
        weights or a gate/up projection.

        The `block_sizes` determines the proportions of the packed tensors.
        The columns are split in equally sized blocks when `block_sizes` is an
        `int`, or in blocks proportional given to the sizes. For instance
        `[2, 1, 1]` will divide an input with dimensionality `1024` in
        `[512, 256, 256]`.
        """
        ...

    def get_weights_col(self, weights: "Weights", prefix: str):
        """
        Get weights at the given prefix and apply column-splitting for tensor
        paralllism.
        """
        return weights.get_multi_weights_col([prefix], 0)

    @abstractmethod
    def get_multi_weights_col(self, weights: "Weights", prefixes: List[str], dim: int):
        """
        Get the weights at the given prefixes, column-split them for tensor
        parallelim, and then concatenate the weights along the given dimension.
        """
        ...

    @abstractmethod
    def get_weights_row(self, weights: "Weights", prefix: str):
        """
        Get the weights at the given prefix and apply row-splitting for tensor
        parallism.
        """
        ...


76
77
78
79
80
81
82
83
84
85
86
class Weight(ABC):
    """Instances of this type implement unquantized/quantized/to-be
    quantized weights."""

    @abstractmethod
    def get_linear(self, bias: torch.Tensor):
        """Create a linear layer from this weight."""
        ...


@dataclass
87
class UnquantizedWeight(Weight):
88
89
90
91
92
93
94
95
96
97
98
    weight: torch.Tensor

    def get_linear(self, bias: torch.Tensor):
        from text_generation_server.layers.linear import FastLinear, FastLinearROCm

        if SYSTEM == "rocm":
            return FastLinearROCm(self.weight, bias)
        else:
            return FastLinear(self.weight, bias)


99
class DefaultWeightsLoader(WeightsLoader):
100
101
    """Weight loader that loads (unquantized) Torch tensors."""

102
    def __init__(self, weight_class: Type[UnquantizedWeight]):
103
104
105
106
107
108
        """Create a loader. Weights will be wrapped using the given `weights_class`,
        normally this will be `UnquantizedWeight`, but a quantizer-specific class
        such as `Fp8Weight` can be used to quantize the weights during loading.
        """
        self.weight_class = weight_class

109
110
111
112
113
    """
    Loader that uses tensors as-is with the exception of applying sharding
    and/or concatenation.
    """

114
115
116
    def get_weights(self, weights: "Weights", prefix: str):
        return weights.get_tensor(f"{prefix}.weight")

117
118
119
120
121
122
    def get_weights_col_packed(
        self,
        weights: "Weights",
        prefix: str,
        block_sizes: Union[int, List[int]],
    ):
123
124
125
126
        return self.weight_class(
            weights.get_packed_sharded(
                f"{prefix}.weight", dim=0, block_sizes=block_sizes
            ),
127
128
129
130
        )

    def get_multi_weights_col(self, weights: "Weights", prefixes: List[str], dim: int):
        w = [weights.get_sharded(f"{p}.weight", dim=0) for p in prefixes]
131
        return self.weight_class(torch.cat(w, dim=dim))
132
133

    def get_weights_row(self, weights: "Weights", prefix: str):
134
135
136
        return self.weight_class(
            weights.get_sharded(f"{prefix}.weight", dim=1),
        )
137
138
139


class Weights:
140
141
142
143
144
145
    def __init__(
        self,
        filenames: List[Path],
        device,
        dtype,
        process_group,
146
        weights_loader: WeightsLoader,
147
        aliases: Optional[Dict[str, List[str]]] = None,
OlivierDehaene's avatar
OlivierDehaene committed
148
        prefix: Optional[str] = None,
149
    ):
150
151
152
153
154
155
156
157
158
        routing = {}
        for filename in filenames:
            with safe_open(filename, framework="pytorch") as f:
                for k in f.keys():
                    if k in routing:
                        raise RuntimeError(
                            f"Key {k} was found in multiple files: {filename} and {routing[k]}"
                        )
                    routing[k] = filename
159
160
161
        if aliases is None:
            aliases = {}
        self.aliases = aliases
162
163
164
165
        self.routing = routing
        self.device = device
        self.dtype = dtype
        self.process_group = process_group
Nicolas Patry's avatar
Nicolas Patry committed
166
        self.prefix = prefix
167
        self.weights_loader = weights_loader
168
169
170
171
172
173
174
175
176
        self._handles = {}

    def _get_handle(self, filename):
        if filename not in self._handles:
            f = safe_open(filename, framework="pytorch")
            self._handles[filename] = f

        return self._handles[filename]

177
    def get_filename(self, tensor_name: str) -> (str, str):
Nicolas Patry's avatar
Nicolas Patry committed
178
179
180
181
182
183
184
185
186
187
        names = [tensor_name]
        if self.prefix is not None:
            prefixed = f"{self.prefix}.{tensor_name}"
            names.append(prefixed)
        for name in names:
            filename = self.routing.get(name, None)
            if filename is not None:
                return str(filename), name

            aliases = self.aliases.get(name, [])
188
189
190
191
            for alias in aliases:
                filename = self.routing.get(alias, None)
                if filename is not None:
                    return str(filename), alias
Nicolas Patry's avatar
Nicolas Patry committed
192
        raise RuntimeError(f"weight {tensor_name} does not exist")
193
194

    def _get_slice(self, tensor_name: str):
195
        filename, tensor_name = self.get_filename(tensor_name)
196
197
198
199
        f = self._get_handle(filename)
        slice_ = f.get_slice(tensor_name)
        return slice_

200
    def has_tensor(self, tensor_name: str):
201
202
203
204
205
206
        try:
            self.get_filename(tensor_name)
        except Exception:
            return False
        return True

207
208
209
    def get_shape(self, tensor_name: str):
        return self._get_slice(tensor_name).get_shape()

210
211
212
    def get_tensor(
        self, tensor_name: str, to_device: bool = True, to_dtype: bool = True
    ) -> torch.Tensor:
213
        filename, tensor_name = self.get_filename(tensor_name)
214
215
        f = self._get_handle(filename)
        tensor = f.get_tensor(tensor_name)
216
        # Special case for gptq which shouldn't convert
217
        # u4 which are disguised as int32. Exl2 uses int16
218
219
220
221
222
223
224
225
226
227
228
        # as well. FP8 uses torch.float8_e4m3fn
        if (
            tensor.dtype
            not in [
                torch.float8_e4m3fn,
                torch.int16,
                torch.int32,
                torch.int64,
            ]
            and to_dtype
        ):
229
            tensor = tensor.to(dtype=self.dtype)
xiaobin's avatar
xiaobin committed
230
231
        if to_device:
            tensor = tensor.to(device=self.device)
232
233
        return tensor

234
235
236
    def get_partial_sharded(
        self, tensor_name: str, dim: int, to_device=True, to_dtype=True
    ):
237
        filename, tensor_name = self.get_filename(tensor_name)
xiaobin's avatar
xiaobin committed
238
239
        f = self._get_handle(filename)
        slice_ = f.get_slice(tensor_name)
240
241
242
243
        world_size = self.process_group.size()
        rank = self.process_group.rank()

        size = slice_.get_shape()[dim]
244
        block_size = (size + world_size - 1) // world_size
245
246
247
248
249
250
251
252
253
        start = rank * block_size
        stop = (rank + 1) * block_size

        if dim == 0:
            tensor = slice_[start:stop]
        elif dim == 1:
            tensor = slice_[:, start:stop]
        else:
            raise NotImplementedError("Let's make that generic when needed")
254
        # Special case for gptq which shouldn't convert
255
        # u4 which are disguised as int32. exl2 uses int16.
256
257
258
259
260
        # FP8 uses torch.float8_e4m3fn.
        if (
            tensor.dtype not in (torch.float8_e4m3fn, torch.int16, torch.int32)
            and to_dtype
        ):
261
            tensor = tensor.to(dtype=self.dtype)
262
263
        if to_device:
            tensor = tensor.to(device=self.device)
264
        return tensor
265

266
    def get_sharded(self, tensor_name: str, dim: int, to_device=True, to_dtype=True):
267
268
269
270
271
272
273
274
        filename, tensor_name = self.get_filename(tensor_name)
        f = self._get_handle(filename)
        slice_ = f.get_slice(tensor_name)
        world_size = self.process_group.size()
        size = slice_.get_shape()[dim]
        assert (
            size % world_size == 0
        ), f"The choosen size {size} is not compatible with sharding on {world_size} shards"
275
276
277
        return self.get_partial_sharded(
            tensor_name, dim, to_device=to_device, to_dtype=to_dtype
        )
278

279
    def get_packed_sharded(
280
281
282
283
284
        self,
        tensor_name: str,
        dim: int,
        block_sizes: Union[int, List[int]],
        to_dtype=True,
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
    ) -> torch.Tensor:
        """
        Get a shard from a tensor that packs multiple tensors.

        When a tensor packs multiple tensors (such as QKV or an up
        projection + gate projection), sharding with `get_sharded` is not
        safe since it would not split the packed tensors across shards.

        This method shards a tensor, such that the packed tensors are
        split across shards.

        The columns are split in equally sized blocks when blocks is an `int`, or
        in blocks proportional given to the sizes. For instance `[2, 1, 1]` will
        divide an input with dimensionality `1024` in `[512, 256, 256]`. This is
        convenient for e.g. splitting QKV without knowing the storage details of
        quantized weights.
        """
        slice_ = self._get_slice(tensor_name)
        total_size = slice_.get_shape()[dim]
304
305
        block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=block_sizes)

xiaobin's avatar
xiaobin committed
306
307
308
        world_size = self.process_group.size()
        rank = self.process_group.rank()

309
        tensors = []
310
311
312
313
        block_offset = 0
        for block_size in block_sizes:
            assert (
                block_size % world_size == 0
314
            ), f"Prepacked tensor cannot be sharded across {world_size} shards"
315
316
317
            shard_block_size = block_size // world_size
            start = rank * shard_block_size
            stop = (rank + 1) * shard_block_size
318
319
320
321
322
323
324
            if dim == 0:
                tensor = slice_[block_offset + start : block_offset + stop]
            elif dim == 1:
                tensor = slice_[:, block_offset + start : block_offset + stop]
            else:
                raise NotImplementedError("Currently only dim=0 or dim=1 is supported")
            tensors.append(tensor)
325
            block_offset += block_size
326
327
        tensor = torch.cat(tensors, dim=dim)
        tensor = tensor.to(device=self.device)
328

329
        # Avoid casting quantizer dtypes.
330
331
332
333
334
335
336
337
338
339
        if (
            tensor.dtype
            not in [
                torch.float8_e4m3fn,
                torch.int16,
                torch.int32,
                torch.int64,
            ]
            and to_dtype
        ):
340
341
342
            tensor = tensor.to(dtype=self.dtype)

        return tensor
xiaobin's avatar
xiaobin committed
343

344
345
346
    def get_weights(self, prefix: str):
        return self.weights_loader.get_weights(self, prefix)

347
348
349
350
351
352
353
    def get_weights_col_packed_qkv(
        self,
        prefix: str,
        num_heads: int,
        num_key_value_heads: int,
    ):
        return self.get_weights_col_packed(
354
            prefix, [num_heads, num_key_value_heads, num_key_value_heads]
355
        )
Nicolas Patry's avatar
Nicolas Patry committed
356

357
358
    def get_weights_col_packed_gate_up(self, prefix: str):
        return self.get_weights_col_packed(prefix, 2)
Nicolas Patry's avatar
Nicolas Patry committed
359

360
    def get_weights_col_packed(self, prefix: str, block_sizes: Union[int, List[int]]):
xiaobin's avatar
xiaobin committed
361
        """
362
363
364
365
366
        The columns are split in equally sized blocks when blocks is an `int`, or
        in blocks proportional given to the sizes. For instance `[2, 1, 1]` will
        divide an input with dimensionality `1024` in `[512, 256, 256]`. This is
        convenient for e.g. splitting QKV without knowing the storage details of
        quantized weights.
xiaobin's avatar
xiaobin committed
367
        """
368
        return self.weights_loader.get_weights_col_packed(self, prefix, block_sizes)
369

370
371
    def get_weights_col(self, prefix: str):
        return self.weights_loader.get_weights_col(self, prefix)
372

373
374
    def get_multi_weights_col(self, prefixes: List[str], dim: int):
        return self.weights_loader.get_multi_weights_col(self, prefixes, dim)
OlivierDehaene's avatar
OlivierDehaene committed
375

xiaobin's avatar
xiaobin committed
376
377
378
379
380
381
382
383
384
385
386
387
388
389
    def get_tensor_shard(self, var, dim):
        world_size = self.process_group.size()
        rank = self.process_group.rank()
        block_size = var.size()[dim] // world_size
        start = rank * block_size
        stop = (rank + 1) * block_size
        if dim == 0:
            tensor = var[start:stop]
        elif dim == 1:
            tensor = var[:, start:stop]
        else:
            raise NotImplementedError("Let's make that generic when needed")
        tensor = tensor.to(dtype=self.dtype)
        tensor = tensor.to(device=self.device)
OlivierDehaene's avatar
OlivierDehaene committed
390
        return tensor
391

392
393
    def get_weights_row(self, prefix: str):
        return self.weights_loader.get_weights_row(self, prefix)
394

395
396
397
398
399
400
401
402
403
404
405
406
407
408
    @contextmanager
    def use_loader(self, weights_loader: WeightsLoader):
        """
        This method is a context manager that can be used to use `Weights` with
        a different loader for the duration of the context.
        """

        old_loader = self.weights_loader
        self.weights_loader = weights_loader
        try:
            yield
        finally:
            self.weights_loader = old_loader

409
410
411
412
    @property
    def loader(self):
        return self.weights_loader

413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439

def _blocks_to_block_sizes(total_size: int, blocks: Union[int, List[int]]) -> List[int]:
    """
    Convert block count or proportions to block sizes.

    This function accepts

    - The number of blocks (int), in which case the block size is
      total_size//blocks; or
    - A list of block sizes (List[int]).

    In the latter case, if sum(blocks) < total_size, the ratios between
    the block sizes will be preserved. For instance, if blocks is
    [2, 1, 1] and total_size is 1024, the returned block sizes are
    [512, 256, 256].
    """
    if isinstance(blocks, list):
        total_blocks = sum(blocks)
        assert (
            total_size % total_blocks == 0
        ), f"Cannot split {total_size} in proportional blocks: {blocks}"
        part_size = total_size // total_blocks
        return [part_size * block for block in blocks]
    else:
        assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}"
        single_size = total_size // blocks
        return [single_size] * blocks