Repository navigation
fix(serve): size tensor parallelism from the target instance's GPUs - #6433
Open
AkshayShah03 wants to merge 1 commit into
Open
AkshayShah03 wants to merge 1 commit into
AkshayShah03 wants to merge 1 commit into
Conversation
_get_default_tensor_parallel_degree used the GPU count of the machine running the build whenever it had any GPUs, and only fell back to the gpu_count passed in for the target instance otherwise. The SageMaker endpoint builds for TGI (NUM_SHARD) and DJL (TENSOR_PARALLEL_DEGREE) pass the target instance's GPU count, so building on an 8-GPU box for a 1-GPU ml.g5.2xlarge set the degree to 8 (the container cannot start), and building on a 1-GPU notebook for a 4-GPU ml.g5.12xlarge set it to 1. Prefer the given gpu_count and only probe local GPUs when it is not provided (local container mode).
AkshayShah03
requested a deployment
to
manual-approval
October 12, 2026 17:55 — with
GitHub Actions
Waiting
AkshayShah03
requested a deployment
to
manual-approval
October 12, 2026 17:56 — with
GitHub Actions
Waiting
AkshayShah03
requested a deployment
to
manual-approval
October 12, 2026 17:56 — with
GitHub Actions
Waiting
This branch is waiting to be deployed
This file contains hidden or 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
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.
Issue #, if available: N/A
Description of changes:
_get_default_tensor_parallel_degreeused the GPU count of the machine running the build whenever that machine had any GPUs. It only used thegpu_countpassed in when the build machine had none:The SageMaker endpoint builds for TGI (
NUM_SHARD) and DJL (TENSOR_PARALLEL_DEGREE) pass the GPU count of the target instance, taken frominstance_type. So:ml.g5.2xlargeset the degree to 8, and the container cannot start;ml.g5.12xlargeset it to 1.This change uses
gpu_countwhen it's given, and only probes local GPUs when it isn't, as in local container mode, where the callers don't pass it.Testing:
tests/unit/model_server/test_djl_utils.py: the target count wins in both directions (1 local → 4 target, 8 local → 1 target), and withoutgpu_countthe local GPUs are used. The first fails before this change.tests/unit/model_serverandtest_model_builder_servers_hf_model_id.py: 134 passed.By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.