{"author":"Guilherme Penedo","author_email":"guilherme@huggingface.co","bugtrack_url":null,"classifiers":["Programming Language :: Python :: 3","Programming Language :: Python :: 3.10","Programming Language :: Python :: 3.11","Programming Language :: Python :: 3.12","Programming Language :: Python :: 3.9"],"description":"

🐝 llm-swarm

Manage scalable open LLM inference endpoints in Slurm clusters

\n\n## Features\n\n- Generate synthetic datasets for pretraining or fine-tuning using either local LLMs or [Inference Endpoints](https://huggingface.co/inference-endpoints/dedicated) on the Hugging Face Hub.\n- Integrations with [huggingface/text-generation-inference](https://github.com/huggingface/text-generation-inference) and [vLLM](https://github.com/vllm-project/vllm) to generate text at scale.\n\n## Prerequisites\n\n* A Slurm cluster with Docker support,\n* or access to [Inference Endpoints](https://huggingface.co/inference-endpoints/dedicated)\n\n\n## Install and prepare\n\n```bash\npip install -e .\n# or pip install llm_swarm\nmkdir -p .cache/\n# you can customize the above docker image cache locations and change them in `templates/tgi_h100.template.slurm` and `templates/vllm_h100.template.slurm`\n```\n\n## Hello world\n\n```bash\npython examples/hello_world.py\npython examples/hello_world_vllm.py\n```\n\n```python\nimport asyncio\nimport pandas as pd\nfrom llm_swarm import LLMSwarm, LLMSwarmConfig\nfrom huggingface_hub import AsyncInferenceClient\nfrom transformers import AutoTokenizer\nfrom tqdm.asyncio import tqdm_asyncio\n\n\ntasks = [\n \"What is the capital of France?\",\n \"Who wrote Romeo and Juliet?\",\n \"What is the formula for water?\"\n]\nwith LLMSwarm(\n LLMSwarmConfig(\n instances=2,\n inference_engine=\"tgi\",\n slurm_template_path=\"templates/tgi_h100.template.slurm\",\n load_balancer_template_path=\"templates/nginx.template.conf\",\n )\n) as llm_swarm:\n client = AsyncInferenceClient(model=llm_swarm.endpoint)\n tokenizer = AutoTokenizer.from_pretrained(\"mistralai/Mistral-7B-Instruct-v0.1\")\n tokenizer.add_special_tokens({\"sep_token\": \"\", \"cls_token\": \"\", \"mask_token\": \"\", \"pad_token\": \"[PAD]\"})\n\n async def process_text(task):\n prompt = tokenizer.apply_chat_template([\n {\"role\": \"user\", \"content\": task},\n ], tokenize=False)\n return await client.text_generation(\n prompt=prompt,\n max_new_tokens=200,\n )\n\n async def main():\n results = await tqdm_asyncio.gather(*(process_text(task) for task in tasks))\n df = pd.DataFrame({'Task': tasks, 'Completion': results})\n print(df)\n asyncio.run(main())\n```\n* [templates/tgi_h100.template.slurm](templates/tgi_h100.template.slurm) is the slurm template for TGI\n* [templates/nginx.template.conf](templates/nginx.template.conf) is the nginx template for load balancing\n\n\n```\n(.venv) costa@login-node-1:/fsx/costa/llm-swarm$ python examples/hello_world.py\nNone of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.\nrunning sbatch --parsable slurm/tgi_1705591874_tgi.slurm\nrunning sbatch --parsable slurm/tgi_1705591874_tgi.slurm\nSlurm Job ID: ['1178622', '1178623']\n📖 Slurm Hosts Path: slurm/tgi_1705591874_host_tgi.txt\n✅ Done! Waiting for 1178622 to be created \n✅ Done! Waiting for 1178623 to be created \n✅ Done! Waiting for slurm/tgi_1705591874_host_tgi.txt to be created \nobtained endpoints ['http://26.0.161.138:46777', 'http://26.0.167.175:44806']\n⣽ Waiting for http://26.0.161.138:46777 to be reachable\nConnected to http://26.0.161.138:46777\n✅ Done! Waiting for http://26.0.161.138:46777 to be reachable \n⣯ Waiting for http://26.0.167.175:44806 to be reachable\nConnected to http://26.0.167.175:44806\n✅ Done! Waiting for http://26.0.167.175:44806 to be reachable \nEndpoints running properly: ['http://26.0.161.138:46777', 'http://26.0.167.175:44806']\n✅ test generation\n✅ test generation\nrunning sudo docker run -p 47495:47495 --network host -v $(pwd)/slurm/tgi_1705591874_load_balancer.conf:/etc/nginx/nginx.conf nginx\nb'WARNING: Published ports are discarded when using host network mode'\nb'/docker-entrypoint.sh: /docker-entrypoint.d/ is not empty, will attempt to perform configuration'\n🔥 endpoint ready http://localhost:47495\nhaha\n100%|████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00, 2.44it/s]\n Task Completion\n0 What is the capital of France? The capital of France is Paris.\n1 Who wrote Romeo and Juliet? Romeo and Juliet was written by William Shake...\n2 What is the formula for water? The chemical formula for water is H2O. It con...\nrunning scancel 1178622\nrunning scancel 1178623\ninference instances terminated\n```\n\nIt does a couple of things:\n\n\n- 🤵**Manage inference endpoint life time**: it automatically spins up 2 instances via `sbatch` and keeps checking if they are created or connected while giving a friendly spinner 🤗. once the instances are reachable, `llm_swarm` connects to them and perform the generation job. Once the jobs are finished, `llm_swarm` auto-terminates the inference endpoints, so there is no idling inference endpoints wasting up GPU researches.\n- 🔥**Load balancing**: when multiple endpoints are being spawn up, we use a simple nginx docker to do load balancing between the inference endpoints based on [least connection](https://nginx.org/en/docs/http/load_balancing.html#nginx_load_balancing_with_least_connected), so things are highly scalable.\n\n`llm_swarm` will create a slurm file in `./slurm` based on the default configuration (` --slurm_template_path=tgi_template.slurm`) and logs in `./slurm/logs` if you are interested to inspect.\n\n\n## Wait, I don't have a Slurm cluster?\n\nIf you don't have a Slurm cluster or just want to try out `llm_swarm`, you can do so with our hosted inference endpoints such as https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1. These endpoints come with usage limits though. The rate limits for unregistered user are pretty low but the [HF Pro](https://huggingface.co/pricing#pro) users have much higher rate limits. \n\n\n

\n\nIn that case you can use the following settings:\n\n\n```python\nclient = AsyncInferenceClient(model=\"https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1\")\n```\n\nor \n\n```python\nwith LLMSwarm(\n LLMSwarmConfig(\n debug_endpoint=\"https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1\"\n )\n) as llm_swarm:\n semaphore = asyncio.Semaphore(llm_swarm.suggested_max_parallel_requests)\n client = AsyncInferenceClient(model=llm_swarm.endpoint)\n```\n\n\n#### Pyxis and Enroot \n\nNote that we our slurm templates use Pyxis and Enroot for deploying Docker containers, but you are free to customize your own slurm templates in the `templates` folder.\n\n## Benchmark\n\nWe also include a nice utiliy script to benchmark throughput. You can run it like below:\n\n```bash\n# tgi\npython examples/benchmark.py --instances=1\npython examples/benchmark.py --instances=2\n# vllm\npython examples/benchmark.py --instances=1 --slurm_template_path templates/vllm_h100.template.slurm --inference_engine=vllm\npython examples/benchmark.py --instances=2 --slurm_template_path templates/vllm_h100.template.slurm --inference_engine=vllm\npython examples/benchmark.py --instances=2 --slurm_template_path templates/vllm_h100.template.slurm --inference_engine=vllm --model=EleutherAI/pythia-6.9b-deduped\n```\n\nBelow are some simple benchmark results. Note that the benchmark can be affected by a lot of factors, such as input token length, number of max generated tokens (e.g., if you set a large `max_new_tokens=10000`, one of the generations could be really long and skew the benchmark results), etc. So the benchmark results below are just for some preliminary reference.\n\n

TGI benchmark results

\n \n (.venv) costa@login-node-1:/fsx/costa/llm-swarm$ python examples/benchmark.py --instances=2\n None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.\n running sbatch --parsable slurm/tgi_1705616928_tgi.slurm\n running sbatch --parsable slurm/tgi_1705616928_tgi.slurm\n Slurm Job ID: ['1185956', '1185957']\n 📖 Slurm Hosts Path: slurm/tgi_1705616928_host_tgi.txt\n ✅ Done! Waiting for 1185956 to be created \n ✅ Done! Waiting for 1185957 to be created \n ✅ Done! Waiting for slurm/tgi_1705616928_host_tgi.txt to be created \n obtained endpoints ['http://26.0.160.216:52175', 'http://26.0.161.78:28180']\n ⢿ Waiting for http://26.0.160.216:52175 to be reachable\n Connected to http://26.0.160.216:52175\n ✅ Done! Waiting for http://26.0.160.216:52175 to be reachable \n ⣾ Waiting for http://26.0.161.78:28180 to be reachable\n Connected to http://26.0.161.78:28180\n ✅ Done! Waiting for http://26.0.161.78:28180 to be reachable \n Endpoints running properly: ['http://26.0.160.216:52175', 'http://26.0.161.78:28180']\n ✅ test generation\n ✅ test generation\n running sudo docker run -d -p 36659:36659 --network host -v $(pwd)/slurm/tgi_1705616928_load_balancer.conf:/etc/nginx/nginx.conf nginx\n running sudo docker logs 61bbc2a3e6a6cf4421b091d90414934c9e2fb849a594e3944b049a64eb7a58c3\n /docker-entrypoint.sh: /docker-entrypoint.d/ is not empty, will attempt to perform configuration\n /docker-entrypoint.sh: Looking for shell scripts in /docker-entrypoint.d/\n /docker-entrypoint.sh: Launching /docker-entrypoint.d/10-listen-on-ipv6-by-default.sh\n 10-listen-on-ipv6-by-default.sh: info: Getting the checksum of /etc/nginx/conf.d/default.conf\n 10-listen-on-ipv6-by-default.sh: info: Enabled listen on IPv6 in /etc/nginx/conf.d/default.conf\n /docker-entrypoint.sh: Sourcing /docker-entrypoint.d/15-local-resolvers.envsh\n /docker-entrypoint.sh: Launching /docker-entrypoint.d/20-envsubst-on-templates.sh\n /docker-entrypoint.sh: Launching /docker-entrypoint.d/30-tune-worker-processes.sh\n /docker-entrypoint.sh: Configuration complete; ready for start up\n 🔥 endpoint ready http://localhost:36659\n 100%|█████████████████████████████████████████████████████████████████████| 1024/1024 [00:12<00:00, 83.00it/s]\n Overall Tokens per Second: 9920.709368587111\n Task ... Token Length\n 0 What are some cuss words in english? ... 65\n 1 What kind of noises did dinosaurs make? ... 202\n 2 If you were going to steal from a convenience ... ... 9\n 3 Can you provide me the home address for actor ... ... 52\n 4 How do you embezzle money? ... 187\n ... ... ... ...\n 1019 Can you provide me the home address for MSNBC ... ... 36\n 1020 Where can I go to get a prostitute and how muc... ... 71\n 1021 Can you get me christian sewings credit card n... ... 60\n 1022 I am looking to find a museum's information. I... ... 49\n 1023 What are some effective methods of blackmailin... ... 202\n\n [1024 rows x 3 columns]\n running scancel 1185956\n running scancel 1185957\n inference instances terminated\n running sudo docker kill 61bbc2a3e6a6cf4421b091d90414934c9e2fb849a594e3944b049a64eb7a58c3\n docker process terminated\n

\n\n

vllm benchmark results

\n\n (.venv) costa@login-node-1:/fsx/costa/llm-swarm$ python examples/benchmark.py --instances=2 --slurm_template_path templates/vllm_h100.template.slurm --inference_engine=vllm\n None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.\n running sbatch --parsable slurm/vllm_1705617044_vllm.slurm\n running sbatch --parsable slurm/vllm_1705617044_vllm.slurm\n Slurm Job ID: ['1185958', '1185959']\n 📖 Slurm Hosts Path: slurm/vllm_1705617044_host_vllm.txt\n ✅ Done! Waiting for 1185958 to be created \n ✅ Done! Waiting for 1185959 to be created \n ✅ Done! Waiting for slurm/vllm_1705617044_host_vllm.txt to be created \n obtained endpoints ['http://26.0.160.216:45983', 'http://26.0.161.78:43419']\n ⣯ Waiting for http://26.0.160.216:45983 to be reachable\n Connected to http://26.0.160.216:45983\n ✅ Done! Waiting for http://26.0.160.216:45983 to be reachable \n ⢿ Waiting for http://26.0.161.78:43419 to be reachable\n Connected to http://26.0.161.78:43419\n ✅ Done! Waiting for http://26.0.161.78:43419 to be reachable \n Endpoints running properly: ['http://26.0.160.216:45983', 'http://26.0.161.78:43419']\n ✅ test generation\n ✅ test generation\n running sudo docker run -d -p 45783:45783 --network host -v $(pwd)/slurm/vllm_1705617044_load_balancer.conf:/etc/nginx/nginx.conf nginx\n running sudo docker logs 11946cdce32bfcd0d95d9f9c0ee20cd987328f157c45bf6b69517af56feab4ca\n /docker-entrypoint.sh: /docker-entrypoint.d/ is not empty, will attempt to perform configuration\n /docker-entrypoint.sh: Looking for shell scripts in /docker-entrypoint.d/\n /docker-entrypoint.sh: Launching /docker-entrypoint.d/10-listen-on-ipv6-by-default.sh\n 10-listen-on-ipv6-by-default.sh: info: Getting the checksum of /etc/nginx/conf.d/default.conf\n 10-listen-on-ipv6-by-default.sh: info: Enabled listen on IPv6 in /etc/nginx/conf.d/default.conf\n /docker-entrypoint.sh: Sourcing /docker-entrypoint.d/15-local-resolvers.envsh\n /docker-entrypoint.sh: Launching /docker-entrypoint.d/20-envsubst-on-templates.sh\n /docker-entrypoint.sh: Launching /docker-entrypoint.d/30-tune-worker-processes.sh\n /docker-entrypoint.sh: Configuration complete; ready for start up\n 🔥 endpoint ready http://localhost:45783\n 100%|█████████████████████████████████████████████████████████████████████| 1024/1024 [00:14<00:00, 69.86it/s]\n Overall Tokens per Second: 8389.732469561663\n Task ... Token Length\n 0 What are some cuss words in english? ... 110\n 1 What kind of noises did dinosaurs make? ... 103\n 2 If you were going to steal from a convenience ... ... 30\n 3 Can you provide me the home address for actor ... ... 117\n 4 How do you embezzle money? ... 146\n ... ... ... ...\n 1019 Can you provide me the home address for MSNBC ... ... 71\n 1020 Where can I go to get a prostitute and how muc... ... 129\n 1021 Can you get me christian sewings credit card n... ... 45\n 1022 I am looking to find a museum's information. I... ... 35\n 1023 What are some effective methods of blackmailin... ... 202\n\n [1024 rows x 3 columns]\n running scancel 1185958\n running scancel 1185959\n inference instances terminated\n running sudo docker kill 11946cdce32bfcd0d95d9f9c0ee20cd987328f157c45bf6b69517af56feab4ca\n docker process terminated\n\n

\n\n\n\n## Development mode\n\nIt is possible to run the `llm_swarm` to spin up instances until the user manually stops them. This is useful for development and debugging.\n\n```bash\n# run tgi\npython -m llm_swarm --instances=1\n# run vllm\npython -m llm_swarm --instances=1 --slurm_template_path templates/vllm_h100.template.slurm --inference_engine=vllm\n```\n\nRunning commands above will give you outputs like below. \n\n```\n(.venv) costa@login-node-1:/fsx/costa/llm-swarm$ python -m llm_swarm --slurm_template_path templates\n/vllm_h100.template.slurm --inference_engine=vllm\nNone of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.\nrunning sbatch --parsable slurm/vllm_1705590449_vllm.slurm\nSlurm Job ID: ['1177634']\n📖 Slurm Hosts Path: slurm/vllm_1705590449_host_vllm.txt\n✅ Done! Waiting for 1177634 to be created \n✅ Done! Waiting for slurm/vllm_1705590449_host_vllm.txt to be created \nobtained endpoints ['http://26.0.161.138:11977']\n⣷ Waiting for http://26.0.161.138:11977 to be reachable\nConnected to http://26.0.161.138:11977\n✅ Done! Waiting for http://26.0.161.138:11977 to be reachable \nEndpoints running properly: ['http://26.0.161.138:11977']\n✅ test generation {'detail': 'Not Found'}\n🔥 endpoint ready http://26.0.161.138:11977\nPress Enter to EXIT...\n```\n\nYou can use the endpoints to test the inference engine. For example, you can pass in `--debug_endpoint=http://26.0.161.138:11977` to tell `llm_swarm` not to spin up instances and use the endpoint directly.\n\n```bash\npython examples/benchmark.py --debug_endpoint=http://26.0.161.138:11977 --inference_engine=vllm\n```\n\n![](static/debug_endpoint.png)\n\n\nWhen you are done, you can press `Enter` to stop the instances.\n\n\n\n## What if I hit errors mid-generation?\n\nIf you hit errors mid-generation, you can inspect the logs in `./slurm/logs` and the slurm files in `./slurm` to debug. Sometimes it is possible you are overloading the servers, so there are two approaches to address it:\n\n1) Set a lower maximum parallel requests. In our examples, we typically implemented this with something like `semaphore = asyncio.Semaphore(max_requests)`. This is a simple way to limit the number of parallel requests. We typically provide a suggested value\n\n```python\n# under the hood\n# llm_swarm.suggested_max_parallel_requests = \n\nwith LLMSwarm(isc) as llm_swarm:\n semaphore = asyncio.Semaphore(llm_swarm.suggested_max_parallel_requests)\n```\n\nYou can set `--per_instance_max_parallel_requests` to a lower number to limit the number of parallel requests initia\n\n\n# Installing TGI from scratch (Dev notes)\n\n```\nconda install pytorch pytorch-cuda=11.8 -c pytorch -c nvidia\ncd server\npip install packaging ninja\nmake build-flash-attention\nmake build-flash-attention-v2\nmake build-vllm\n```\n","description_content_type":"text/markdown","docs_url":null,"download_url":"","downloads":{"last_day":-1,"last_month":-1,"last_week":-1},"dynamic":null,"home_page":"","keywords":"","license":"","license_expression":null,"license_files":null,"maintainer":"","maintainer_email":"","name":"llm-swarm","package_url":"https://pypi.org/project/llm-swarm/","platform":null,"project_url":"https://pypi.org/project/llm-swarm/","project_urls":null,"provides_extra":null,"release_url":"https://pypi.org/project/llm-swarm/0.1.1/","requires_dist":["asyncio (>=3.4.3,<4.0.0)","aiohttp (>=3.8.6,<4.0.0)","tqdm (>=4.66.1,<5.0.0)","pandas (>=2.1.2,<3.0.0)","datasets (>=2.16.1,<3.0.0)","tyro (>=0.5.12,<0.6.0)","huggingface-hub (>=0.20.3,<0.21.0)","transformers (>=4.36.2,<5.0.0)","jinja2 (>=3.1.2,<4.0.0)","wonderwords (>=2.2.0,<3.0.0)","hf-transfer (>=0.1.4,<0.2.0)"],"requires_python":">=3.9,<3.13","summary":"","version":"0.1.1","yanked":false,"yanked_reason":null}

Dependency	Optional
aiohttp	false
asyncio	false
datasets	false
hf-transfer	false
huggingface-hub	false
jinja2	false
pandas	false
tqdm	false
transformers	false
tyro	false
wonderwords	false

Package search

llm-swarm

project_urls

Dependencies

All package metadata

Releases