#

テストスイートで llm.int8 () 量子化を使用して GPT-Neox を評価する

このコードは、一連のタスクで llm.int8 () 量子化を使用して GPT-Neox を評価します。

14import torch
15from torch import nn
16
17from labml import monit
18from labml_nn.neox.evaluation import run_eval_harness
19from labml_nn.neox.model import LayerGenerator

#

22def main():

#

端末

24    device = torch.device('cuda:0')

#

float16 のレイヤーを CPU にロードします。レイヤーをGPUにロードした後にその場でこれを行うと、CUDAメモリの断片化が発生するため、後でレイヤーをint8に変換します（フラグメンテーションにより約3GBのメモリが失われる可能性があります

）。

29    layer_generator = LayerGenerator(is_clone_layers=True,
30                                     dtype=torch.float16,
31                                     device=torch.device('cpu'),
32                                     )

#

レイヤーをロード

34    layers = list(layer_generator.load())

#

これにより、CUDA メモリの断片化が減少します。

37    for layer in monit.iterate('Convert to int8', layers, is_children_silent=True):
38        layer_generator.post_load_prepare(layer,
39                                          device=device,
40                                          is_llm_int8=True,
41                                          llm_int8_threshold=6.0,
42                                          )
43        layer.to(device)

#

nn.Sequential モデル作成

46    model = nn.Sequential(*layers)

#

評価用ハーネスを実行

49    print(run_eval_harness(model, 'half_precision', [], device))

#

53if __name__ == '__main__':
54    main()