(key,
tp_dim=-1,
no_prefix=0,
is_expert_weights=False,
tp_size=None,
tp_rank=None)
| 1558 | torch_dtype = str_dtype_to_torch(dtype) |
| 1559 | |
| 1560 | def load(key, |
| 1561 | tp_dim=-1, |
| 1562 | no_prefix=0, |
| 1563 | is_expert_weights=False, |
| 1564 | tp_size=None, |
| 1565 | tp_rank=None): |
| 1566 | if not no_prefix: |
| 1567 | key = f'{model_prefix}.' + key |
| 1568 | ptr_idx = safetensors_map[key] if key in safetensors_map else 0 |
| 1569 | |
| 1570 | if key not in safetensors_ptrs[ptr_idx].keys(): |
| 1571 | return None |
| 1572 | |
| 1573 | tensor_slice = safetensors_ptrs[ptr_idx].get_slice(key) |
| 1574 | tensor_shape = tensor_slice.get_shape() |
| 1575 | if tp_dim == -1: |
| 1576 | res = tensor_slice[:] |
| 1577 | elif tp_dim >= 0 and tp_dim < len(tensor_shape): |
| 1578 | if is_expert_weights: |
| 1579 | if tp_size is None: |
| 1580 | tp_size = mapping.moe_tp_size |
| 1581 | if tp_rank is None: |
| 1582 | tp_rank = mapping.moe_tp_rank |
| 1583 | else: |
| 1584 | if tp_size is None: |
| 1585 | tp_size = mapping.tp_size |
| 1586 | if tp_rank is None: |
| 1587 | tp_rank = mapping.tp_rank |
| 1588 | dim_size = tensor_shape[tp_dim] |
| 1589 | if dim_size % tp_size != 0: |
| 1590 | logger.error( |
| 1591 | f"Current weight {key}'s shape {tensor_shape} is invalid at dimension {tp_dim} for TP size {tp_size}" |
| 1592 | ) |
| 1593 | indices = [slice(None)] * len(tensor_shape) |
| 1594 | indices[tp_dim] = slice(dim_size * tp_rank // tp_size, |
| 1595 | dim_size * (tp_rank + 1) // tp_size) |
| 1596 | res = tensor_slice[indices] |
| 1597 | else: |
| 1598 | raise ValueError( |
| 1599 | f"Invalid TP dim {tp_dim} for weight {key}'s shape {tensor_shape}" |
| 1600 | ) |
| 1601 | return res.to(torch_dtype).contiguous( |
| 1602 | ) if "block_sparse_moe.gate" not in key and "block_sparse_moe.router" not in key else res.to( |
| 1603 | torch.float32) |
| 1604 | |
| 1605 | def load_and_set(target, |
| 1606 | key, |
no test coverage detected