Lower weight compression memory footprint by sorting weights according to their size

[nikita-savelyevv]

Changes

Sort weights for compression:

all_weight_params = sorted(all_weight_params, key=lambda wp: wp.num_weights, reverse=True)

Reason for changes

During weights compression, memory footprint gradually increases when new low-bit constants are created. At the same time there are temporary spikes in memory footprint which happen during compressed weight computation. For example, here:

if invert_scale:
    scale = fns.power(scale, -1)
    compressed_weights = weight * scale
else:
    compressed_weights = weight / scale
if zero_point is not None:
    compressed_weights += zero_point.astype(weight.dtype)
compressed_weights = fns.round(compressed_weights)
compressed_weights = fns.clip(compressed_weights, level_low, level_high).astype(dtype)

Multiple temporary full precision arrays are needed to be created here. After that they get garbage-collected. However, as it was said, this creates temporary spikes in memory footprint. Taking this into account, it makes sense to compress large constants first so that there are not many low-bit constants taking up memory yet. This mostly is affected by embedding matrices.

Please see memory figures below. They were obtain during 8-bit weights compression. Figures were gathered with the memory_logger.py, memory_type=SYSTEM_NORMALIZED.

Backend	Model	Before	After
OV	qwen2-7b
PT	qwen2-7b

For example, for qwen2-7b OV model there is a reduction from ~12GB peak footprint to ~7GB.

Much lower values for OV backend compared to PT are because OV models are read using mmap which allows to avoid allocating memory for the whole full-precision model.

Related tickets

144501

[andreyanufr]

@andreyanufr Could you please confirm that AWQ and Scale Estimation algorithms are agnostic to the weight ordering? @alexsu52 And also GPTQ?

AWQ and Scale Estimation algorithms are agnostic to the weight ordering, but not GPTQ.

[nikita-savelyevv]

Weight compression build with id 103 has passed