[FEATURE] Why not Batch Predict?

[urmichm]

Is your feature request related to a problem?
TextSimilarityCrossEncoderModel predictions are performed one by one inside a loop using getPredictor().predict(input) rather than processing the inputs in batches using batchPredict(listOfInputs)

This approach looks inefficient. Is there a specific reason for processing the predictions one by one?

To my understanding, Predictor.java creates internally a list/batch with a single element and calls batchPredict anyway.

Similar behaviour is presented on TextEmbeddingModel

What solution would you like?
Use batchProcessing to improve efficiency

[aaslan1997]

main...aaslan1997:ml-commons:main#diff-df7084e40b9a935bccf807556debb3c55fd387263febfff87badc26d7f2b8774

I started working on this, so in theory my code works for batch predicting but i get some CUDA errors if it runs on GPU on some inputs (not always). I'm not sure if the problem is in ml-commons or in djl.

"root_cause": [
{
"type": "engine_exception",
"reason": "engine_exception: default_program(95): error: extra text after expected end of number\n aten_mul[(long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)] = __float2bfloat16(v * -3.389531389251535e+38.f);\n ^\n\ndefault_program(99): error: extra text after expected end of number\n aten_add[(long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)] = __float2bfloat16(v_1 / 8.f + v_2 * -3.389531389251535e+38.f);\n ^\n\n2 errors detected in the compilation of "default_program".\n\nnvrtc compilation failed: \n\n#define NAN __int_as_float(0x7fffffff)\n#define POS_INFINITY int_as_float(0x7f800000)\n#define NEG_INFINITY int_as_float(0xff800000)\n\n\ntemplate\n__device T maximum(T a, T b) {\n return isnan(a) ? a : (a > b ? a : b);\n}\n\ntemplate\n__device T minimum(T a, T b) {\n return isnan(a) ? a : (a < b ? a : b);\n}\n\n\n#define __BFLOAT16_TO_US(var) (reinterpret_cast<unsigned short>(&(var)))\n#define __BFLOAT16_TO_CUS(var) \\n (reinterpret_cast<const unsigned short>(&(var)))\n\ntypedef struct align(2) {\n unsigned short x;\n}\n__nv_bfloat16_raw;\n\n#if defined(__cplusplus)\nstruct align(2) __nv_bfloat16 {\n host device __nv_bfloat16() {}\n\n host device __nv_bfloat16& operator=(const __nv_bfloat16_raw& hr) {\n x = hr.x;\n return this;\n }\n\n protected:\n unsigned short x;\n};\n\n#if defined(CUDACC)\n__device unsigned short __internal_float2bfloat16(\n const float f,\n unsigned int& sign,\n unsigned int& remainder) {\n unsigned int x;\n\n x = __float_as_uint(f);\n\n if ((x & 0x7fffffffU) > 0x7f800000U) {\n sign = 0U;\n remainder = 0U;\n return static_cast(0x7fffU);\n }\n sign = x >> 31;\n remainder = x << 16;\n return static_cast(x >> 16);\n}\n\n/ Definitions of intrinsics */\n__device __nv_bfloat16 __float2bfloat16(const float a) {\n __nv_bfloat16 val;\n#if CUDA_ARCH >= 800\n asm("{ cvt.rn.bf16.f32 %0, %1;}\n" : "=h"(__BFLOAT16_TO_US(val)) : "f"(a));\n#else\n __nv_bfloat16_raw r;\n unsigned int sign;\n unsigned int remainder;\n r.x = internal_float2bfloat16(a, sign, remainder);\n if ((remainder > 0x80000000U) ||\n ((remainder == 0x80000000U) && ((r.x & 0x1U) != 0U))) {\n r.x++;\n }\n val = r;\n#endif\n return val;\n}\n\n__device float __bfloat162float(const __nv_bfloat16 a) {\n float val;\n asm("{ mov.b32 %0, {0,%1};}\n" : "=f"(val) : "h"(__BFLOAT16_TO_CUS(a)));\n return val;\n}\n#endif /* defined(CUDACC) /\n#endif / defined(__cplusplus) /\n#undef __BFLOAT16_TO_US\n#undef __BFLOAT16_TO_CUS\n\nextern "C" global\nvoid fused_mul_div_add(__nv_bfloat16 tattention_scores_2, _nv_bfloat16* tv, __nv_bfloat16* aten_add, __nv_bfloat16* aten_mul) {\n{\nif (blockIdx.x<1ll ? 1 : 0) {\nif ((long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)<7ll ? 1 : 0) {\nif (blockIdx.x<1ll ? 1 : 0) {\n float v = _bfloat162float(tv[(long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)]);\n aten_mul[(long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)] = __float2bfloat16(v * -3.389531389251535e+38.f);\n } } }if ((long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)<588ll ? 1 : 0) {\n float v_1 = __bfloat162float(tattention_scores_2[(long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)]);\n float v_2 = _bfloat162float(tv[((long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)) % 7ll]);\n aten_add[(long long)(threadIdx.x) + 512ll * (long long)(blockIdx.x)] = __float2bfloat16(v_1 / 8.f + v_2 * -3.389531389251535e+38.f);\n }}\n}\n"
}
],

Maybe someone has some ideas or can help me with that?

[urmichm]

More details on the float issue
deepjavalibrary/djl#3089

which is related to pytorch here
pytorch/pytorch#107503

[urmichm]

@aaslan1997 let's use your changes as the base for the fix