| 505 | # FSDP has issues with `inference_mode` |
| 506 | @torch.no_grad() |
| 507 | def validate( |
| 508 | fabric: lightning.Fabric, |
| 509 | model: GPT, |
| 510 | val_dataloader: DataLoader, |
| 511 | tokenizer: Tokenizer, |
| 512 | eval: EvalArgs, |
| 513 | data: List[DataModule], |
| 514 | ) -> torch.Tensor: |
| 515 | fabric.print('Validating ...') |
| 516 | model.eval() |
| 517 | losses = torch.zeros(min(len(val_dataloader), eval.max_iters)) |
| 518 | for k, batch in enumerate(val_dataloader): |
| 519 | if k >= eval.max_iters: |
| 520 | break |
| 521 | input_ids, targets = batch['input_ids'], batch['labels'] |
| 522 | logits = model(input_ids) |
| 523 | losses[k] = chunked_cross_entropy( |
| 524 | logits[..., :-1, :], targets[..., 1:], chunk_size=0 |
| 525 | ) |
| 526 | |
| 527 | val_loss = losses.mean() |
| 528 | |
| 529 | # produce an example: |
| 530 | instruction = ( |
| 531 | 'Recommend a movie for me to watch during the weekend and explain the reason.' |
| 532 | ) |
| 533 | fabric.print(instruction) |
| 534 | |
| 535 | # Simplistically take the first datamodule's prompstyle. They shouldn't be different... |
| 536 | prompt = data[0].prompt_style.apply(instruction) |
| 537 | encoded = tokenizer.encode(prompt, device=fabric.device) |
| 538 | |
| 539 | # TODO: Our GPT block does not have the newer method of #set_kv_cache. |
| 540 | # with fabric.init_tensor(): |
| 541 | # # do not set `max_seq_length=max_returned_token` because memory is not a concern here |
| 542 | # model.set_kv_cache(batch_size=1) |
| 543 | |
| 544 | output = generate( |
| 545 | model, |
| 546 | encoded, |
| 547 | max_returned_tokens=len(encoded) + eval.max_new_tokens, |
| 548 | temperature=0.8, |
| 549 | eos_id=tokenizer.eos_id, |
| 550 | ) |
| 551 | |
| 552 | # TODO: Our GPT block does not have the new method of #clear_kv_cache |
| 553 | # model.clear_kv_cache() |
| 554 | |
| 555 | output = tokenizer.decode(output) |
| 556 | fabric.print(output) |
| 557 | |
| 558 | model.train() |
| 559 | return val_loss |
| 560 | |
| 561 | |
| 562 | def get_lr_scheduler(optimizer, warmup_steps: int, max_steps: int): |