Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix max tensor_parallel_degree #1182

Merged
merged 1 commit into from
Oct 12, 2023
Merged

Fix max tensor_parallel_degree #1182

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Fix max tensor_parallel_degree 
This provides a fix for a "max" tensor_parallel_degree. In the future, this area
probably needs greater cleanup and maybe some refactoring.
  • Loading branch information
@zachgk
zachgk committed Oct 12, 2023
commit 26d4a4ffda1a3f3476d7827d12bf99e3c5767da4
11 changes: 10 additions & 1 deletion wlm/src/main/java/ai/djl/serving/wlm/ModelInfo.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -805,7 +805,16 @@ public String[] getLoadOnDevices() {
int gpuCount = eng.getGpuCount();
String v = Utils.getenv("TENSOR_PARALLEL_DEGREE", "-1");
v = prop.getProperty("option.tensor_parallel_degree", v);
int tpDegree = Integer.parseInt(v);
int tpDegree;
if ("max".equals(v)) {
if (gpuCount > 0) {
tpDegree = gpuCount;
} else {
tpDegree = NeuronUtils.getNeuronCores();
}
} else {
tpDegree = Integer.parseInt(v);
}
if (gpuCount > 0) {
int gpuPerWorker = 1;
if ("Python".equals(engineName)) {
Expand Down
Loading