hl.triton_kernel codegen at the wrong places
Check examples/distributed/allreduce_bias_rmsnorm.py, when Helion codegen looped reduction, the hl.triton_kernel(symm_mem_sync, ...) will be codegend outside the loop. But it should be put inside the code. This cause accuracy check fail.
To repro:
1. Add reduction_loops=[512] to the helion config of the one_shot_allreduce_bias_rmsnorm_kernel kernel
2. run `torchrun --nproc-per-node=8 examples/distributed/allreduce_bias_rmsnorm.py`
Example output code: https://gist.github.com/shunting314/2a82a1c2404f647861cecbbe70cda74a#file-cnnpqtecfztyd6t7y7o5ano6tzbumflhkutgselcqlvnx2wr3f3u-py-L138
FYI @jansel , @yf225
0 条评论