diff --git a/docs/neox/evaluation/llm_int8.html b/docs/neox/evaluation/llm_int8.html new file mode 100644 index 00000000..60926b00 --- /dev/null +++ b/docs/neox/evaluation/llm_int8.html @@ -0,0 +1,177 @@ + + + + + + + + + + + + + + + + + + + + + + + llm_int8.py + + + + + + + + + + +
+
+
+
+

+ home + neox + evaluation +

+

+ + Sponsor + + Github + + Twitter +

+

+ + View code on Github +

+
+
+
+
+
+ # +
+ +
+
+ 
1import torch
+2from torch import nn
+3
+4from labml import monit
+5from labml_nn.neox.evaluation import run_eval_harness
+6from labml_nn.neox.model import LayerGenerator
+7
+8if __name__ == '__main__':
+9    device = torch.device('cuda:0')
+10    layer_generator = LayerGenerator(is_clone_layers=True,
+11                                     dtype=torch.float16,
+12                                     device=torch.device('cpu'),
+13                                     )
+
+
+
+
+
+ # +
+

Load layers

+ +
+
+ 
15    layers = list(layer_generator.load())
+
+
+
+
+
+ # +
+

This reduces CUDA memory fragmentation

+ +
+
+ 
18    for layer in monit.iterate('Convert to int8', layers, is_children_silent=True):
+19        layer_generator.post_load_prepare(layer,
+20                                          device=device,
+21                                          is_llm_int8=True,
+22                                          llm_int8_threshold=6.0,
+23                                          )
+24        layer.to(device)
+25
+26    with monit.section('Sequential'):
+27        model = nn.Sequential(*layers)
+28
+29    print(run_eval_harness(model, 'half_precision', [], device))
+
+
+
+ Trending Research Papers + labml.ai +
+
+ + + + \ No newline at end of file diff --git a/docs/neox/model.html b/docs/neox/model.html index 8f793cb5..7759d28e 100644 --- a/docs/neox/model.html +++ b/docs/neox/model.html @@ -230,7 +230,7 @@ + is the base for , which defaults to
@@ -253,7 +253,7 @@
#
-

To store for the features

+

To store for the features

@@ -265,7 +265,7 @@
#
-

Cache and

+

Cache and

@@ -278,7 +278,7 @@
#
-

Base for

+

Base for

@@ -357,7 +357,7 @@
#
-

Initialize

+

Initialize

@@ -369,7 +369,7 @@
#
-

+

@@ -382,7 +382,7 @@
#
-

Initialize and cache

+

Initialize and cache

@@ -399,7 +399,7 @@
#
-

Get position indexes

+

Get position indexes

@@ -411,7 +411,7 @@
#
-

+

@@ -423,8 +423,8 @@
#
-

Concatenate so that for row we have

-

+

Concatenate so that for row we have

+

@@ -436,7 +436,7 @@
#
-

Calculate and in fp32

+

Calculate and in fp32

@@ -501,7 +501,7 @@ #

RoPE embeddings

-

for

+

for

@@ -1566,7 +1566,10 @@ is the data type of the model
  • device is the device of the model
    • -

    Returns the layers as a generator

    +
  • is_llm_int8 + specifies whether to use int8 quantization
    • +
  • llm_int8_threshold + is the threshold used to separate outlier features
    • @@ -1575,7 +1578,10 @@ 459 filter_layers: Optional[Set] = None, 460 is_clone_layers: bool = True, 461 dtype: torch.dtype = torch.float, -462 device: torch.device = torch.device('cpu')):
    +462 device: torch.device = torch.device('cpu'), +463 is_llm_int8: bool = False, +464 llm_int8_threshold: float = 6.0, +465 ):
    @@ -1586,34 +1592,39 @@
    - 
    482        if filter_layers is None:
-483            filter_layers = set(range(n_layers + 3))
-484
-485        self.n_vocab = n_vocab
-486        self.n_hidden = n_hidden
-487        self.n_layers = n_layers
-488        self.n_heads = n_heads
-489        self.filter_layers = filter_layers
-490        self.is_clone_layers = is_clone_layers
-491        self.dtype = dtype
-492        self.device = device
-493
-494        self.pre_created_layers = dict(
-495            transformer_layer=None,
-496        )
    + 
    486        if filter_layers is None:
