Xla parallel proxy #12
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Originally imported from gemma repository https://github.com/google/gemma_pytorch.git At commit version: cf8658c.
This will allow to execute models in a parallel way and interact with them from the caller thread. To see how it works in output mode, you can launch the test with debug enabled this way: DEBUG=1 pytest -s tests/test_parallel_proxy.py
tengomucho
force-pushed
the
xla-parallel-proxy
branch
from
7317e4b to
c8d7bbe
Compare
tengomucho
force-pushed
the
xla-parallel-proxy
branch
from
c8d7bbe to
747ad57
Compare
mfuntowicz
requested changes
Apr 5, 2024
optimum/tpu/xla_parallel_proxy.py
Outdated
| xmp.spawn(_mp_fn, args=(args), join=True, daemon=False) | ||
|
|
||
|
|
||
| class ModelProxy: |
There was a problem hiding this comment.
I would rename this, it's not clear imo it handles the MP logic, wdyt?
What about something like DistributedModel or DistributedTpuModel ?
optimum/tpu/xla_parallel_proxy.py
Outdated
|
|
||
| def send(self, command: ModelCommand, data: Dict = None): | ||
| # First wait until model is ready to receive commands | ||
| debug(f" MM Command {command} waiting for model to be ready") |
There was a problem hiding this comment.
Let's remove the DEBUG/debug reference and leverage python logging
mfuntowicz
approved these changes
Apr 8, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do?
This PR does several things:
adds an example with
gemma-2bthat show the inference as it worksAdds and implementation that will allow to launch and communicate with models running in parallel using
ModelProxyAdds a test and that to a CI workflow.
Did you write any new necessary tests?