[AIR][GPU Batch Prediction] Add basic support for GPU batch prediction #26251
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
jiaodong
commented
Jul 2, 2022
| output = self.model(model_input) | ||
|
|
||
| def untensorize(torch_tensor): | ||
| numpy_array = torch_tensor.cpu().detach().numpy() | ||
| numpy_array = torch_tensor.detach().cpu().numpy() |
There was a problem hiding this comment.
this is minor nit improvement so we don't construct autograd edge on cpu then detach it. End result is the same.
| self.use_gpu = use_gpu | ||
| if use_gpu: | ||
| # Ensure input tensor and model live on GPU for GPU inference | ||
| self.model.to(torch.device("cuda")) |
There was a problem hiding this comment.
model init and config is moved to constructor rather than upon every predict() call.
| with tf.device("GPU:0"): | ||
| self.model = self.model_definition() | ||
| else: | ||
| self.model = self.model_definition() |
There was a problem hiding this comment.
model init and config is moved to constructor rather than upon every predict() call.
jiaodong
added
the
tests-ok
ericl
reviewed
Jul 5, 2022
ericl
reviewed
Jul 5, 2022
ericl
reviewed
Jul 5, 2022
ericl
reviewed
Jul 5, 2022
ericl
reviewed
Jul 5, 2022
ericl
added
@author-action-required
jiaodong
added
tests-ok
ericl
reviewed
Jul 6, 2022
ericl
reviewed
Jul 6, 2022
ericl
reviewed
Jul 6, 2022
ericl
approved these changes
Jul 6, 2022
jiaodong
removed
the
tests-ok
amogkam
reviewed
Jul 11, 2022
.buildkite/pipeline.gpu.large.yml
Outdated
| @@ -39,6 +39,8 @@ | |||
| - DATA_PROCESSING_TESTING=1 TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh | |||
| - pip install -Ur ./python/requirements_ml_docker.txt | |||
| - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=gpu python/ray/air/... | |||
| - pip install -U horovod | |||
There was a problem hiding this comment.
Ah instead of manually installing here you can set the INSTALL_HOROVOD=1 env var for install-dependencies.sh
There was a problem hiding this comment.
oh nice let me change it
| "//python/ray/air:__pkg__", | ||
| "//python/ray/air:__subpackages__", | ||
| "//python/ray/train:__pkg__", | ||
| "//python/ray/train:__subpackages__", |
There was a problem hiding this comment.
For my own understanding, what does this do?
There was a problem hiding this comment.
i noticed that ray train tests are actually using modules from air folder, such as the test_horovod. So I setup both ML and AIR bazel targets to be visible to each other so tests and directly import and use if needed.
|
@jiaodong actually looks like the lint failure is related. This PR removes |
jiaodong
added
the
tests-ok
Stefan-1313
pushed a commit
to Stefan-1313/ray_mod
that referenced
this pull request
Aug 18, 2022
ray-project#26251) This PR adds GPU support for pytorch and tensorflow predictor, as well as automatic setting `use_gpu` flag in `BatchPredictor`. Notable changes: - Added `use_gpu` flag in the constructor of `TorchPredictor` and `TensorflowPredictor` (note it's slightly different from our latest design doc that puts this flag at `predict()` call) - Added `use_gpu` flag to `SklearnPredictor` so its interface is compatible with `BatchPredictor` - Code to move both model weights and input tensor to default visible GPU at index 0 if flag is set - parametrized existing predictor tests to use GPU for both CPU & GPU coverage - Changed BUILD CI tests with an added `gpu` tag (I'm not 100% sure if that's a right way tho) Follow ups: ray-project#26249 is created in case our host has multiple GPU devices. It's a bit out of scope for this PR, but for GPU batch inference ideally we should be able to evenly use all GPU devices on host where CPU & DRAM are busy with pre-fetching + data movement to GPU. We might approximately do the same by scheduling same # of Predictor instances on the host, but that's worth verifying once benchmarks are set. Signed-off-by: Stefan van der Kleij <[email protected]>
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.
Why are these changes needed?
This PR adds GPU support for pytorch and tensorflow predictor, as well as automatic setting
use_gpuflag inBatchPredictor.Notable changes:
use_gpuflag in the constructor ofTorchPredictorandTensorflowPredictor(note it's slightly different from our latest design doc that puts this flag atpredict()call)use_gpuflag toSklearnPredictorso its interface is compatible withBatchPredictorgputag (I'm not 100% sure if that's a right way tho)Follow ups:
#26249 is created in case our host has multiple GPU devices. It's a bit out of scope for this PR, but for GPU batch inference ideally we should be able to evenly use all GPU devices on host where CPU & DRAM are busy with pre-fetching + data movement to GPU. We might approximately do the same by scheduling same # of Predictor instances on the host, but that's worth verifying once benchmarks are set.
Related issue number
Checks
scripts/format.shto lint the changes in this PR.