+487            filter_layers = set(range(n_layers + 3))
+488
+489        self.n_vocab = n_vocab
+490        self.n_hidden = n_hidden
+491        self.n_layers = n_layers
+492        self.n_heads = n_heads
+493        self.filter_layers = filter_layers
+494        self.is_clone_layers = is_clone_layers
+495        self.dtype = dtype
+496        self.device = device
+497        self.is_llm_int8 = is_llm_int8
+498        self.llm_int8_threshold = llm_int8_threshold
+499
+500        self.pre_created_layers = dict(
+501            transformer_layer=None,
+502        )
    -
    +
    #
    - +

    Prepares the layer for usage

    +

    We move the layer to the device and convert it to the correct data type

    +
    • layer + is the layer to prepare
      • +

      Returns the prepared layer

    +
    - 
    498    def _prepare_layer(self, layer: NeoXModule):
-499        layer = layer.to(self.device, self.dtype)
-500        return layer
    + 
    504    def _prepare_layer(self, layer: NeoXModule):
    @@ -1624,33 +1635,35 @@
    - 
    502    def _create_and_cache_layer(self, name: str, creator: Callable[[], NeoXModule]):
-503        if self.pre_created_layers[name] is None or not self.is_clone_layers:
-504            layer = creator()
-505        else:
-506            layer = copy.deepcopy(self.pre_created_layers[name])
-507
-508        layer: NeoXModule = self._prepare_layer(layer)
-509
-510        if self.pre_created_layers[name] is None:
-511            self.pre_created_layers[name] = layer
-512
-513        return layer
    + 
    513        return layer.to(self.device, self.dtype)
    -
    +
    #
    - +

    ### Layer transformations after loading the checkpoint

    +

    This function implements layer transformations after loading the checkpoint.

    +

    Currently, it only applies the int8 quantization.

    +
    • layer + is the layer to prepare
      • +
    • is_llm_int8 + specifies whether to use int8 quantization
      • +
    • device + is the device of the model
      • +
    • llm_int8_threshold + is the threshold used to separate outlier features
      • +

      Returns the prepared layer

    +
    - 
    515    def _create_transformer_layer(self):
-516        return self._create_and_cache_layer(
-517            'transformer_layer',
-518            lambda: TransformerLayer(self.n_hidden, self.n_heads)
-519        )
    + 
    515    @torch.no_grad()
+516    def post_load_prepare(self, layer: NeoXModule, *,
+517                          is_llm_int8: bool = None,
+518                          device: torch.device = None,
+519                          llm_int8_threshold: float = None,
+520                          ):
    @@ -1658,11 +1671,16 @@
    #
    - +

    Get default values if not specified

    +
    - 
    521    def _create_embedding_layer(self):
-522        return Embedding(self.n_vocab, self.n_hidden)
    + 
    537        if is_llm_int8 is None:
+538            is_llm_int8 = self.is_llm_int8
+539        if device is None:
+540            device = self.device
+541        if llm_int8_threshold is None:
+542            llm_int8_threshold = self.llm_int8_threshold
    @@ -1670,11 +1688,12 @@
    #
    - +

    Skip if not using int8 quantization

    +
    - 
    524    def _create_final_norm_layer(self):
-525        return FinalNorm(self.n_hidden)
    + 
    545        if not is_llm_int8:
+546            return layer
    @@ -1682,11 +1701,12 @@
    #
    - +

    Only convert the linear layers in the transformer layers

    +
    - 
    527    def _create_readout_layer(self):
-528        return ReadoutLayer(self.n_hidden, self.n_vocab)
    + 
    549        if not isinstance(layer, TransformerLayer):
+550            return layer
    @@ -1694,11 +1714,12 @@
    #
    - +

    Use make_llm_int8_linear + defined in utilities.

    +
    - 
    530    @torch.no_grad()
-531    def get_layers(self) -> Generator[Tuple[NeoXModule, Tuple[str, str]], None, None]:
    + 
    553        from labml_nn.neox.utils.llm_int8 import make_llm_int8_linear
    @@ -1706,14 +1727,23 @@
    #
    -

    Embedding layer

    +

    Convert the linear layers

    - 
    533        if 0 in self.filter_layers:
