Performance problem with multi head training

[bhcao]

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.

[lwalew]

Hi @bhcao , apologies for the late reply. Our internal testing reveals that for models of a "normal" size the slowdown scales as expected. Indeed, if restricting the training to homogenous batches which only ever use a single head within a batch, we could obtain slightly better training speeds. But this is unfortunately not something we will be prioritizing at the moment.