Add support for transformers-neuronx continuous batching #488
Conversation
|
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. |
dacorvo
force-pushed
the
continuous_batching
branch
6 times, most recently
from
7c11422
to
3862e28
Compare
dacorvo
force-pushed
the
continuous_batching
branch
from
3862e28
to
f45d620
Compare
optimum/exporters/neuron/base.py
Outdated
| @@ -379,7 +379,7 @@ class NeuronDecoderConfig(NeuronConfig): | |||
| be passed to export the model, | |||
| - NEURONX_CLASS (`str`) -- the name of the transformers-neuronx class to instantiate for the model. | |||
| It is a full class name defined relatively to the transformers-neuronx module, e.g. `gpt2.model.GPT2ForSampling` | |||
| [`~optimum.utils.DummyInputGenerator`] specifying how to create dummy inputs. | |||
| - CONTINUOUS_BATCHING (`bool`) -- Whether the model supports continuous batching or not. | |||
There was a problem hiding this comment.
| - CONTINUOUS_BATCHING (`bool`) -- Whether the model supports continuous batching or not. | |
| - CONTINUOUS_BATCHING (`bool`, defaults to `False`) -- Whether the model supports continuous batching or not. |
| if seq_ids is None: | ||
| seq_ids = torch.arange(input_ids.shape[0]) | ||
| else: | ||
| assert seq_ids.shape[0] == input_ids.shape[0] |
There was a problem hiding this comment.
nit: maybe raise an explicit error with a message saying that the shapes should match in this case
There was a problem hiding this comment.
I am using an assert because these methods are always called internally, so it is rather to catch internal programming errors.
optimum/neuron/modeling.py
Outdated
| input_ids: torch.Tensor, | ||
| attention_mask: torch.Tensor, | ||
| seq_ids: Optional[List[int]] = None, | ||
| **kwargs, |
There was a problem hiding this comment.
Is this necessary? I tend to avoid it whenever possible.
There was a problem hiding this comment.
It is a leftover from the GenerationMixin, but I moved away from this pattern and I 100% agree with you on this, at least for us where we are supposed to be a bit more certain about what we accept or not. Will remove it.
| auto_cast_type = auto_cast_type.replace("p", "") | ||
| tnx_kwargs = { | ||
| "batch_size": batch_size, | ||
| "tp_degree": num_cores, |
There was a problem hiding this comment.
The TP degree is always exactly the number of neuron cores used?
There was a problem hiding this comment.
Yes. I use the num_cores name in optimum-neuron because I use it for two things: to set the TP degree (here) and also to restrict the number of cores I reserve at initialization (otherwise the TNX runtime takes them all).
What does this PR do?
This adds support for
transformers-neuronxcontinuous batching.This also activates automatically the continuous batching feature for any model that supports it as soon as the batch size is > 1.
This allows to workaround the fact that static batching is broken in
transformers-neuronxfor these models: aws-neuron/transformers-neuronx#79.But this is also a game changer for NeuronX TGI, as it unlocks:
LlamaandMistralmodels thanks totransformers-neuronx"continuous batching" graphs that share KV cache between nodes,max-batch-sizeparameter.