Forward step.
(data_iterator, model, args, timers)
| 76 | |
| 77 | |
| 78 | def forward_step(data_iterator, model, args, timers): |
| 79 | """Forward step.""" |
| 80 | |
| 81 | # Get the batch. |
| 82 | timers('batch generator').start() |
| 83 | tokens, labels, loss_mask, attention_mask, position_ids = get_batch( |
| 84 | data_iterator, args, timers) |
| 85 | timers('batch generator').stop() |
| 86 | # Forward model. |
| 87 | logits, *mems = model(tokens, position_ids, attention_mask) |
| 88 | losses = mpu.vocab_parallel_cross_entropy(logits.contiguous().float(), labels) |
| 89 | # scaling loss mask |
| 90 | loss_mask = loss_mask.view(-1) |
| 91 | |
| 92 | losses = losses.view(-1) * loss_mask |
| 93 | loss = torch.sum(losses) / loss_mask.sum() |
| 94 | return loss, {} |
| 95 | |
| 96 | |
| 97 | class FakeDataset(TensorDataset): |
no test coverage detected