-534            with monit.section('Embedding layer'):
-535                layer = self._prepare_layer(self._create_embedding_layer())
-536            yield layer, ('layer_00-model_00-model_states.pt', 'layer_00-model_01-model_states.pt')
    + 
    556        with monit.section('Convert to int8'):
+557            layer.attention.output = make_llm_int8_linear(layer.attention.output,
+558                                                          device=device,
+559                                                          threshold=llm_int8_threshold)
+560            layer.attention.qkv_lin = make_llm_int8_linear(layer.attention.qkv_lin,
+561                                                           device=device,
+562                                                           threshold=llm_int8_threshold)
+563            layer.ffn.dense_h_h4 = make_llm_int8_linear(layer.ffn.dense_h_h4,
+564                                                        device=device,
+565                                                        threshold=llm_int8_threshold)
+566            layer.ffn.dense_h4_h = make_llm_int8_linear(layer.ffn.dense_h4_h,
+567                                                        device=device,
+568                                                        threshold=llm_int8_threshold)
    @@ -1721,27 +1751,29 @@
    #
    -

    Transformer layers

    +

    - 
    539        for i in range(self.n_layers):
    + 
    570        return layer
    -
    +
    #
    -

    Transformer layer

    +

    Creates and caches a layer

    +

    Copying cached layers is faster than initializing new layers because it takes time to initialize parameters.

    +
    • name + is the name of the layer
      • +
    • creator + is the function to create the layer
      • +

      Returns the created layer or a copy of the cached layer

    - 
    541            if i + 1 in self.filter_layers:
-542                with monit.section(f'Transformer Layer {i}'):
-543                    yield self._create_transformer_layer(), \
-544                          (f'layer_{i + 2 :02d}-model_00-model_states.pt',
-545                           f'layer_{i + 2 :02d}-model_01-model_states.pt')
    + 
    572    def _create_and_cache_layer(self, name: str, creator: Callable[[], NeoXModule]):
    @@ -1749,14 +1781,17 @@
    #
    -

    Final normalization layer

    - +
    - 
    548        if self.n_layers + 1 in self.filter_layers:
-549            with monit.section('Final norm layer'):
-550                layer = self._prepare_layer(self._create_final_norm_layer())
-551            yield layer, ('layer_47-model_00-model_states.pt', 'layer_47-model_01-model_states.pt')
    + 
    584        if not self.is_clone_layers:
+585            return self._prepare_layer(creator())
+586
+587        if self.pre_created_layers[name] is None:
+588            self.pre_created_layers[name] = self._prepare_layer(creator())
+589
+590        layer = copy.deepcopy(self.pre_created_layers[name])
+591        return layer
    @@ -1764,14 +1799,14 @@
    #
    -

    Readout layer

    - +
    - 
    554        if self.n_layers + 2 in self.filter_layers:
-555            with monit.section('Readout layer'):
-556                layer = self._prepare_layer(self._create_readout_layer())
-557            yield layer, ('layer_48-model_00-model_states.pt', 'layer_48-model_01-model_states.pt')
    + 
    593    def _create_transformer_layer(self):
+594        return self._create_and_cache_layer(
+595            'transformer_layer',
+596            lambda: TransformerLayer(self.n_hidden, self.n_heads)
+597        )
    @@ -1782,20 +1817,177 @@
    - 
    559    @property
-560    def total_layers(self):
-561        return self.n_layers + 3
-562
-563    @torch.no_grad()
-564    def load(self) -> Generator[NeoXModule, None, None]:
-565        with torch.no_grad():
-566            with monit.section("Layers"):
-567                for i, (layer, files) in enumerate(self.get_layers()):
-568                    if files is not None:
-569                        layer.load_state(*checkpoint.load_checkpoint_files(files))
-570
-571                    monit.progress(min(0.99, (i + 1) / self.total_layers))
-572                    yield layer
    + 
    599    def _create_embedding_layer(self):
+600        return Embedding(self.n_vocab, self.n_hidden)
    +
    + +
    +
    +
    + # +
    + +
    +
    + 
    602    def _create_final_norm_layer(self):
