Performance problem with multi head training

Because the graph used by `jax.jit` is static, it will repeatedly calculate all heads during each step, resulting in significant waste, which is particularly obvious when training many heads. We use almost the same multi head implementation as MLIPv2 on a small FeNNol model with 3 heads, it slows down the training speed by 20%. One solution may be to ensure that each step has only one head and select only the corresponding parameter outside the jitted function.

This is a reminder while I have not studied this issue on MLIPv2 in detail. I will conduct performance analysis and write codes when available.