[llama4][auxiliary-loss-free load balancing] update expert_bias without backward hooks

[hann-wang]

Changes:

• Introduce a finalize_model_grads_func attribute in TrainSpec.
• Set finalize_model_grads_func to update_router_expert_bias for MoE models.
• finalize_model_grads_func is called AFTER gradient accumulation steps.
• enable_tp2ep is reserved for [llama4] enable expert parallel on the same device mesh as tp (tp2ep) #1269

Reasons:

• Friendly for torch.compile and activation checkpointing.
• The original implementation updates expert_bias on each microbatches during gradient accumulation.

cc @tianyu-l