+603        return FinalNorm(self.n_hidden)
    +
    +
    +
    +
    +
    + # +
    + +
    +
    + 
    605    def _create_readout_layer(self):
+606        return ReadoutLayer(self.n_hidden, self.n_vocab)
    +
    +
    +
    +
    +
    + # +
    +

    Generator to get layers

    + +
    +
    + 
    608    @torch.no_grad()
+609    def get_layers(self) -> Generator[Tuple[NeoXModule, Tuple[str, str]], None, None]:
    +
    +
    +
    +
    +
    + # +
    +

    Embedding layer

    + +
    +
    + 
    614        if 0 in self.filter_layers:
+615            with monit.section('Embedding layer'):
+616                layer = self._prepare_layer(self._create_embedding_layer())
+617            yield layer, ('layer_00-model_00-model_states.pt', 'layer_00-model_01-model_states.pt')
    +
    +
    +
    +
    +
    + # +
    +

    Transformer layers

    + +
    +
    + 
    620        for i in range(self.n_layers):
    +
    +
    +
    +
    +
    + # +
    +

    Transformer layer

    + +
    +
    + 
    622            if i + 1 in self.filter_layers:
+623                with monit.section(f'Transformer Layer {i}'):
+624                    yield self._create_transformer_layer(), \
+625                          (f'layer_{i + 2 :02d}-model_00-model_states.pt',
+626                           f'layer_{i + 2 :02d}-model_01-model_states.pt')
    +
    +
    +
    +
    +
    + # +
    +

    Final normalization layer

    + +
    +
    + 
    629        if self.n_layers + 1 in self.filter_layers:
+630            with monit.section('Final norm layer'):
+631                layer = self._prepare_layer(self._create_final_norm_layer())
+632            yield layer, ('layer_47-model_00-model_states.pt', 'layer_47-model_01-model_states.pt')
    +
    +
    +
    +
    +
    + # +
    +

    Readout layer

    + +
    +
    + 
    635        if self.n_layers + 2 in self.filter_layers:
+636            with monit.section('Readout layer'):
+637                layer = self._prepare_layer(self._create_readout_layer())
+638            yield layer, ('layer_48-model_00-model_states.pt', 'layer_48-model_01-model_states.pt')
+639
+640        for k in self.pre_created_layers.keys():
+641            self.pre_created_layers[k] = None
    +
    +
    +
    +
    +
    + # +
    +

    Returns the total number of layers

    + +
    +
    + 
    643    @property
+644    def total_layers(self):
    +
    +
    +
    +
    +
    + # +
    + +
    +
    + 
    648        return self.n_layers + 3
    +
    +
    +
    +
    +
    + # +
    +

    Generator to load layers

    + +
    +
    + 
    650    @torch.no_grad()
+651    def load(self) -> Generator[NeoXModule, None, None]:
    +
    +
    +
    +
    +
    + # +
    + +
    +
    + 
    655        with monit.section("Layers"):
+656            for i, (layer, files) in enumerate(self.get_layers()):
+657                if files is not None:
+658                    layer.load_state(*checkpoint.load_checkpoint_files(files))
+659
+660                layer = self.post_load_prepare(layer)
+661
+662                monit.progress(min(0.99, (i + 1) / self.total_layers))
+663                yield layer
    diff --git a/docs/neox/readme.html b/docs/neox/readme.html new file mode 100644 index 00000000..149419b7 --- /dev/null +++ b/docs/neox/readme.html @@ -0,0 +1,129 @@ + + + + + + + + + + + + + + + + + + + + + + + GPT-NeoX + + + + + + + + + + +
    +
    +
    +
    +

    + home + neox +

    +

    + + Sponsor + + Github + + Twitter +

    +

    + + View code on Github +

    +
    +
    +
    +
    +
    + # +
    + MarkdownException + Small: parse error +
    +
    + +
    +
    +
    + Trending Research Papers + labml.ai +
    +
    + + + + \ No newline at end of file diff --git a/docs/neox/samples/llm_int8.html b/docs/neox/samples/llm_int8.html new file mode 100644 index 00000000..f365cefb --- /dev/null +++ b/docs/neox/samples/llm_int8.html @@ -0,0 +1,362 @@ + + + + + + + + + + + + + + + + + + + + + + + Generate Text with GPT-NeoX using LLM.int8() quantization + + + + + + + + + + +
    +
    +
    +
    +

    + home + neox + samples +

    +

    + + Sponsor + + Github + + Twitter +

    +

    + + View code on Github +

    +
    +
    +
    +
    +
    + # +
    +

    Generate Text with GPT-NeoX using LLM.int8() quantization

    +

    This shows how to generate text from GPT-NeoX using LLM.int8() quantization.

    +

    This needs a GPU with more than 45GB memory.

    + +
    +
    + 
    15from typing import List
