add 8bit quant all reduce support

Author: samsja

open_diloco/train_fsdp.py

@@ -97,7 +97,7 @@ class HvConfig(BaseConfig):
     announce_maddrs: list[str] | None = None
     matchmaking_time: float | None = None
     averaging_timeout: float | None = None
-    hivemind_compression: Literal["none", "fp16", "scaled-fp16"] = "none"
+    hivemind_compression: Literal["none", "fp16", "scaled-fp16", "uniform8bit", "quantile8bit"] = "none"
     all_reduce_strategy: AllReduceStrategy = AllReduceStrategy.WAIT_FOR_ALL
     timeout_waiting_for_peers: float | None = None
     skip_load_from_peers: bool = False

open_diloco/utils.py

@@ -108,6 +108,16 @@ def get_compression_kwargs(hivemind_compression: str) -> dict:
 
         ret_kwargs["grad_compression"] = NoCompression()
         ret_kwargs["state_averaging_compression"] = NoCompression()
+    elif hivemind_compression == "uniform8bit":
+        from hivemind import Uniform8BitQuantization
+
+        ret_kwargs["grad_compression"] = Uniform8BitQuantization()
+        ret_kwargs["state_averaging_compression"] = Uniform8BitQuantization()
+    elif hivemind_compression == "quantile8bit":
+        from hivemind import Quantile8BitQuantization
+
+        ret_kwargs["grad_compression"] = Quantile8BitQuantization()
+        ret_kwargs["state_averaging_compression"] = Quantile8BitQuantization()
     else:
         raise ValueError(
             f"Invalid hivemind_compression: {hivemind_compression}. Please choose 'none', 'fp16', or 'scaled-fp16'."