Skip to content

Commit

Permalink
Use new litserve accelerator="auto" setting (#1336)
Browse files Browse the repository at this point in the history
  • Loading branch information
rasbt authored Apr 24, 2024
1 parent 9d032dd commit 65cc8d0
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 4 deletions.
7 changes: 4 additions & 3 deletions litgpt/deploy/serve.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def setup(self, device: str) -> None:

fabric = L.Fabric(
accelerator=device.type,
devices=1 if device.type=="cpu" else [device.index], # TODO: Update once LitServe supports "auto"
devices=1 if device.type=="cpu" else [device.index],
precision=precision,
)
checkpoint_path = self.checkpoint_dir / "lit_model.pth"
Expand Down Expand Up @@ -99,7 +99,7 @@ def run_server(
top_k: int = 200,
max_new_tokens: int = 50,
devices: int = 1,
accelerator: str = "cuda",
accelerator: str = "auto",
port: int = 8000
) -> None:
"""Serve a LitGPT model using LitServe
Expand All @@ -114,7 +114,8 @@ def run_server(
generated text but can also lead to more incoherent texts.
max_new_tokens: The number of generation steps to take.
devices: How many devices/GPUs to use.
accelerator: The type of accelerator to use. For example, "cuda" or "cpu".
accelerator: The type of accelerator to use. For example, "auto", "cuda", "cpu", or "mps".
The "auto" setting (default) chooses a GPU if available, and otherwise uses a CPU.
port: The network port number on which the model is configured to be served.
"""
check_valid_checkpoint_dir(checkpoint_dir, model_filename="lit_model.pth")
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ dependencies = [
"torch>=2.2.0",
"lightning==2.3.0.dev20240328",
"jsonargparse[signatures]>=4.27.6",
"litserve==0.0.0.dev2", # imported by litgpt.deploy
"litserve>=0.1.0" # imported by litgpt.deploy
]

[project.urls]
Expand Down

0 comments on commit 65cc8d0

Please sign in to comment.