+16
+17import torch
+18from torch import nn
+19
+20from labml import monit
+21from labml_nn.neox.model import LayerGenerator
+22from labml_nn.neox.samples.generate import PROMPT, infer
+23from labml_nn.neox.utils import get_tokens, print_tokens
+24from labml_nn.neox.utils.cache import get_cache
    +
    +
    +
    +
    +
    + # +
    +

    Generate text

    + +
    +
    + 
    27def generate():
    +
    +
    +
    +
    +
    + # +
    +

    Setup cache to cache intermediate key/value pairs for faster generation

    + +
    +
    + 
    33    cache = get_cache()
+34    cache.set('use_cache', True)
    +
    +
    +
    +
    +
    + # +
    +

    Device

    + +
    +
    + 
    37    device = torch.device('cuda:0')
    +
    +
    +
    +
    +
    + # +
    +

    Load layers in float16 into CPU. We convert the layers to int8 later, because doing that on the fly after loading layers to GPU causes CUDA memory fragmentation (about 3GB memory can get lost due to fragmentation).

    + +
    +
    + 
    42    layer_generator = LayerGenerator(is_clone_layers=True,
+43                                     dtype=torch.float16,
+44                                     device=torch.device('cpu'),
    +
    +
    +
    +
    +
    + # +
    +

    is_llm_int8=True,

    + +
    +
    + 
    46                                     )
+47    layers = list(layer_generator.load())
    +
    +
    +
    +
    +
    + # +
    +

    This reduces CUDA memory fragmentation

    + +
    +
    + 
    50    for layer in monit.iterate('Convert to int8', layers, is_children_silent=True):
+51        layer_generator.post_load_prepare(layer,
+52                                          device=device,
+53                                          is_llm_int8=True,
+54                                          llm_int8_threshold=6.0,
+55                                          )
+56        layer.to(device)
    +
    +
    +
    +
    +
    + # +
    +

    Create nn.Sequential + model

    + +
    +
    + 
    59    model = nn.Sequential(*layers)
    +
    +
    +
    +
    +
    + # +
    +

    Clear cache and print memory summary for debugging

    + +
    +
    + 
    62    torch.cuda.empty_cache()
+63    print(torch.cuda.memory_summary())
    +
    +
    +
    +
    +
    + # +
    +

    Get token ids

    + +
    +
    + 
    66    ids = get_tokens(PROMPT)
    +
    +
    +
    +
    +
    + # +
    +

    Run the model

    + +
    +
    + 
    69    cache.set('state_ids', (None, 1))
+70    with monit.section('Infer'):
+71        next_token = infer(model, ids, device)[-1]
    +
    +
    +
    +
    +
    + # +
    +

    Append the predicted token

    + +
    +
    + 
    74    ids += [next_token]
    +
    +
    +
    +
    +
    + # +
    +

    Predict 100 tokens

    + +
    +
    + 
    77    for i in range(1, 100):
    +
    +
    +
    +
    +
    + # +
    +

    Set the state to use cached activations

    + +
    +
    + 
    79        cache.set('state_ids', (i, i + 1))
    +
    +
    +
    +
    +
    + # +
    +

    Get next token. Note that we only feed the last token to the model because we cache the key/value pairs of previous tokens.

    + +
    +
    + 
    82        with monit.section('Infer'):
+83            next_token = infer(model, [next_token], device)[-1]
    +
    +
    +
    +
    +
    + # +
    +

    Append the predicted token

    + +
    +
    + 
    85        ids += [next_token]
    +
    +
    +
    +
    +
    + # +
    +

    Print

    + +
    +
    + 
    87        print_tokens(ids, [ids])
    +
    +
    +
    +
    +
    + # +
    +

    + +
    +
    + 
    91if __name__ == '__main__':
