fix: change jaxbench 42k to avoid producing zero output - #76
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
| x = jnp.matmul(x, weight) + bias | ||
| x = jnp.max(x, axis=1, keepdims=True) | ||
| x = x - jnp.mean(x, axis=1, keepdims=True) | ||
| x = jnp.max(x, axis=1, keepdims=True) |
There was a problem hiding this comment.
why do we need to change the order of operations? Looks like the new order is not consistent with the order in the folder name.
| x = jnp.matmul(x, weight) + bias | ||
| x = jnp.max(x, axis=1, keepdims=True) | ||
| x = x - jnp.mean(x, axis=1, keepdims=True) | ||
| x = jnp.max(x, axis=1, keepdims=True) |
There was a problem hiding this comment.
The original operation will always output zero no matter what input you give. Agent will find this issue and produce a kernel that just produce zero.The kernel will be super fast and the speedup can be 100×. So the original operation is actually not a good test case.
There was a problem hiding this comment.
I understand that you changed x,weight,bias to non-zero values, but I don't understand why do you change matmul->max->subtract->gelu to matmul->subtract->max->gelu? My confusion part is the order in computation().
There was a problem hiding this comment.
Looks like the original code still keeps matmul-max-subtract-gelu. https://github.com/AI-Hypercomputer/accelerator-agents/blob/main/JAXBench/benchmark/42k_Gemm_Max_Subtract_GELU/baseline.py#L23-L28
There was a problem hiding this comment.
Changing the input is not enough. No matter what input, the final result will be zero. What I did is actually changing the operation. We may change the name of this operation if you think it is needed.
There was a problem hiding this comment.
There might be another way to make this operation makes more sense: we should not change the x to be its maximum. I will make a new commit.
…out overwriting x
NinaCai
left a comment
NinaCai
left a comment
There was a problem hiding this comment.
Out of curiosity: How do we decide the benchmark here? Is it extracted from maxtext or just a combinations of different operations? If such combination is never used in any model, why do we want to benchmark them?
It is adapted from KernelBench https://github.com/ScalingIntelligence/KernelBench/blob/main/KernelBench/level2/80_Gemm_Max_Subtract_GELU.py. |
2 participants
No description provided.