[Fp8] Fix experts

[vasqu]

As per title, the current fp8 experts implementation is wrong - likely related to #42456 (cc @3outeille)

Before:

• The number of experts was treated as the number of top k experts that are hit, i.e. num_experts = top_k_weights.shape[1]
  • This is simply wrong as now our indices which can cover all experts is limited to top_k
  • This can trigger device assert errors as one hot might not have the correct indices for the number of classes indicated by num_experts
• The hit expert is now used which is correct if we passed the selection before
  • Only due to the indices being correctly limited

After:

• We now follow the current eager implementation, e.g. Mixtral:

  transformers/src/transformers/models/mixtral/modeling_mixtral.py

  Lines 82 to 86 in 88a5623

  	final_hidden_states = torch.zeros_like(hidden_states)
  	with torch.no_grad():
  	expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
  	expert_mask = expert_mask.permute(2, 1, 0)
  	expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

• The portion where we iterate through experts should only have changed the linear calls of

  transformers/src/transformers/models/mixtral/modeling_mixtral.py

  Lines 88 to 98 in 88a5623

  	for expert_idx in expert_hit:
  	expert_idx = expert_idx[0]
  	if expert_idx == self.num_experts:
  	continue
  	top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
  	current_state = hidden_states[token_idx]
  	gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
  	current_hidden_states = self.act_fn(gate) * up
  	current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
  	current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
  	final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

  • The key is to redirect the correctly hit expert to the top k weight via

    transformers/src/transformers/models/mixtral/modeling_mixtral.py

    Line 92 in 88a5623

    top_k_pos, token_idx = torch.where(expert_mask[expert_idx])

I'm not sure how long it has been this way, but it definitely is not correct atm. I'd be curious if I'm overseeing something here but without this fix #42028 will not work for the original model - my suspicion is on mixtral having a bias towards low indices which might have saved this oversight.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Yes absolutely! I did not check but if a test is missing let's add one!

[ArthurZucker]

@3outeille note that this breaks TP but same as it is broken for now as well

[SunMarc]

LGTM

[SunMarc]

thanks ! maybe we can add a small comment to say that it was mostly copied from deepspeed_v3 modeling, so that we should propagate the changes here also in the future

[vasqu]

Added a comment, referencing mixtral

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: finegrained_fp8

[vasqu]

This acts as a sanity integration check but it depends on the minimax m2 PR (#42028) so I will remove this skip when merging that PR

I think this is the easiest way as these weights force the issue

[github-actions]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43154&sha=d0d860