r"""A helper function to partially load a tensor. This can save memory sometimes as this allows tensor parallel shards to stream into memory ( without being concatenated into a full model first). Args: target (``torch.Tensor``): The target tensor to load the values into.
(
target: torch.Tensor, parallel_dim: int, num_shards: int, shard_id: int,
value: torch.Tensor, mode: str = "set",
)
| 486 | |
| 487 | |
| 488 | def tensor_load_shard( |
| 489 | target: torch.Tensor, parallel_dim: int, num_shards: int, shard_id: int, |
| 490 | value: torch.Tensor, mode: str = "set", |
| 491 | ) -> None: |
| 492 | r"""A helper function to partially load a tensor. This can save memory |
| 493 | sometimes as this allows tensor parallel shards to stream into memory ( |
| 494 | without being concatenated into a full model first). |
| 495 | |
| 496 | Args: |
| 497 | target (``torch.Tensor``): The target tensor to load the values into. |
| 498 | parallel_dim (int): Tensor parallel dimension of the tensor. |
| 499 | num_shards (int): Number of tensor parallel shards of the value. |
| 500 | shard_id (int): The shard id of the current value. |
| 501 | value (``torch.Tensor``): The value to be loaded into the target |
| 502 | tensor. |
| 503 | mode (str): The supported values are ``set`` and ``add``. If ``set``, |
| 504 | the old value in the target tensor is overrided with the new value. |
| 505 | If ``add``, the new value is added to the old value. |
| 506 | """ |
| 507 | assert parallel_dim < target.ndim or parallel_dim == -1 |
| 508 | target_slices = [] |
| 509 | for i in range(target.ndim): |
| 510 | if i == parallel_dim: |
| 511 | dim_st = target.size(i) // num_shards * shard_id |
| 512 | dim_ed = target.size(i) // num_shards * (shard_id + 1) |
| 513 | target_slices.append(slice(dim_st, dim_ed)) |
| 514 | else: |
| 515 | target_slices.append(slice(None)) |
| 516 | if parallel_dim == -1 and shard_id != 0 and mode in ["set", "add"]: |
| 517 | return |
| 518 | if mode == "set": |
| 519 | target[target_slices] = value |
| 520 | elif mode == "add": |
| 521 | target[target_slices] += value |
| 522 | else: |
| 523 | raise NotImplementedError(f"Unknown mode: {mode}.") |
| 524 | |
| 525 | |
| 526 | class ShardedTensorLoader: |