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