[CLEANUP]

Author: Kye

exa/structs/__init__.py

@@ -1,10 +1,15 @@
 from exa.structs.parallelize_models_gpus import (
     prepare_model_for_ddp_inference,
+    setup_distributed_environment,
+    initialize_process_group,
 )
 from exa.structs.model_thread_router import ModelThreadWorker, Router
 
 __all__ = [
     "ModelThreadWorker",
     "Router",
     "prepare_model_for_ddp_inference",
+    "setup_distributed_environment",
+    "initialize_process_group",
+    "prepare_model_for_ddp_inference",
 ]

exa/utils/__init__.py

@@ -33,4 +33,4 @@
     "fused_all_gather_v1",
     "fused_all_gather_v2",
     "calculate_workers",
-]
+]

exa/utils/gpu_ops.py

@@ -6,8 +6,6 @@
     memory_reserved,
 )
 
-#
-
 
 def get_world_size_rank():
     """