LLM.INT8() ප්රමාණකරණය භාවිතා කරමින් GPT-නියෝක්ස් සමඟ පෙළ ජනනය කරන්න

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

25def generate():

31    cache = get_cache()
32    cache.set('use_cache', True)

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

40    layer_generator = LayerGenerator(is_clone_layers=True,
41                                     dtype=torch.float16,
42                                     device=torch.device('cpu'),
43                                     is_llm_int8=False,
44                                     )
45    layers = list(layer_generator.load())

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

57    model = nn.Sequential(*layers)

60    torch.cuda.empty_cache()
61    print(torch.cuda.memory_summary())

64    ids = get_tokens(PROMPT)

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

73    ids += [next_token]

76    for i in range(1, 100):

78        cache.set('state_ids', (i, i + 1))

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

84        ids += [next_token]

86        print_tokens(ids, [ids])

90if __name__ == '__main__':
91    generate()