Corrupted outputs with Marlin int4 kernels as parallelization increases

[dacorvo]

When using MarlinInt4WeightQBitsTensor and its associated optimized gemm kernel, there are issues with the weight/scales/zero-point readback as soon as parallelization increases.

The consequence is that output features higher than 128 are corrupted when a sufficient amount of inputs are parallelized.

Test to reproduce the issue here:

optimum-quanto/test/tensor/weights/optimized/test_marlin_int4_weight_qbits_tensor.py

Line 134 in 852bb9c

@pytest.mark.xfail(reason="Bug in Marlin kernel", strict=False)

[inarikami]

Could be related, but I noticed the latest release of optimum-quanto (v0.2.5) corrupts transformer weights during qfloat8 quantization. Downgrading to 0.2.4 solved this issue. Not sure what the exact cause is but will look into it

Code that caused corruption in 0.2.5 but not earlier versions:

pipe = FluxPipeline.from_pretrained(...

quantize(pipe.transformer, weights=qfloat8)
freeze(pipe.transformer)

quantize(pipe.text_encoder, weights=qfloat8)
freeze(pipe.text_encoder)

quantize(pipe.text_encoder_2, weights=qfloat8)
freeze(pipe.text_encoder_2)

[Leommm-byte]

Could be related, but I noticed the latest release of optimum-quanto (v0.2.5) corrupts transformer weights during qfloat8 quantization. Downgrading to 0.2.4 solved this issue. Not sure what the exact cause is but will look into it

Code that caused corruption in 0.2.5 but not earlier versions:

pipe = FluxPipeline.from_pretrained(...

quantize(pipe.transformer, weights=qfloat8)
freeze(pipe.transformer)

quantize(pipe.text_encoder, weights=qfloat8)
freeze(pipe.text_encoder)

quantize(pipe.text_encoder_2, weights=qfloat8)
freeze(pipe.text_encoder_2)

Yeah, same here. I was confused at first because the generated image was just pure noise so I downgraded to this version

https://github.com/huggingface/optimum-quanto.git@65ace79d6af6ccc27afbb3576541cc36b3e3a98b

and it worked fine. (This was the 0.25.0.dev0)

[dacorvo]

@inarikami @Leommm-byte this cannot be related, as the new Marlin kernel is only available for qint4 and is not used by default.
If you have a script allowing to reproduce the corruption on 0.0.25, feel free to open an issue.

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[ahadnagy]

I just quickly ran this through Compute Sanitizer, and it seems like there's a race condition in the Marlin kernel:

compute-sanitizer --tool=racecheck --target-processes=all pytest [REDUCTED]//repos/optimum-quanto/test/tensor/weights/optimized/test_marlin_int4_weight_qbits_tensor.py::test_marlin_int4_weight_qbits_tensor_linear_failing
========= COMPUTE-SANITIZER
=================================================================================================================================================== test session starts ===================================================================================================================================================
platform linux -- Python 3.12.7, pytest-7.4.4, pluggy-1.5.0
rootdir: [REDUCTED]//repos/optimum-quanto
collected 16 items                                                                                                                                                                                                                                                                                                        

test/tensor/weights/optimized/test_marlin_int4_weight_qbits_tensor.py ========= Error: Race reported between Read access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x4eb0
=========     and Write access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x51c0 [59520 hazards]
========= 
========= Error: Race reported between Read access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x4e80
=========     and Write access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x51b0 [58752 hazards]
========= 
========= Error: Race reported between Read access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x4e40
=========     and Write access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x51a0 [57984 hazards]
========= 
========= Error: Race reported between Read access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x4120
=========     and Write access at void Marlin<(int)256, (int)3, (int)16, (int)4, (int)4, (int)8>(const int4 *, const int4 *, int4 *, const int4 *, const int4 *, int, int, int, int *)+0x4630 [30336 hazards]
...
...

I'll dig deeper when I can find a spare few hours. :)

[dacorvo]

@ahadnagy thank you for your feedback: ping me if you need any help.

[ahadnagy]

So far I've found two possible race conditions:

• Between fetch_to_shared() and fetch_to_registers():
  

  optimum-quanto/optimum/quanto/library/extensions/cuda/marlin/marlin_cuda_kernel.cu

  Line 395 in 66dca87

  auto fetch_to_shared = [&] (int pipe, int a_off, bool pred = true) {

  
  

  optimum-quanto/optimum/quanto/library/extensions/cuda/marlin/marlin_cuda_kernel.cu

  Line 451 in 66dca87

  ldsm4(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);

• Between fetch_to_registers() and thread_block_reduce()
  

  optimum-quanto/optimum/quanto/library/extensions/cuda/marlin/marlin_cuda_kernel.cu

  Line 445 in 66dca87

  reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];

  
  

  optimum-quanto/optimum/quanto/library/extensions/cuda/marlin/marlin_cuda_kernel.cu

  Line 451 in 66dca87

  ldsm4(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);

The second one appears at relatively small sizes as well:

@pytest.mark.parametrize("tokens", [16])
@pytest.mark.parametrize("in_features", [512])
@pytest.mark.parametrize("out_features", [256])

@dacorvo Could these be related to the memory issues you suspected? Not sure if these are false positives or not (never seen a false positive with this tool so far), I'll have to dig into the indexing and pipelining to make sense of it as this kernel is quite involved.

But first, I'm gonna send identities through in a more minimal reproducer, hopefully a pattern will emerge.

[ahadnagy]

Race condition confirmed. I tested the kernel with identity weights, and the output is randomly messed up.
Here's a reproducer:
https://github.com/ahadnagy/optimum-quanto/blob/4a6d678239d4a62b94ebce06bbfe0e6e7235abb0/test/tensor/weights/optimized/test_marlin_int4_weight_qbits_tensor.py#L175-L207

I'm not sure how long it's gonna take to debug this. Maybe it's worth submitting an issue to the Marlin repo?

[dacorvo]

@ahadnagy thank you for investigating this. This is not the original marlin kernel: it has been modified to integrate a shift in addition to the scale.
I isolated the modifications to the original kernel in this commit:
616980f
Maybe that would help identifying if the issue comes from the original kernel or from the modifications.