From dc71574a630c307b83dc2a389d6520375211d482 Mon Sep 17 00:00:00 2001 From: Alexander Borzunov Date: Mon, 28 Nov 2022 03:44:41 +0400 Subject: [PATCH] Use public swarm by default (#92) This PR makes servers and clients use public swarm's bootstrap peers if no other initial peers are specified. If you'd like a server to start a new swarm, provide the `--new_swarm` CLI argument. --- .github/workflows/run-tests.yaml | 20 ++++++++++---------- cli/run_server.py | 19 ++++++++++++++----- src/client/remote_model.py | 5 +++-- src/constants.py | 8 ++++++++ src/server/block_selection.py | 3 +++ src/server/server.py | 12 ++++++++---- 6 files changed, 46 insertions(+), 21 deletions(-) create mode 100644 src/constants.py diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml index dcbcf9f..4dd4e82 100644 --- a/.github/workflows/run-tests.yaml +++ b/.github/workflows/run-tests.yaml @@ -72,15 +72,15 @@ jobs: export REF_NAME=bigscience/bloom-560m python -m cli.run_server --converted_model_name_or_path $MODEL_NAME --block_indices 0:12 \ - --identity tests/test.id --host_maddrs /ip4/127.0.0.1/tcp/31337 --throughput 1 \ + --new_swarm --identity tests/test.id --host_maddrs /ip4/127.0.0.1/tcp/31337 --throughput 1 \ --torch_dtype float32 --compression NONE --attn_cache_size 0.2GiB &> server1.log & SERVER1_PID=$! - + sleep 5 # wait for the first server to initialize DHT - + export INITIAL_PEERS=/ip4/127.0.0.1/tcp/31337/p2p/QmS9KwZptnVdB9FFV7uGgaTq4sEKBwcYeKZDfSpyKDUd1g # ^-- server 1 multiaddr is determined by --identity and --host_maddrs - + python -m cli.run_server --converted_model_name_or_path $MODEL_NAME --block_indices 12:22 \ --initial_peers $INITIAL_PEERS --throughput 1 --torch_dtype float32 &> server2.log & SERVER2_PID=$! @@ -94,20 +94,20 @@ jobs: python -m cli.run_server --converted_model_name_or_path $MODEL_NAME --block_indices 4:16 \ --torch_dtype float32 --initial_peers $INITIAL_PEERS --throughput 1 &> server4.log & SERVER4_PID=$! - + python -m cli.run_server --converted_model_name_or_path $MODEL_NAME --num_blocks 3 \ --initial_peers $INITIAL_PEERS --throughput 1 --torch_dtype float32 &> server5.log & SERVER5_PID=$! - + tail -n 100 -f server*.log & LOGGER_PID=$! sleep 30 # wait for servers to download layers - + kill -0 $SERVER1_PID $SERVER2_PID $SERVER3_PID $SERVER4_PID $SERVER5_PID # ensure all servers survived init - + PYTHONPATH=. pytest tests --durations=0 --durations-min=1.0 -v - + kill -0 $SERVER1_PID $SERVER2_PID $SERVER3_PID $SERVER4_PID $SERVER5_PID # ensure all servers survived tests - + kill -s SIGINT $SERVER1_PID $SERVER2_PID $SERVER3_PID $SERVER4_PID $SERVER5_PID $LOGGER_PID echo "Done!" diff --git a/cli/run_server.py b/cli/run_server.py index fcef351..6ab7aa2 100644 --- a/cli/run_server.py +++ b/cli/run_server.py @@ -6,6 +6,7 @@ from hivemind.utils.limits import increase_file_limit from hivemind.utils.logging import get_logger, use_hivemind_log_handler from humanfriendly import parse_size +from src.constants import PUBLIC_INITIAL_PEERS from src.server.server import Server use_hivemind_log_handler("in_root_logger") @@ -27,10 +28,10 @@ def main(): parser.add_argument('--block_indices', type=str, default=None, help="Specific block indices to serve") parser.add_argument('--prefix', type=str, default=None, help="Announce all blocks with this prefix. By default," "use the same name as in the converted model.") - parser.add_argument('--host_maddrs', nargs='+', default=['/ip4/0.0.0.0/tcp/0'], required=False, - help='Multiaddrs to listen for external connections from other p2p instances; default: all IPv4 and TCP: /ip4/0.0.0.0/tcp/0') + parser.add_argument('--host_maddrs', nargs='+', default=['/ip4/0.0.0.0/tcp/0', '/ip6/::/tcp/0'], required=False, + help='Multiaddrs to listen for external connections from other peers. Default: all IPv4/IPv6 interfaces, a random free TCP port') parser.add_argument('--announce_maddrs', nargs='+', default=None, required=False, - help='Visible multiaddrs the host announces for external connections from other p2p instances') + help='Visible multiaddrs the host announces for external connections from other peers') parser.add_argument('--compression', type=str, default='NONE', required=False, help='Tensor compression communication') @@ -71,8 +72,13 @@ def main(): help='Server will report blocks to DHT once in this many seconds') parser.add_argument('--expiration', type=float, required=False, default=None, help='DHT entries will expire after this many seconds') - parser.add_argument('--initial_peers', type=str, nargs='*', required=False, default=[], - help='multiaddrs of one or more active DHT peers (if you want to join an existing DHT)') + + group = parser.add_mutually_exclusive_group() + group.add_argument('--initial_peers', type=str, nargs='*', required=False, default=PUBLIC_INITIAL_PEERS, + help='Multiaddrs of one or more DHT peers from the target swarm. Default: connects to the public swarm') + group.add_argument('--new_swarm', action='store_true', + help='Start a new private swarm (i.e., do not connect to any initial peers)') + parser.add_argument('--increase_file_limit', action='store_true', help='On *nix, this will increase the max number of processes ' 'a server can spawn before hitting "Too many open files"; Use at your own risk.') @@ -112,6 +118,9 @@ def main(): attn_cache_size, (int, type(None)) ), "unrecognized value for attention_cache_bytes, examples: 1.5GB or 1500MB or 1572864000 (bytes)" + if args.pop("new_swarm"): + args["initial_peers"] = [] + use_auth_token = args.pop("use_auth_token") args["use_auth_token"] = True if use_auth_token in ("True", "true", "") else use_auth_token diff --git a/src/client/remote_model.py b/src/client/remote_model.py index 158b7a1..b846d6e 100644 --- a/src/client/remote_model.py +++ b/src/client/remote_model.py @@ -1,5 +1,5 @@ # this code is in active development, interfaces may change -from typing import Optional, Tuple +from typing import List, Optional import hivemind import torch @@ -17,6 +17,7 @@ from src.bloom.model import ( ) from src.client.remote_generation import RemoteGenerationMixin from src.client.remote_sequential import RemoteSequential +from src.constants import PUBLIC_INITIAL_PEERS from src.utils.misc import DUMMY use_hivemind_log_handler("in_root_logger") @@ -29,7 +30,7 @@ class DistributedBloomConfig(BloomConfig): To create a distributed model, one must provide dht_prefix and either initial_peers or dht. """ - initial_peers: Tuple[str, ...] = () # a list of initial peers for hivemind DHT + initial_peers: List[str] = PUBLIC_INITIAL_PEERS # a list of initial peers for hivemind DHT dht_prefix: str # a prefix for all dht keys that correspond to this model (usually equal to model name) dht: Optional[hivemind.DHT] = None # a running DHT instance, e.g. when using the same DHT for multiple models chunk_size_for_efficient_fp16_on_cpu: int = 10000 # a chunk size for a LM head for efficient half-precision on CPU diff --git a/src/constants.py b/src/constants.py new file mode 100644 index 0000000..9f7197d --- /dev/null +++ b/src/constants.py @@ -0,0 +1,8 @@ +PUBLIC_INITIAL_PEERS = [ + "/dns/bootstrap1.petals.ml/tcp/31337/p2p/QmedTaZXmULqwspJXz44SsPZyTNKxhnnFvYRajfH7MGhCY", + "/dns6/bootstrap1.petals.ml/tcp/31337/p2p/QmedTaZXmULqwspJXz44SsPZyTNKxhnnFvYRajfH7MGhCY", + "/dns/bootstrap2.petals.ml/tcp/31338/p2p/QmQGTqmM7NKjV6ggU1ZCap8zWiyKR89RViDXiqehSiCpY5", + "/dns6/bootstrap2.petals.ml/tcp/31338/p2p/QmQGTqmM7NKjV6ggU1ZCap8zWiyKR89RViDXiqehSiCpY5", + "/dns/bootstrap3.petals.ml/tcp/31339/p2p/QmX82nfE57CSkNgyEC7pPMPBzjcFLLJXdHhvp1AXKVPvJD", + "/dns6/bootstrap3.petals.ml/tcp/31339/p2p/QmX82nfE57CSkNgyEC7pPMPBzjcFLLJXdHhvp1AXKVPvJD", +] diff --git a/src/server/block_selection.py b/src/server/block_selection.py index 8710025..c60b2ee 100644 --- a/src/server/block_selection.py +++ b/src/server/block_selection.py @@ -106,6 +106,9 @@ def should_choose_other_blocks( throughputs[span.start : span.end] += span.throughput new_throughput = throughputs.min() + if new_throughput < initial_throughput or new_throughput < eps: + return False + actual_quality = initial_throughput / new_throughput logger.info(f"Swarm balance quality: {actual_quality * 100:.1f}%") diff --git a/src/server/server.py b/src/server/server.py index 174762d..d065802 100644 --- a/src/server/server.py +++ b/src/server/server.py @@ -5,7 +5,7 @@ import multiprocessing as mp import random import threading import time -from typing import Dict, List, Optional, Sequence, Union +from typing import Dict, List, Optional, Union import numpy as np import psutil @@ -18,6 +18,7 @@ from hivemind.utils.logging import get_logger, use_hivemind_log_handler from src import BloomConfig, declare_active_modules from src.bloom.from_pretrained import DTYPE_MAP, load_pretrained_block +from src.constants import PUBLIC_INITIAL_PEERS from src.data_structures import CHAIN_DELIMITER, UID_DELIMITER, ServerState from src.dht_utils import get_remote_module_infos from src.server import block_selection @@ -39,6 +40,8 @@ class Server(threading.Thread): def __init__( self, + *, + initial_peers: List[str], prefix: Optional[str], converted_model_name_or_path: str, throughput: Union[float, str], @@ -53,7 +56,6 @@ class Server(threading.Thread): cache_dir: Optional[str] = None, attn_cache_size: Optional[int] = None, device: Optional[Union[str, torch.device]] = None, - initial_peers: Sequence[str] = (), compression=CompressionType.NONE, stats_report_interval: Optional[int] = None, custom_module_path=None, @@ -66,7 +68,6 @@ class Server(threading.Thread): mean_block_selection_delay: float = 0.5, use_auth_token: Optional[str] = None, load_in_8bit: bool = False, - *, start: bool, **kwargs, ): @@ -104,7 +105,10 @@ class Server(threading.Thread): self.dht = DHT(initial_peers=initial_peers, start=True, **kwargs) visible_maddrs_str = [str(a) for a in self.dht.get_visible_maddrs()] - logger.info(f"Running DHT node on {visible_maddrs_str}, initial peers = {initial_peers}") + if initial_peers == PUBLIC_INITIAL_PEERS: + logger.info("Connecting to the public Petals swarm") + else: + logger.info(f"Running DHT node on {visible_maddrs_str}, initial peers = {initial_peers}") device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.device = device