Forward step.
(data_iterator, model, args, timers, mems)
| 205 | |
| 206 | |
| 207 | def forward_step(data_iterator, model, args, timers, mems): |
| 208 | """Forward step.""" |
| 209 | |
| 210 | # Get the batch. |
| 211 | timers('batch generator').start() |
| 212 | timers('data loader').start() |
| 213 | rand = random.Random(args.iteration * mpu.get_data_parallel_world_size() + mpu.get_data_parallel_rank()) |
| 214 | if data_iterator[1] and rand.random() < args.multi_task_ratio: |
| 215 | data = next(data_iterator[1]) if data_iterator[1] else None |
| 216 | data["mode"] = "multi-task" |
| 217 | else: |
| 218 | data = next(data_iterator[0]) if data_iterator[0] else None |
| 219 | # print_rank_0("data iterator") |
| 220 | timers('data loader').stop() |
| 221 | tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data, args) |
| 222 | timers('batch generator').stop() |
| 223 | |
| 224 | # print_rank_0("get batch") |
| 225 | |
| 226 | def print_masked_text(batch_id): |
| 227 | block_position_ids = position_ids[:, 1] |
| 228 | position_ids_ = position_ids[:, 0] |
| 229 | sep = attention_mask.item() if torch.numel(attention_mask) == 1 else attention_mask[batch_id].item() |
| 230 | text, last_segment = "", [] |
| 231 | for i, token_id in enumerate(tokens[batch_id, :sep].tolist()): |
| 232 | token = tokenizer.IdToToken(token_id) |
| 233 | if token.startswith('[MASK') or token.endswith('MASK]'): |
| 234 | if last_segment: |
| 235 | text += tokenizer.DecodeIds(last_segment) |
| 236 | last_segment = [] |
| 237 | text += f" [{position_ids_[batch_id, i].item()}, {token}]" |
| 238 | else: |
| 239 | last_segment.append(token_id) |
| 240 | if last_segment: |
| 241 | text += tokenizer.DecodeIds(last_segment) |
| 242 | print(text.encode('utf-8')) |
| 243 | last_index = None |
| 244 | for i in range(sep, tokens.size(1)): |
| 245 | if tokenizer.IdToToken(tokens[batch_id, i].item()).startswith("<|startofpiece"): |
| 246 | if last_index is not None: |
| 247 | print(tokenizer.DecodeIds(tokens[batch_id, last_index: i].tolist()).encode('utf-8'), "|", |
| 248 | tokenizer.DecodeIds(labels[batch_id, last_index: i].tolist()).encode('utf-8'), |
| 249 | position_ids_[batch_id, last_index: i].tolist(), |
| 250 | block_position_ids[batch_id, last_index:i].tolist()) |
| 251 | last_index = i |
| 252 | if last_index is not None: |
| 253 | print(tokenizer.DecodeIds(tokens[batch_id, last_index:].tolist()).encode('utf-8'), "|", |
| 254 | tokenizer.DecodeIds(labels[batch_id, last_index:].tolist()).encode('utf-8'), |
| 255 | position_ids_[batch_id, last_index:].tolist(), block_position_ids[batch_id, last_index:].tolist()) |
| 256 | |
| 257 | if data is not None and "mode" in data: |
| 258 | mode = data['mode'] |
| 259 | else: |
| 260 | mode = 'bert' |
| 261 | |
| 262 | logits, *mems = model(tokens, position_ids, attention_mask, *mems) |
| 263 | losses = mpu.vocab_parallel_cross_entropy(logits.contiguous().float(), |
| 264 | labels) |