(model, testenc, dev = torch.device('cuda:0'))
| 5 | |
| 6 | @torch.no_grad() |
| 7 | def llama_eval(model, testenc, dev = torch.device('cuda:0')): |
| 8 | print('Evaluating ...') |
| 9 | model.to(dev) |
| 10 | testenc = testenc.input_ids |
| 11 | nsamples = testenc.numel() // model.seqlen |
| 12 | |
| 13 | use_cache = model.config.use_cache |
| 14 | model.config.use_cache = False |
| 15 | layers = model.model.layers |
| 16 | |
| 17 | model.model.embed_tokens = model.model.embed_tokens.to(dev) |
| 18 | layers[0] = layers[0].to(dev) |
| 19 | |
| 20 | dtype = next(iter(model.parameters())).dtype |
| 21 | inps = torch.zeros((nsamples, model.seqlen, model.config.hidden_size), dtype=dtype, device=dev) |
| 22 | cache = {'i': 0, 'attention_mask': None} |
| 23 | |
| 24 | class Catcher(nn.Module): |
| 25 | |
| 26 | def __init__(self, module): |
| 27 | super().__init__() |
| 28 | self.module = module |
| 29 | |
| 30 | def forward(self, inp, **kwargs): |
| 31 | inps[cache['i']] = inp |
| 32 | cache['i'] += 1 |
| 33 | cache['attention_mask'] = kwargs['attention_mask'] |
| 34 | cache['position_ids'] = kwargs['position_ids'] |
| 35 | raise ValueError |
| 36 | |
| 37 | layers[0] = Catcher(layers[0]) |
| 38 | for i in range(nsamples): |
| 39 | batch = testenc[:, (i * model.seqlen):((i + 1) * model.seqlen)].to(dev) |
| 40 | try: |
| 41 | model(batch) |
| 42 | except ValueError: |
| 43 | pass |
| 44 | layers[0] = layers[0].module |
| 45 | |
| 46 | layers[0] = layers[0].cpu() |
| 47 | model.model.embed_tokens = model.model.embed_tokens.cpu() |
| 48 | torch.cuda.empty_cache() |
| 49 | |
| 50 | outs = torch.zeros_like(inps) |
| 51 | attention_mask = cache['attention_mask'] |
| 52 | position_ids = cache['position_ids'] |
| 53 | |
| 54 | for i in tqdm.tqdm(range(len(layers))): |
| 55 | # print(i, end=" ") |
| 56 | layer = layers[i].cpu() |
| 57 | layer = layer.to(dev) |
| 58 | for j in range(nsamples): |
| 59 | outs[j] = layer(inps[j].unsqueeze(0), attention_mask=attention_mask, position_ids=position_ids)[0] |
| 60 | layers[i] = layer.cpu() |
| 61 | del layer |
| 62 | torch.cuda.empty_cache() |
| 63 | inps, outs = outs, inps |
| 64 | # print(".") |
no test coverage detected