how to generate router_logits in moe models using model.generate()?
Author: Jimmy-Lu
Pull request: no
Labels: name: Generation

### System Info

- `transformers` version: 4.41.2
- Platform: Linux-5.4.0-144-generic-x86_64-with-glibc2.31
- Python version: 3.10.0
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.31.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <yes>
- Using distributed or parallel set-up in script?: <yes>


### Who can help?

@gante

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

from transformers import AutoModelForCausalLM, AutoTokenizer
device = "cuda" # the device to load the model onto

model = AutoModelForCausalLM.from_pretrained(
   "/localssd/swlu/Qwen1.5-MoE-A2.7B-Chat",
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("/localssd/swlu/Qwen1.5-MoE-A2.7B-Chat")

prompt = "Give me a short introduction to large language model."
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(device)

generated_ids = model.generate(
    model_inputs.input_ids,
    max_new_tokens=512,
    return_dict_in_generate = True,
    output_router_logits = True
)
print("outputs:", generated_ids.router_logits)

### Expected behavior

I want to get router_logits of moe models using model.generate() with the code above.
But got:
AttributeError: 'GenerateDecoderOnlyOutput' object has no attribute 'router_logits'

Comments:
@amyeroberts  can you explain the label Generation, is this issue still open, as I have started looking into it . Please let me know

Hi @amanjam, 

> can you explain the label Generation

Are you asking about the reason for a label or why `Generation` applies to this issue? 

> is this issue still open

Yes, it it still open. You can check this by looking at the PR title, where there's a green `Open` badge 

<img width="1266" alt="image" src="https://github.com/huggingface/transformers/assets/22614925/2eb7e72f-fb1b-419e-82a2-68ef6bc43978">



 

hi @amyeroberts 
Yes , I am asking why generation applies to this issue? Please explain it a bit

@amanjam To help clarify, it might be easier for you to explain where any confusion might be coming from. The issue title specifically asks about generation, the issue involves the `generate` method and the error message is about the `GenerateDecoderOnlyOutput`

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.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @Jimmy-Lu 👋 Thank you for opening this issue 🤗 

Forwarding MoE-specific outputs is not yet enabled in `generate`, but we'd gladly accept a PR for it!

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.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.