This code evaluate GPT-NeoX using, on a suite of tasks.
12import argparse
13
14import torch
15from torch import nn
16
17from labml_nn.neox.evaluation import run_eval_harness
18from labml_nn.neox.model import LayerGenerator
21def main():
Argument parser
23 parser = argparse.ArgumentParser()
24
25 parser.add_argument("--flash", action='store_true', help="whether to use Flash Attention")
26
27 opt = parser.parse_args()
Device
30 device = torch.device('cuda:0')
Load layers
32 layers = list(LayerGenerator(is_clone_layers=True,
33 filter_layers=None,
34 dtype=torch.float16,
35 device=device,
36 is_flash_attention=opt.flash,
37 ).load())
Create nn.Sequential
model
40 model = nn.Sequential(*layers)
43 print(run_eval_harness(model, 'half_precision', ['lambada'], device))
47if __name__ == '__main__':
48 main()