| 438 | |
| 439 | |
| 440 | def chunked_cross_entropy( |
| 441 | logits: Union[torch.Tensor, List[torch.Tensor]], targets: torch.Tensor, chunk_size: int = 128 |
| 442 | ) -> torch.Tensor: |
| 443 | # with large max_sequence_lengths, the beginning of `backward` allocates a large memory chunk which can dominate |
| 444 | # the memory usage in fine-tuning settings with low number of parameters. |
| 445 | # as a workaround hack, the cross entropy computation is chunked to force it to deallocate on the go, reducing |
| 446 | # the memory spike's magnitude |
| 447 | |
| 448 | # lm_head was chunked (we are fine-tuning) |
| 449 | if isinstance(logits, list): |
| 450 | # don't want to chunk cross entropy |
| 451 | if chunk_size == 0: |
| 452 | logits = torch.cat(logits, dim=1) |
| 453 | logits = logits.reshape(-1, logits.size(-1)) |
| 454 | targets = targets.reshape(-1) |
| 455 | return torch.nn.functional.cross_entropy(logits, targets, ignore_index=-1) |
| 456 | |
| 457 | # chunk cross entropy |
| 458 | logit_chunks = [logit_chunk.reshape(-1, logit_chunk.size(-1)) for logit_chunk in logits] |
| 459 | target_chunks = [target_chunk.reshape(-1) for target_chunk in targets.split(logits[0].size(1), dim=1)] |
| 460 | loss_chunks = [ |
| 461 | torch.nn.functional.cross_entropy(logit_chunk, target_chunk, ignore_index=-1, reduction="none") |
| 462 | for logit_chunk, target_chunk in zip(logit_chunks, target_chunks) |
| 463 | ] |
| 464 | return torch.cat(loss_chunks).mean() |
| 465 | |
| 466 | # no chunking at all |
| 467 | logits = logits.reshape(-1, logits.size(-1)) |
| 468 | targets = targets.reshape(-1) |
| 469 | if chunk_size == 0: |
| 470 | return torch.nn.functional.cross_entropy(logits, targets, ignore_index=-1) |
| 471 | |
| 472 | # lm_head wasn't chunked, chunk cross entropy |
| 473 | logit_chunks = logits.split(chunk_size) |
| 474 | target_chunks = targets.split(chunk_size) |
| 475 | loss_chunks = [ |
| 476 | torch.nn.functional.cross_entropy(logit_chunk, target_chunk, ignore_index=-1, reduction="none") |
| 477 | for logit_chunk, target_chunk in zip(logit_chunks, target_chunks) |
| 478 | ] |
| 479 | return torch.cat(loss_chunks).mean() |
| 480 | |
| 481 | |
| 482 | def map_old_state_dict_weights(state_dict: Dict, mapping: Mapping, prefix: str) -> Dict: |