(data_iterator, args, timers)
| 38 | |
| 39 | |
| 40 | def get_batch(data_iterator, args, timers): |
| 41 | # Items and their type. |
| 42 | keys = ['text', 'loss_mask'] |
| 43 | datatype = torch.int64 |
| 44 | |
| 45 | # Broadcast data. |
| 46 | timers('data loader').start() |
| 47 | if data_iterator is not None: |
| 48 | data = next(data_iterator) |
| 49 | else: |
| 50 | data = None |
| 51 | timers('data loader').stop() |
| 52 | |
| 53 | data_b = mpu.broadcast_data(keys, data, datatype) |
| 54 | # Unpack. |
| 55 | tokens_ = data_b['text'].long() |
| 56 | loss_mask = data_b['loss_mask'].float() |
| 57 | |
| 58 | labels = tokens_[:, 1:].contiguous() |
| 59 | loss_mask = loss_mask[:, 1:].contiguous() |
| 60 | tokens = tokens_[:, :-1].contiguous() |
| 61 | |
| 62 | attention_mask = None |
| 63 | |
| 64 | # Get the masks and postition ids. |
| 65 | attention_mask, loss_mask, position_ids = get_masks_and_position_ids( |
| 66 | tokens, |
| 67 | loss_mask=loss_mask, |
| 68 | attention_mask=attention_mask, |
| 69 | args=args |
| 70 | ) |
| 71 | # Convert |
| 72 | if args.fp16: |
| 73 | attention_mask = attention_mask.half() |
| 74 | |
| 75 | return tokens, labels, loss_mask, attention_mask, position_ids |
| 76 | |
| 77 | |
| 78 | def forward_step(data_iterator, model, args, timers): |
no test coverage detected