Assisted decoding multi-gpu

[zucchini-nlp]

What does this PR do?

Fixes #35099, we need to manually move inputs to the correct device as assistant and target models can be initialized on different devices. We already move the input to assistant's device when get_candidate_inputs but don't do it in target models' forward

Btw, seems like we are now having more troubles with certain generation techniques or cache not working in multi-gpu setting. And we don't have much tests for multiple gpus, so I am not sure if that's intended or we can add more integration tests. Also if those tests will be run as part of pull request CI or not. cc @gante for this question

cc @jmamou

[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.

[jmamou]

in a previous PR https://github.com/keyboardAnt/transformers/pull/4/files I proposed to do it just after get_candidates at the same place we move candidate_logits to self.device

[zucchini-nlp]

oke, might be better to move it there for easier readability :)

[ArthurZucker]

Thanks! Regarding multi-gpu, we don't have fast tests, but we could add them with github actions.

[jmamou]

@zucchini-nlp @ArthurZucker
For AG with multi-GPU, it would be beneficial to include in the documentation that for optimal speedup, the assistant model should be placed on the default/first device of the target model (target.device). This setup avoids the overhead of transferring the candidate token IDs tensor from assistant.device to target.device after the speculative iteration and back to assistant.device with the validated token IDs following target validation

[zucchini-nlp]

@jmamou yes, feel free to add it in the docs in subsequent open PRs if relevant