Adding ScatterMoE kernel support for Granite models.

[shawntan]

What does this PR do?

Adds ScatterMoE kernel support for Granite MoE models.
Started in #40365 but has significantly deviated in approach, so starting a new pull request.

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@MekkCyber already started to provide some comments in #40365.

[MekkCyber]

Thanks a lot for this pr @shawntan ! can you open a pr here : https://github.com/huggingface/kernels-community to add your kernel without the build folder, so that we can review the source code, and once merged we will upload the builds on the hub, and then we can merge this pr so everyone can use the kernel. If you don't have time i'm happy to help 🤗

[shawntan]

So far I've simply been copying from torch-ext to build/torch-universal to "build" the kernel.

Should I follow the example here: https://github.com/huggingface/kernels-community/tree/main/trimul_gpumode

It's another triton based model with universal "build".

[shawntan]

Should I be targetting main or the latest release?

[shawntan]

Some changes have been made to GraniteMoE in the latest main that affects how the kernel is used: e6a8e7d...7938e91#diff-49f545f0c27ad25a565417d120d925f6bfcde532d6e0ea5539dcec182ad978aaL373

I haven't been able to test out the current main branch successfully, there seems to be many breaking changes, so I need some help in modifying the current PR:

1. Should I target the main branch or the latest release?
2. Should I be modifying GraniteMoE to the way it is supposed to be?
3. Changing GraniteMoeMoE back and regenerating the modeling_ files of the downstream classes would fix the issue, but what other models with MoEs are affected and how was the change decided?

[shawntan]

@MekkCyber will need some help here. Which branch should I target for the PR: main or the latest release?

main has issues with the hf API URLs. The current state of the PR targets the latest release.

[MekkCyber]

Hey @shawntan, let's target main! I think you did some faulty rebase, let's only keep the relevant files in the PR.
I'm not sure what you mean by modifying GraniteMoE but we shouldn't touch the modeling files, the kernel should be applied using the decorator solely.

[shawntan]

I've reverted the PR back to the original one against main

My issue was this line:

main: https://github.com/huggingface/transformers/blob/main/src/transformers/models/granitemoe/modeling_granitemoe.py#L235

v4.57.1: https://github.com/huggingface/transformers/blob/v4.57.1/src/transformers/models/granitemoe/modeling_granitemoe.py#L320

router_logits has been removed. I cannot figure out why or how it got removed, but I think it breaks the Granite models downstream of that class. It's also hard for me to test this since pulling from the hub is also broken in the latest main.

[MekkCyber]

router_logits has been removed

Yes it was removed in the MoE refactor for vllm compatibility with transformers but this shouldn't break anything in the transformers implementation

It's also hard for me to test this since pulling from the hub is also broken in the latest main.

What do you mean ?

[shawntan]

router_logits has been removed

Yes it was removed in the MoE refactor for vllm compatibility with transformers but this shouldn't break anything in the transformers implementation

I see. I will need to change the layer definition in the kernel since it will only produce one output, instead of the tuple it returns right now. I am also confused as to how the auxiliary loss will work if the MoE doesn't return the router_logits. I will look into it.

Update:
Looking at it, the current changes break the output_router_logits=True path.

The attribute is checked for during the forward call:

transformers/src/transformers/models/granitemoe/modeling_granitemoe.py

Line 725 in ac81541

if output_router_logits:

But it is never passed to the model:

transformers/src/transformers/models/granitemoe/modeling_granitemoe.py

Lines 696 to 704 in ac81541

	outputs = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	cache_position=cache_position,
	**kwargs,
	)

There is also no pathway for passing the router logits, or recomputing them in GraniteMoeModel:

transformers/src/transformers/models/granitemoe/modeling_granitemoe.py

Lines 492 to 551 in ac81541

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Cache] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	use_cache: Optional[bool] = None,
	cache_position: Optional[torch.LongTensor] = None,
	**kwargs: Unpack[TransformersKwargs],
	) -> MoeModelOutputWithPast:
	if (input_ids is None) ^ (inputs_embeds is not None):
	raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
	if use_cache and past_key_values is None:
	past_key_values = DynamicCache(config=self.config)
	if inputs_embeds is None:
	inputs_embeds = self.embed_tokens(input_ids)
	if cache_position is None:
	past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
	cache_position = torch.arange(
	past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
	)
	if position_ids is None:
	position_ids = cache_position.unsqueeze(0)
	causal_mask = create_causal_mask( # ONLY DIFF WITH MIXTRAL: NO SLIDING
	config=self.config,
	input_embeds=inputs_embeds,
	attention_mask=attention_mask,
	cache_position=cache_position,
	past_key_values=past_key_values,
	position_ids=position_ids,
	)
	inputs_embeds = inputs_embeds * self.embedding_multiplier
	hidden_states = inputs_embeds
	# create position embeddings to be shared across the decoder layers
	position_embeddings = self.rotary_emb(hidden_states, position_ids)
	for decoder_layer in self.layers[: self.config.num_hidden_layers]:
	hidden_states = decoder_layer(
	hidden_states,
	position_embeddings=position_embeddings,
	attention_mask=causal_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	use_cache=use_cache,
	cache_position=cache_position,
	**kwargs,
	)
	hidden_states = self.norm(hidden_states)
	return MoeModelOutputWithPast( # only diff with Mistral is the output type, we need MoE
	last_hidden_state=hidden_states,
	past_key_values=past_key_values,
	)

What's the ideal way that the HF team is thinking of allowing for this while still maintaining vllm compatibility? I can make the necessary changes.

It's also hard for me to test this since pulling from the hub is also broken in the latest main.

What do you mean ?

Hmm, sorry, I can't seem to reproduce the same issue I saw before with the jinja files.

Still, there are some incompatibilities from the Mamba end of things, right now with both repositories on main, the error being thrown is this:

  File "/proj/checkpoints/shawntan/transformers/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py", line 47, in <module>
    from mamba_ssm.ops.triton.selective_state_update import selective_state_update
  File "/u/shawntan/mamba/mamba_ssm/__init__.py", line 6, in <module>
    from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
  File "/u/shawntan/mamba/mamba_ssm/models/mixer_seq_simple.py", line 20, in <module>
    from mamba_ssm.utils.generation import GenerationMixin
  File "/u/shawntan/mamba/mamba_ssm/utils/generation.py", line 14, in <module>
    from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer
ImportError: cannot import name 'GreedySearchDecoderOnlyOutput' from 'transformers.generation' (/proj/checkpoints/shawntan/transformers/src/transformers/generation/__init__.py)

[shawntan]

TL;DR:

1. If I change the kernel to suit the current state of GraniteMoeMoE, it would need 1 output, but the problem would be the calculating of the aux loss would be broken.
2. I don't believe that is the intended behaviour, so my question is: How should the GraniteMoeMoE class be modified to best support the vLLM compatibility while not breaking previous functionality (aux loss)?

@MekkCyber

[ArthurZucker]

Hey! #41580 was merged, IDK if you need more changes but indeed for collecting router logits you might need a small change! You could put the collection on the gate linear layer?

[github-actions]

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

run-slow: granitemoe, granitemoehybrid, granitemoeshared