Migrate logits computation and gather to model_runner

[esmeetu]

1. Finish TODO: Use NCCL instead of ray for control-plane communication to remove serialization overhead #2221 (comment)
2. This might replace Fix vocab_size inconsistency for sampler #2398 . (refer to [Misc][Log] Add log for tokenizer length not equal to vocabulary size #3500)

Description

This PR migrate logits computation and gather to model_runner. This change make Sampler simple and clean.
Furthermore, i want to remove sample method from model file, like llama.py. Because model and sample are at different stage, and should decouple with each other. But in this comment of PR #3183 : #3183 (comment) , that model will scale logits in sampler, so i will keep sample method as it is. This PR will better support #3183 model integration as well.

Pipeline:

1. Prepare inputs
2. Model
3. LogitsProcessor (support logits scale and Custom LogitsProcessor Functions)
4. Sampler
5. Output

TODO

• Fix llama, opt for tests
• Fix tests (sampler and logits_processor)
• Fix other model files
  - [ ] Fix Fix vocab_size inconsistency for sampler #2398

[grandiose-pizza]

Hi @esmeetu , Important update.

Regarding this comment, #3183 (comment)

I found the bug. It will be useful for this PR.

Please add this change while making the PR for scaling logits: e04e56d

This is done as during multiGPU setting logits are None for get_tensor_model_parallel_rank() that are > 0 (it will be 0 & 1 in a 2GPU setting) as per this:

vllm/vllm/model_executor/parallel_utils/communication_op.py

Line 102 in a33ce60

output_tensor = None

[Yard1]

While I agree with this change in principle, I think it's important to ensure we have an API that can support different usecases. For example, I would suggest making the logit generation a layer (or some other sort of abstraction on the model level). The fact that it has now been moved out of the model into the model runner makes the code harder to understand, especially considering the sampler remains a layer.

In other words, I would just suggest adding a new logit generator layer to the models (or the sampler, though models would be better - the output of a model should be logits IMO) and not putting that logic inside of model runner.

[esmeetu]

@Yard1 @zhuohan123 I redesign this PR, which make logits processor an individual layer. PTAL!
After this design is supported, i will fix other model files.

[Yard1]

This need changes to work with the LoRA path

[esmeetu]

@Yard1 All CI passed, please review this again. cc @zhuohan123
And should we resolve #2398 into this PR? Is there any better solution?

[zhuohan123]

Thanks for the fix! This looks very good. Much better than what we have before :) Can you merge the PR after fixing the merge conflicts?