+92    generate()
    +
    +
    +
    + Trending Research Papers + labml.ai +
    +
    + + + + \ No newline at end of file diff --git a/docs/neox/utils/llm_int8.html b/docs/neox/utils/llm_int8.html new file mode 100644 index 00000000..7084da7d --- /dev/null +++ b/docs/neox/utils/llm_int8.html @@ -0,0 +1,247 @@ + + + + + + + + + + + + + + + + + + + + + + + LLM.int8() on GPT-NeoX + + + + + + + + + + +
    +
    +
    +
    +

    + home + neox + utils +

    +

    + + Sponsor + + Github + + Twitter +

    +

    + + View code on Github +

    +
    +
    +
    +
    +
    + # +
    +

    LLM.int() on GPT-NeoX

    +

    This implements a utility function to transform a nn.Linear + layer to LLM.int8() linear layer.

    +

    LLM.int8() paper shows you can use int8 quantization while handling outliers to reduce memory footprint without performance degradation in large language models. They convert weights and inputs to scaled 8-bit integers and does matrix multiplication producing int32 results which is then converted back to float16 and rescaled. They show that in large langauge models, some features can give extreme values (outliers) that dominate the model's output. These features get clamped in 8-bit integer space which causes the model performance to degrade. As a solution they pick these outliers (greater than a specified threshold) and compute their multiplications separately in float16 space. Since the percentage of outliers is around 0.01% this doesn't increase memory usage, and prevents the model from degrading performance.

    +

    The code to transform GPT-NoeX layers is defined in model.py.

    +

    Here are example uses of GPT-NeoX with int8 quantization.

    + + +
    +
    + 
    33
    +
    +
    +
    +
    +
    + # +
    +

    Import bitsandbytes + package

    + +
    +
    + 
    34try:
+35    from bitsandbytes.nn import Linear8bitLt, Int8Params
+36except ImportError:
+37    raise ImportError('''Please install `bitsandbytes` with `pip install bitsandbytes -U`''')
+38
+39import torch
+40from torch import nn
    +
    +
    +
    +
    +
    + # +
    +

    Transform a nn.Linear + layer to LLM.int8() linear layer

    +
    • linear_module + is the nn.Linear + layer to transform
      • +
    • device + is the device of the model
      • +
    • threshold + is the threshold to use for outlier detection
    + +
    +
    + 
    43def make_llm_int8_linear(linear_module: nn.Linear, device: torch.device, threshold: float = 6.0):
    +
    +
    +
    +
    +
    + # +
    +

    + +
    +
    + 
    53    assert isinstance(linear_module, nn.Linear)
    +
    +
    +
    +
    +
    + # +
    +

    Create an empty Linear8bitLt module

    + +
    +
    + 
    56    int8_lin = Linear8bitLt(
+57        linear_module.in_features,
+58        linear_module.out_features,
+59        linear_module.bias is not None,
+60        has_fp16_weights=False,
+61        threshold=threshold,
+62    )
    +
    +
    +
    +
    +
    + # +
    +

    Quantize the weights

    + +
    +
    + 
    65    int8_lin._parameters['weight'] = Int8Params(linear_module.weight.data.cpu(),
+66                                                requires_grad=False,
+67                                                has_fp16_weights=False).to(device)
    +
    +
    +
    +
    +
    + # +
    +

    Set the bias in float16 space

    + +
    +
    + 
    70    if linear_module.bias is not None:
