| 43 | |
| 44 | @torch.no_grad() |
| 45 | def llama_eval(model, testenc, dev, seqlen = 2048): |
| 46 | from tqdm import tqdm |
| 47 | print('Evaluating ...') |
| 48 | |
| 49 | testenc = testenc.input_ids |
| 50 | nsamples = testenc.numel() // seqlen |
| 51 | |
| 52 | use_cache = model.config.use_cache |
| 53 | model.config.use_cache = False |
| 54 | layers = model.model.layers |
| 55 | |
| 56 | model.model.embed_tokens = model.model.embed_tokens.to(dev) |
| 57 | layers[0] = layers[0].to(dev) |
| 58 | |
| 59 | dtype = next(iter(model.parameters())).dtype |
| 60 | inps = torch.zeros((nsamples, seqlen, model.config.hidden_size), dtype=dtype, device=dev) |
| 61 | cache = {'i': 0, 'attention_mask': None} |
| 62 | |
| 63 | class Catcher(nn.Module): |
| 64 | |
| 65 | def __init__(self, module): |
| 66 | super().__init__() |
| 67 | self.module = module |
| 68 | |
| 69 | def forward(self, inp, **kwargs): |
| 70 | inps[cache['i']] = inp |
| 71 | cache['i'] += 1 |
| 72 | cache['attention_mask'] = kwargs['attention_mask'] |
| 73 | cache['position_ids'] = kwargs['position_ids'] |
| 74 | raise ValueError |
| 75 | |
| 76 | layers[0] = Catcher(layers[0]) |
| 77 | for i in range(nsamples): |
| 78 | batch = testenc[:, (i * seqlen):((i + 1) * seqlen)].to(dev) |
| 79 | try: |
| 80 | model(batch) |
| 81 | except ValueError: |
| 82 | pass |
| 83 | layers[0] = layers[0].module |
| 84 | |
| 85 | layers[0] = layers[0].cpu() |
| 86 | model.model.embed_tokens = model.model.embed_tokens.cpu() |
| 87 | torch.cuda.empty_cache() |
| 88 | |
| 89 | outs = torch.zeros_like(inps) |
| 90 | attention_mask = cache['attention_mask'] |
| 91 | position_ids = cache['position_ids'] |
| 92 | |
| 93 | for i in tqdm(range(len(layers))): |
| 94 | # print('layer', i) |
| 95 | layer = layers[i].to(dev) |
| 96 | layer = layers[i].to(dtype) |
| 97 | for j in range(nsamples): |
| 98 | # print("dtype", inps[j].unsqueeze(0).dtype) |
| 99 | outs[j] = layer(inps[j].unsqueeze(0), attention_mask=attention_mask, position_ids=position_ids)[0] |
| 100 | layers[i] = layer.cpu() |
| 101 | del layer |
| 102 | torch.cuda.empty_cache() |