| 125 | |
| 126 | |
| 127 | class SFTMetric: |
| 128 | def __init__(self, device): |
| 129 | self.n_step = 0 |
| 130 | self.right = torch.Tensor([0]).to(device=device) |
| 131 | self.total = torch.Tensor([0]).to(device=device) |
| 132 | self.total_loss = torch.Tensor([0]).to(device=device) |
| 133 | self.world_size = dist.get_world_size() |
| 134 | |
| 135 | def __call__(self, logits, labels, loss): |
| 136 | return self.update(logits, labels, loss) |
| 137 | |
| 138 | def update(self, logits, labels, loss): |
| 139 | self.n_step += 1 |
| 140 | with torch.no_grad(): |
| 141 | shift_preds = logits[..., :-1, :].argmax(dim=-1) |
| 142 | shift_labels = labels[..., 1:] |
| 143 | self.right += (shift_preds == shift_labels).masked_fill(shift_labels.eq(-100), 0).sum().item() |
| 144 | self.total += (shift_labels != -100).sum().item() |
| 145 | self.total_loss += loss.item() |
| 146 | |
| 147 | def get_metric(self, reset=True): |
| 148 | dist.all_reduce(self.right, op=torch.distributed.ReduceOp.SUM) |
| 149 | dist.all_reduce(self.total, op=torch.distributed.ReduceOp.SUM) |
| 150 | dist.all_reduce(self.total_loss, op=torch.distributed.ReduceOp.SUM) |
| 151 | |
| 152 | acc = (self.right / self.total).item() |
| 153 | loss = self.total_loss.item() / (self.world_size * self.n_step) |
| 154 | |
| 155 | if reset: |
| 156 | self.n_step = 0 |
| 157 | self.right.fill_(0) |
| 158 | self.total.fill_(0) |
| 159 | self.total_loss.fill_(0) |
| 160 | return acc, loss |
| 161 | |
| 162 | |
| 163 | def train(args): |