Skip to content

fix(serve): size tensor parallelism from the target instance's GPUs - #6433

Open
AkshayShah03 wants to merge 1 commit into
aws:masterfrom
AkshayShah03:fix/tensor-parallel-degree-target-gpus
Open

AkshayShah03 wants to merge 1 commit into
aws:masterfrom
AkshayShah03:fix/tensor-parallel-degree-target-gpus

Conversation

@AkshayShah03

Copy link
Copy Markdown

Issue #, if available: N/A

Description of changes:

_get_default_tensor_parallel_degree used the GPU count of the machine running the build whenever that machine had any GPUs. It only used the gpu_count passed in when the build machine had none:

tot_gpus = len(available_gpus) if available_gpus else gpu_count

The SageMaker endpoint builds for TGI (NUM_SHARD) and DJL (TENSOR_PARALLEL_DEGREE) pass the GPU count of the target instance, taken from instance_type. So:

  • building on an 8-GPU box for a 1-GPU ml.g5.2xlarge set the degree to 8, and the container cannot start;
  • building on a 1-GPU notebook for a 4-GPU ml.g5.12xlarge set it to 1.

This change uses gpu_count when 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:

  • New tests in tests/unit/model_server/test_djl_utils.py: the target count wins in both directions (1 local → 4 target, 8 local → 1 target), and without gpu_count the local GPUs are used. The first fails before this change.
  • tests/unit/model_server and test_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.

_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).

This branch is waiting to be deployed

1 waiting deployment
manual-approval — c92ffc84 Waiting Oct 12, 2026 by AkshayShah03 via wait-for-approval #1987
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant