Repository navigation
Conversation
| vec = self._vectorize_block(block_weight).float() | ||
| centroid = centroids[block_idx] | ||
| chunk_rows = max(1, _CDIST_CHUNK_BYTES // (centroid.shape[0] * 4)) | ||
| clusters = torch.cat( |
There was a problem hiding this comment.
Maybe nit: We now have a lot of this pattern of "split some operation across chunks, then re-concatenate the results". Can we create a helper for this that takes in the blockwise operation as a callable, along with the axis to split across and the chunk size to use?
There was a problem hiding this comment.
I tried this out and while it does work, I feel like it introduces a level of abstraction that isn't really needed. The chunking operations are already pretty much just a single for loop, inside of which we can take whatever actions we want.
Implementing the helper means extracting each set of separate actions into their own functions or lambda function and defining the generic scaffold helper function which itself contains some logic for slicing, zipping, mapping, and concatting. The resulting changes lead to more lines of code than before and additional functions to be unit tested, and don't look as clear as the original code with single for loop in my opinion.
964d73f to
2607341
Compare
…ices_on_original_device # Conflicts: # src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
Currently palettizer indices are computed by moving centroids and weights to cpu before running torch.cdist. This PR keeps centroids/weights on the original device when computing indices to allow gpu/mps speedup, and moves indices to cpu at the end of the computation.
This does mean that extra gpu memory is needed to temporarily store the indices before they are moved to cpu, though they will be with uint8 datatype.
As the torch.cdist operation itself could potentially lead to much greater gpu memory usage due to computing the distance of each weight between each centroid,
_CDIST_CHUNK_BYTESis used to chunk the computation to cap the usage at 0.5 GB.