Build the memory references. The memory buffers are built using the build_memory_buffer API. This API will allocate a weight buffer pointer to the memory buffer according to the weight_buffer_meta. Args: weight_buffer_meta: memory_buffers: Returns:
(weight_buffer_meta: Dict[str, Dict], memory_buffers: Dict[torch.dtype, MemoryBuffer])
| 114 | |
| 115 | |
| 116 | def build_memory_reference(weight_buffer_meta: Dict[str, Dict], memory_buffers: Dict[torch.dtype, MemoryBuffer]): |
| 117 | """Build the memory references. The memory buffers are built using the build_memory_buffer API. |
| 118 | This API will allocate a weight buffer pointer to the memory buffer according to the weight_buffer_meta. |
| 119 | |
| 120 | Args: |
| 121 | weight_buffer_meta: |
| 122 | memory_buffers: |
| 123 | |
| 124 | Returns: |
| 125 | |
| 126 | """ |
| 127 | start_idx = {} |
| 128 | weight_buffers = {} |
| 129 | for dtype in memory_buffers.keys(): |
| 130 | start_idx[dtype] = 0 |
| 131 | |
| 132 | for name, meta_info in sorted(weight_buffer_meta.items()): |
| 133 | shape = meta_info['shape'] |
| 134 | dtype = meta_info['dtype'] |
| 135 | |
| 136 | buffer = memory_buffers[dtype].get(shape, start_index=start_idx[dtype]) |
| 137 | start_idx[dtype] += calc_padded_numel(shape, dtype) |
| 138 | weight_buffers[name] = buffer |
| 139 | |
| 140 | return weight_buffers |
| 141 | |
| 142 | |
| 143 | class MemoryBufferModuleWrapper: |
no test coverage detected