Enable greedy sampling #70
Conversation
|
@michaelbenayoun Could you please take a look when you get a chance? |
|
|
||
| return np.array(results) | ||
|
|
||
|
|
There was a problem hiding this comment.
Parameterize this test to test all the generative models we support.
There was a problem hiding this comment.
We are testing sampling for GPT and BART models at the moment. In the revision, I've currently included t5-small, and will include more models over the next few weeks.
There was a problem hiding this comment.
Where is it tested for GPT and BART?
Anyways, alright let's do that!
There was a problem hiding this comment.
GPT and BART will be committed after we merge this CR. We are currently fixing a few bugs,
|
|
||
| return np.array(results) | ||
|
|
||
|
|
There was a problem hiding this comment.
Where is it tested for GPT and BART?
Anyways, alright let's do that!
optimum/neuron/generation/utils.py
Outdated
| next_token_logits = outputs.logits[:, -1, :] | ||
|
|
||
| # pre-process distribution | ||
| next_tokens_scores = logits_processor(input_ids, next_token_logits) |
There was a problem hiding this comment.
In the non-XLA version, input_ids is of dim bs x seq_length. With the padding we introduced here, dimensions change. Also, we don't know which exact processor is being used and whether it is supported by XLA -> I think we can't expect each logit processor to be XLA-compatible, so we should probably compute on CPU.
| next_tokens_scores = logits_processor(input_ids, next_token_logits) | |
| if is_torch_tpu_available(): | |
| input_ids_ = input_ids.to('cpu')[:, :seq_length] | |
| next_token_logits_ = next_token_logits.to('cpu') | |
| next_tokens_scores = logits_processor(input_ids_, next_token_logits_) | |
| next_tokens_scores = next_tokens_scores.to(input_ids.device) | |
| else: | |
| next_tokens_scores = logits_processor(input_ids, next_token_logits) |
There was a problem hiding this comment.
@aashiqmuhamed want to commit this?
There was a problem hiding this comment.
Yes, I'm committing without the is_torch_tpu_available() check, since we expect to run on Trainium by default.
optimum/neuron/generation/utils.py
Outdated
|
|
||
| import torch | ||
| import torch.distributed as dist | ||
| import torch_xla.core.xla_model as xm |
There was a problem hiding this comment.
This causes an import error.
Can you do:
from ..utils import is_torch_xla_available
if is_torch_xla_available():
import torch_xla.core.xla_model
Basically we want to be able to import and test code on regular machines when we do not need all of this.
There was a problem hiding this comment.
Got it, is the optimum neuron library designed for both CPU and Trainium?
There was a problem hiding this comment.
Only Trainium but it is more to be able to run non-trainum dependent tests on regular machines. Without this the test will fail if we dont have torch_xla installed even though we want to test something unrelated.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. |
|
The EC2 runners cannot be used because they need secrets from this repo, which you do not have on your fork... Can you run the following command please: This will fix the styling error you currently have. |
This CR enables greedy sampling in model.generate on XLA devices such as Trainium and TPU. This addresses issues such as huggingface/transformers#18661 and huggingface/transformers#12322.
The implementation is inspired by the corresponding Tensorflow generate function in transformers. The CR uses conditional statements to support greedy sampling, and also implements kv-cache functionality that is XLA compatible.