huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 04:44:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: set sharded true if WORLD_SIZE is set

Browse Source
This commit is contained in:
drbh 2024-06-12 17:12:18 +00:00
parent cdbf802860
commit 9854f20225
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
server/text_generation_server/cli.py
View File

@ -44,6 +44,9 @@ def serve(
otlp_endpoint: Optional[str] = None, otlp_endpoint: Optional[str] = None,
max_input_tokens: Optional[int] = None, max_input_tokens: Optional[int] = None,
): ):
# derive sharded from environment variables if not provided
sharded = sharded or os.getenv("WORLD_SIZE", None) is not None
if sharded: if sharded:
assert ( assert (
os.getenv("RANK", None) is not None os.getenv("RANK", None) is not None