sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

faster gather for bert (#4526)

Browse Source
This commit is contained in:
wozeparrot
2024-05-10 22:28:48 -07:00
committed by GitHub
parent 922e6e056a
commit d2c347fc74
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/models/bert.py
+5 -2
View File
@@ -59,8 +59,11 @@ class BertForMLPerf:
output = self.model(input_ids, attention_mask, segment_ids)
clsf_logits = self.clsf_output(self.clsf_pooling_activation(self.clsf_pooler(output[:, 0]))).cast(dtypes.float32)
masked_positions = masked_positions[:, :, None].expand(-1, -1, output.shape[-1])
h_masked = Tensor.gather(output, masked_positions, 1)
# gather only the masked_positions we care about
counter = Tensor.arange(output.shape[1], requires_grad=False, device=output.device).reshape(1, 1, output.shape[1]).expand(*masked_positions.shape, output.shape[1])
onehot = counter == masked_positions.unsqueeze(2).expand(*masked_positions.shape, output.shape[1])
h_masked = onehot @ output
h_masked = self.lm_norm(self.lm_transform_activation(self.lm_transform(h_masked)))
lm_logits = self.lm_output(h_masked) + self.lm_output_bias