Browse Source

Update hivemind to 1.1.8, enable efficient bfloat16 encoding (#311)

This PR:

1. Updates hivemind to 1.1.8 (includes https://github.com/learning-at-home/hivemind/pull/565)
2. Enables efficient bfloat16 serialization by default (`USE_LEGACY_BFLOAT16 = False`)
3. Removes logging code that was included to hivemind in https://github.com/learning-at-home/hivemind/pull/542
Alexander Borzunov 2 năm trước cách đây
mục cha
commit
0a313bf6c5
3 tập tin đã thay đổi với 12 bổ sung19 xóa
  1. 1 1
      setup.cfg
  2. 11 0
      src/petals/__init__.py
  3. 0 18
      src/petals/utils/logging.py

+ 1 - 1
setup.cfg

@@ -37,7 +37,7 @@ install_requires =
     huggingface-hub>=0.11.1,<1.0.0
     transformers>=4.25.1,<5.0.0
     speedtest-cli==2.1.3
-    hivemind==1.1.7
+    hivemind==1.1.8
     tensor_parallel==1.0.23
     humanfriendly
     async-timeout>=4.0.2

+ 11 - 0
src/petals/__init__.py

@@ -1,6 +1,17 @@
+import os
+
+import hivemind
+
 from petals.client import *
 from petals.utils.logging import initialize_logs as _initialize_logs
 
 __version__ = "1.1.4"
 
+
+def _override_bfloat16_mode_default():
+    if os.getenv("USE_LEGACY_BFLOAT16") is None:
+        hivemind.compression.base.USE_LEGACY_BFLOAT16 = False
+
+
 _initialize_logs()
+_override_bfloat16_mode_default()

+ 0 - 18
src/petals/utils/logging.py

@@ -4,16 +4,6 @@ import os
 from hivemind.utils import logging as hm_logging
 
 
-def in_jupyter() -> bool:
-    """Check if the code is run in Jupyter or Colab"""
-
-    try:
-        __IPYTHON__
-        return True
-    except NameError:
-        return False
-
-
 def initialize_logs():
     """Initialize Petals logging tweaks. This function is called when you import the `petals` module."""
 
@@ -21,14 +11,6 @@ def initialize_logs():
     if os.getenv("PETALS_LOGGING", "True").lower() in ("false", "0"):
         return
 
-    if in_jupyter():
-        os.environ["HIVEMIND_COLORS"] = "True"
-    importlib.reload(hm_logging)
-
-    # Remove log handlers from previous import of hivemind.utils.logging and extra handlers on Colab
-    hm_logging.get_logger().handlers.clear()
-    hm_logging.get_logger("hivemind").handlers.clear()
-
     hm_logging.use_hivemind_log_handler("in_root_logger")
 
     # We suppress asyncio error logs by default since they are mostly not relevant for the end user,