Changing output to follow Transformers v5 #63
Conversation
| return layer_output | ||
|
|
There was a problem hiding this comment.
so we don't need the router_logits in the kernel ?
There was a problem hiding this comment.
From what I've understood, the new transformers v5 tracks the router_logits via some annotations to the main class: _can_record_outputs and OutputRecorder (huggingface/transformers#40822 (comment)). @ArthurZucker has confirmed those annotations will be added to GraniteMoe, making the output signature same as all the other MoE models, where the router logits are similarly tracked via those annotations.
I'm not sure if the underlying mechanism that tracks the logits for v5 will work with the kernel, but it's hard to test since for GraniteMoeHybrid, we need mamba_ssm to be compatible with v5.
There was a problem hiding this comment.
Thanks for clarifications @shawntan, I don't get though why it's hard to test ?
There was a problem hiding this comment.
Ok, Mamba's been updated, I've tested it. Sorry about that.
Everything seems good to go!
2 participants
Modifying kernel to match transformers v5 based on: huggingface/transformers#40822 (comment)