refactor: 优化参数传递，清理导入样式

astrai/parallel/__init__.py

@@ -1,14 +1,13 @@
+from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
 from astrai.parallel.setup import (
-    get_world_size,
-    get_rank,
     get_current_device,
+    get_rank,
+    get_world_size,
     only_on_rank,
     setup_parallel,
     spawn_parallel_fn,
 )
 
-from astrai.parallel.module import RowParallelLinear, ColumnParallelLinear
-
 __all__ = [
     "get_world_size",
     "get_rank",

astrai/parallel/module.py

@@ -1,10 +1,10 @@
+from typing import Dict
+
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
-import torch.distributed as dist
-
 from torch import Tensor
-from typing import Dict
 
 
 class ParallelModel(nn.Module):

astrai/parallel/setup.py

@@ -1,12 +1,12 @@
 import os
+from contextlib import contextmanager
+from functools import wraps
+from typing import Callable, List, Optional
+
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from functools import wraps
-from contextlib import contextmanager
-from typing import Callable, List, Optional
-
 
 def get_current_device():
     return os.environ["LOCAL_DEVICE"]