Build the memory buffer given weight_buffer_meta Args: weight_buffer_meta: contains mapping from name to a dictionary containing shape and dtype of the tensors Returns: a large memory buffer for each dtype that can hold all the tensors
(weight_buffer_meta: Dict[str, Dict])
| 68 | |
| 69 | |
| 70 | def build_memory_buffer(weight_buffer_meta: Dict[str, Dict]) -> Dict[torch.dtype, MemoryBuffer]: |
| 71 | """Build the memory buffer given weight_buffer_meta |
| 72 | |
| 73 | Args: |
| 74 | weight_buffer_meta: contains mapping from name to a dictionary containing shape and dtype of the tensors |
| 75 | |
| 76 | Returns: a large memory buffer for each dtype that can hold all the tensors |
| 77 | |
| 78 | """ |
| 79 | memory_buffers = {} |
| 80 | total_numel_map = {} # map from dtype to the total numel |
| 81 | for name, meta_info in sorted(weight_buffer_meta.items()): |
| 82 | shape = meta_info["shape"] |
| 83 | dtype = meta_info["dtype"] |
| 84 | |
| 85 | assert isinstance(shape, torch.Size) |
| 86 | assert isinstance(dtype, torch.dtype) |
| 87 | |
| 88 | if dtype not in total_numel_map: |
| 89 | total_numel_map[dtype] = 0 |
| 90 | |
| 91 | total_numel_map[dtype] += calc_padded_numel(shape, dtype) |
| 92 | |
| 93 | for dtype, total_numel in total_numel_map.items(): |
| 94 | memory_buffers[dtype] = MemoryBuffer(total_numel, total_numel, dtype) |
| 95 | |
| 96 | return memory_buffers |
| 97 | |
| 98 | |
| 99 | def build_memory_reference_from_module(module: torch.nn.Module, memory_buffers: Dict[torch.dtype, MemoryBuffer], maintain_weight=True): |
no test coverage detected