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