This class wraps a GPT layer for specific operations.
| 4 | |
| 5 | # Define WrappedGPT class |
| 6 | class WrappedGPT: |
| 7 | """ |
| 8 | This class wraps a GPT layer for specific operations. |
| 9 | """ |
| 10 | |
| 11 | def __init__(self, layer, layer_id=0, layer_name="none"): |
| 12 | self.layer = layer |
| 13 | self.dev = self.layer.weight.device |
| 14 | self.rows = layer.weight.data.shape[0] |
| 15 | self.columns = layer.weight.data.shape[1] |
| 16 | |
| 17 | self.scaler_row = torch.zeros((self.columns), device=self.dev) |
| 18 | # self.activations = [torch.zeros((self.columns), device=self.dev)] |
| 19 | self.activations = [] |
| 20 | self.nsamples = 0 |
| 21 | |
| 22 | self.layer_id = layer_id |
| 23 | self.layer_name = layer_name |
| 24 | |
| 25 | def add_batch(self, inp, out, tar): |
| 26 | """ |
| 27 | tar: batch_size * seq_len, inp corresponding to the position where tar == -100 will be ignored |
| 28 | """ |
| 29 | if len(inp.shape) == 2: |
| 30 | inp = inp.unsqueeze(0) |
| 31 | if len(tar.shape) == 2: |
| 32 | tar = tar.unsqueeze(0) |
| 33 | |
| 34 | tmp = inp.shape[0] # bs |
| 35 | |
| 36 | mask = tar.ne(-100) |
| 37 | if isinstance(self.layer, nn.Linear): |
| 38 | if len(inp.shape) == 3: |
| 39 | inp = inp.reshape((-1, inp.shape[-1])) |
| 40 | mask = mask.flatten() |
| 41 | inp = inp[mask] # remove -100's |
| 42 | inp = inp.t() |
| 43 | |
| 44 | self.scaler_row *= self.nsamples / (self.nsamples + tmp) |
| 45 | self.nsamples += tmp |
| 46 | |
| 47 | inp = inp.type(torch.float32) |
| 48 | self.scaler_row += torch.norm(inp, p=2, dim=1) ** 2 / self.nsamples |
| 49 | self.activations.append(inp) |
no outgoing calls
no test coverage detected