+71        int8_lin._parameters['bias'] = nn.Parameter(linear_module.bias.data,
+72                                                    requires_grad=False)
    +
    +
    +
    +
    +
    + # +
    +

    + +
    +
    + 
    75    return int8_lin
    +
    +
    +
    + Trending Research Papers + labml.ai +
    +
    + + + + \ No newline at end of file diff --git a/docs/sitemap.xml b/docs/sitemap.xml index c257ce02..8fb4e95e 100644 --- a/docs/sitemap.xml +++ b/docs/sitemap.xml @@ -134,98 +134,119 @@ https://nn.labml.ai/neox/checkpoint.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/index.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 + 1.00 + + + + + https://nn.labml.ai/neox/utils/llm_int8.html + 2022-08-19T16:30:00+00:00 1.00 https://nn.labml.ai/neox/utils/cache.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/utils/index.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/utils/text_dataset.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/utils/trainer.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/utils/finetune.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/model.html - 2022-08-10T16:30:00+00:00 + 2022-08-19T16:30:00+00:00 + 1.00 + + + + + https://nn.labml.ai/neox/samples/llm_int8.html + 2022-08-19T16:30:00+00:00 1.00 https://nn.labml.ai/neox/samples/generate.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/samples/index.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/samples/finetune.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/tokenizer.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 + 1.00 + + + + + https://nn.labml.ai/neox/evaluation/llm_int8.html + 2022-08-19T16:30:00+00:00 1.00 https://nn.labml.ai/neox/evaluation/half_precision.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/neox/evaluation/index.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 @@ -463,7 +484,7 @@ https://nn.labml.ai/index.html - 2022-08-08T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 @@ -575,7 +596,7 @@ https://nn.labml.ai/optimizers/adam_fp16.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 @@ -659,21 +680,21 @@ https://nn.labml.ai/scaling/index.html - 2022-08-09T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/scaling/zero3/index.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 https://nn.labml.ai/scaling/zero3/finetune_neox.html - 2022-08-10T16:30:00+00:00 + 2022-08-11T16:30:00+00:00 1.00 diff --git a/labml_nn/neox/model.py b/labml_nn/neox/model.py index 829132f7..376af19d 100644 --- a/labml_nn/neox/model.py +++ b/labml_nn/neox/model.py @@ -518,7 +518,22 @@ class LayerGenerator: device: torch.device = None, llm_int8_threshold: float = None, ): - # If we are using int8 quantization, we need to convert the layer to int8 + """ + + ### Layer transformations after loading the checkpoint + + This function implements layer transformations after loading the checkpoint. + + Currently, it only applies the int8 quantization. + + :param layer: is the layer to prepare + :param is_llm_int8: specifies whether to use int8 quantization + :param device: is the device of the model + :param llm_int8_threshold: is the threshold $\alpha$ used to separate outlier features + :return: the prepared layer + """ + + # Get default values if not specified if is_llm_int8 is None: is_llm_int8 = self.is_llm_int8 if device is None: @@ -526,6 +541,7 @@ class LayerGenerator: if llm_int8_threshold is None: llm_int8_threshold = self.llm_int8_threshold + # Skip if not using int8 quantization if not is_llm_int8: return layer @@ -536,7 +552,7 @@ class LayerGenerator: # Use `make_llm_int8_linear` defined in [utilities](./utils/llm_int8.html). from labml_nn.neox.utils.llm_int8 import make_llm_int8_linear - # + # Convert the linear layers with monit.section('Convert to int8'): layer.attention.output = make_llm_int8_linear(layer.attention.output, device=device, diff --git a/labml_nn/neox/samples/llm_int8.py b/labml_nn/neox/samples/llm_int8.py index ff48b537..cea6a58e 100644 --- a/labml_nn/neox/samples/llm_int8.py +++ b/labml_nn/neox/samples/llm_int8.py @@ -1,3 +1,17 @@ +""" +--- +title: Generate Text with GPT-NeoX using LLM.int8() quantization +summary: > + Generate Text with GPT-NeoX using LLM.int8() quantization +--- + +# Generate Text with GPT-NeoX using LLM.int8() quantization + +This shows how to generate text from GPT-NeoX using [LLM.int8() quantization](../utils/llm_int8.html). + +This needs a GPU with more than 45GB memory. +""" + from typing import List import torch @@ -5,31 +19,10 @@ from torch import nn from labml import monit from labml_nn.neox.model import LayerGenerator +from labml_nn.neox.samples.generate import PROMPT, infer from labml_nn.neox.utils import get_tokens, print_tokens from labml_nn.neox.utils.cache import get_cache -# Prompt to complete -PROMPT = 'Einstein was born in the German Empire, but moved to Switzerland in 1895, forsaking his German' - - -def infer(model: nn.Module, ids: List[int], device: torch.device): - """ - ### Predict the next token - - :param layers: is the list of layers - :param ids: are the input token ids - :param device: is the device of the model - """ - - with torch.no_grad(): - # Get the tokens - x = torch.tensor(ids)[None, :].to(device) - # Eval model - x = model(x) - - # Return predicted token - return x[0].max(dim=-1)[1].tolist() - def generate(): """ @@ -43,12 +36,14 @@ def generate(): # Device device = torch.device('cuda:0') + # Load layers in float16 into CPU. We convert the layers to int8 later, because doing that + # on the fly after loading layers to GPU causes CUDA memory fragmentation + # (about 3GB memory can get lost due to fragmentation). layer_generator = LayerGenerator(is_clone_layers=True, dtype=torch.float16, device=torch.device('cpu'), # is_llm_int8=True, ) - # Load layers layers = list(layer_generator.load()) # This reduces CUDA memory fragmentation @@ -60,10 +55,11 @@ def generate(): ) layer.to(device) + # Create `nn.Sequential` model model = nn.Sequential(*layers) + # Clear cache and print memory summary for debugging torch.cuda.empty_cache() - print(torch.cuda.memory_summary()) # Get token ids diff --git a/labml_nn/neox/utils/llm_int8.py b/labml_nn/neox/utils/llm_int8.py index 1f684e28..58048760 100644 --- a/labml_nn/neox/utils/llm_int8.py +++ b/labml_nn/neox/utils/llm_int8.py @@ -1,8 +1,36 @@ """ -* [Generate](../samples/llm_int8.html) -* [Evaluation](../evaluation/llm_int8.html) +--- +title: LLM.int8() on GPT-NeoX +summary: > + Transform nn.Linear layers to 8-bit integer layers. +--- + +# LLM.int() on GPT-NeoX + +This implements a utility function to transform a `nn.Linear` layer to LLM.int8() linear layer. + +[LLM.int8() paper](https://papers.labml.ai/paper/eb2bcaee1d0011edaa66a71c10a887e7) + shows you can use int8 quantization while handling outliers to +reduce memory footprint without performance degradation in large language models. +They convert weights and inputs to scaled 8-bit integers and does matrix multiplication +producing int32 results which is then converted back to float16 and rescaled. +They show that in large langauge models, some features can give extreme values (outliers) +that dominate the model's output. +These features get clamped in 8-bit integer space which causes the model performance to degrade. +As a solution they pick these outliers (greater than a specified threshold) +and compute their multiplications separately in float16 space. +Since the percentage of outliers is around 0.01% this doesn't increase memory usage, +and prevents the model from degrading performance. + +The code to transform GPT-NoeX layers is defined in [model.py](../model.html#post_load_prepare). + +Here are example uses of GPT-NeoX with int8 quantization. + +* [Generate Text](../samples/llm_int8.html) +* [Run Evaluation Tests](../evaluation/llm_int8.html) """ +# Import [`bitsandbytes`](https://github.com/timdettmers/bitsandbytes) package try: from bitsandbytes.nn import Linear8bitLt, Int8Params except ImportError: @@ -13,7 +41,18 @@ from torch import nn def make_llm_int8_linear(linear_module: nn.Linear, device: torch.device, threshold: float = 6.0): - # Create a Linear8bitLt module + """ + ## Transform a `nn.Linear` layer to LLM.int8() linear layer + + :param linear_module: is the `nn.Linear` layer to transform + :param device: is the device of the model + :param threshold: is the threshold $\alpha$ to use for outlier detection + """ + + # + assert isinstance(linear_module, nn.Linear) + + # Create an empty Linear8bitLt module int8_lin = Linear8bitLt( linear_module.in_features, linear_module.out_features, @@ -22,15 +61,15 @@ def make_llm_int8_linear(linear_module: nn.Linear, device: torch.device, thresho threshold=threshold, ) - # Set the weights + # Quantize the weights int8_lin._parameters['weight'] = Int8Params(linear_module.weight.data.cpu(), requires_grad=False, has_fp16_weights=False).to(device) - # Set the bias. - # We don't have to convert this to Int8 since it doesn't use a lot of memory. + # Set the bias in float16 space if linear_module.bias is not None: int8_lin._parameters['bias'] = nn.Parameter(linear_module.bias.data, requires_grad=False) + # return int8_lin