deftdawg/bitsandbytes-rocm
1
0
Fork 0
forked from mrq/bitsandbytes-rocm
Code Pull Requests Activity

Fixed cpu blockwise quantization for small input tensors.

Browse Source
This commit is contained in:
Tim Dettmers 2022-09-13 10:37:53 -07:00
parent d8dbf3a9b5
commit c05dd42ddd
2 changed files with 22 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
csrc/cpu_ops.cpp
View File

@ -30,11 +30,12 @@ void quantize_cpu(float *code, float *A, float *absmax, unsigned char *out, long
// between 16k and 64k on Linux (we reach this when running BLOOM-176B with a large batch size) // between 16k and 64k on Linux (we reach this when running BLOOM-176B with a large batch size)
for(long long offset = 0; offset < num_blocks; offset+=thread_wave_size) for(long long offset = 0; offset < num_blocks; offset+=thread_wave_size)
{ {
pthread_t *threads = (pthread_t *) malloc(sizeof(pthread_t) * thread_wave_size); long long valid_chunks = num_blocks - offset >= thread_wave_size ? thread_wave_size : num_blocks - offset;
pthread_t *threads = (pthread_t *) malloc(sizeof(pthread_t) * valid_chunks);
struct quantize_block_args **args = (quantize_block_args **) malloc(thread_wave_size * sizeof(quantize_block_args *)); struct quantize_block_args **args = (quantize_block_args **) malloc(valid_chunks * sizeof(quantize_block_args *));
for(long long i = 0; i < thread_wave_size; i++) for(long long i = 0; i < valid_chunks; i++)
args[i] = (quantize_block_args *) malloc(sizeof(quantize_block_args)); args[i] = (quantize_block_args *) malloc(sizeof(quantize_block_args));
int chunks_processed = 0; int chunks_processed = 0;
@ -56,14 +57,14 @@ void quantize_cpu(float *code, float *A, float *absmax, unsigned char *out, long
pthread_create(&threads[chunks_processed], NULL, &quantize_block, (void *) arg); pthread_create(&threads[chunks_processed], NULL, &quantize_block, (void *) arg);
chunks_processed += 1; chunks_processed += 1;
if(chunks_processed == thread_wave_size){ break; } if(chunks_processed == valid_chunks){ break; }
} }
for (int i = 0; i < thread_wave_size; i++) for (int i = 0; i < valid_chunks; i++)
int err = pthread_join(threads[i], NULL); int err = pthread_join(threads[i], NULL);
free(threads); free(threads);
for (int i = 0; i < thread_wave_size; i++) for (int i = 0; i < valid_chunks; i++)
free(args[i]); free(args[i]);
free(args); free(args);

2
tests/test_functional.py
View File

@ -2133,7 +2133,7 @@ def test_blockwise_cpu_large():
reldiffs = [] reldiffs = []
batch = 128 batch = 128
seq = 128 seq = 128
hidden = 14336 for hidden in [128, 14336]:
for blocksize in [4096, 16384]: for blocksize in [4096, 16384]:
for i in range(2): for i in range(2):
A1 = torch.randn(batch, seq, hidden, device='cpu') A1 = torch.randn(batch, seq, hidden, device='cpu')