ITADN

foreach optimizer pattern for much faster @compile with large number of parameter leafs

#2929Opencsvance 创建于 2026-05-29
C
csvancecommented
Currently if you try to `@compile` the optimizer update step for a model with many parameter leafs, it takes a massive amount of time because Reactant/XLA has to handle the exact same element wise update across every leaf individually. With models like ResNet50 and Adam optimizer variants, I've observed compile times close to an hour just for the optimizer step. The solution is to flatten all parameters and handle the update in a single step which is how PyTorch handles it with its foreach flag. There isn't much documentation about this as far as I can tell, so it might be worth some sort of abstraction/utility functions either here or in downstream packages to automatically handle this seamlessly. If this sort of thing just worked without hour long compile times, it would be much easier to pickup Reactant.jl and be immediately productive.
7 条评论