[DPO/KTO] Mixtral Load Balancing Loss

[claralp]

Question towards @lewtun @kashif:

We noticed that the load balancing loss (aux_loss) that is implemented in MoEs modeling_mixtral.py#L1244 is not added to the loss implemented in DPO/KTO trainers.

Isn't this needed so that the load balancing between experts is still guaranteed after fine-tuning or does the KL loss penalization sufficiently prevents the router weights from changing too much?

[kashif]

great question @claralp so my initial intuition was that we are changing the weight of the model and we keep the routing of the experts in-tact. If one were to also add the aux loss, I am not sure currently how that would fit into the preference tuning objective...

[PhilipMay]

my initial intuition was that we are changing the weight of the model and we keep the routing of the experts in-tact.

@kashif So your suggestion is to use a lora / qlora adapter that does not target the router layers?

[claralp]

@kashif, so I guess it is important to set output_router_logits=False for preference tuning.
Had a look into the code and it looks like setting this does not disable the router weights from being trained.
So where do you disable the weights for routing of experts in preference tuning?
If that is a required manual step, maybe a note should be added to the trl MoE examples.

[kashif]

one clarification... in the preference tuning losses we are not explicitly back-proping over the categorical cross entropy loss of the models but rather via the log-probs obtained from the logits... so in this case these methods are not doing anything to the router's weights in MoE models

[claralp]

@kashif we could observe a gradient in the gate weights even when backpropagating via the log-probs obtained from the logits (as in DPO/KTO).
If the gradient calculation is not explicitly switched off for the 'gate' weights or if they are excluded from the target_modules in the PEFT config (default), these gating weights seem to be trained by DPO/KTO

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.