[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

applications/ColossalChat/coati/distributed/comm.py

@@ -1,5 +1,6 @@
-from typing import Any, Dict
 import copy
+from typing import Any, Dict
+
 import ray
 import ray.util.collective as cc
 import torch
@@ -31,6 +32,7 @@ def ray_broadcast_object(obj: Any, src: int = 0, device=None, group_name: str =
             obj = c10d._tensor_to_object(obj, size_tensor.item())
     return obj
 
+
 def ray_broadcast_tensor_dict(
     tensor_dict: Dict[str, torch.Tensor],
     src: int = 0,

applications/ColossalChat/coati/distributed/zero_bubble/consumer.py

@@ -7,17 +7,16 @@
 import ray.util.collective as cc
 import torch
 import torch.distributed as dist
+from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
 from coati.distributed.profiling_utils import CustomProfiler
+from coati.distributed.utils import bind_batch, post_recv, unbind_batch
 from tqdm import tqdm
 
 from colossalai.booster import Booster
 from colossalai.booster.plugin import HybridParallelPlugin
 from colossalai.initialize import launch
 from colossalai.utils import get_current_device
 
-from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
-from coati.distributed.utils import bind_batch, post_recv, unbind_batch
-
 
 class BaseConsumer:
     def __init__(

applications/ColossalChat/coati/distributed/zero_bubble/distributor.py

@@ -3,12 +3,11 @@
 import ray
 import ray.util.collective as cc
 import torch
+from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
 from coati.distributed.profiling_utils import CustomProfiler
 
 from colossalai.utils import get_current_device
 
-from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
-
 
 @ray.remote
 class Distributor:

applications/ColossalChat/coati/distributed/zero_bubble/grpo_consumer.py

@@ -5,9 +5,9 @@
 import torch
 import wandb
 from coati.distributed.comm import SharedVariableActor
-from coati.distributed.zero_bubble.consumer import BaseConsumer
 from coati.distributed.loss import PolicyLoss
 from coati.distributed.utils import memory_efficient_logprob
+from coati.distributed.zero_bubble.consumer import BaseConsumer
 from coati.trainer.utils import all_reduce_mean, all_reduce_sum
 from transformers import AutoModelForCausalLM, AutoTokenizer
 

applications/ColossalChat/coati/distributed/zero_bubble/producer.py

@@ -11,20 +11,19 @@
 import tqdm
 import wandb
 from coati.dataset.loader import RawConversationDataset, collate_fn_grpo
+from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
+from coati.distributed.inference_backend import BACKEND_MAP
 from coati.distributed.profiling_utils import CustomProfiler
 from coati.distributed.reward.reward_fn import boxed_math_reward_fn, code_reward_fn, math_reward_fn
 from coati.distributed.reward.verifiable_reward import VerifiableReward
+from coati.distributed.utils import pre_send, safe_append_to_jsonl_file
 from ray.util.collective import allreduce
 from ray.util.collective.types import ReduceOp
 from torch.utils.data import DataLoader, DistributedSampler
 from transformers import AutoTokenizer
 
 from colossalai.utils import get_current_device
 
-from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
-from coati.distributed.inference_backend import BACKEND_MAP
-from coati.distributed.utils import pre_send, safe_append_to_jsonl_file
-
 try:
     from vllm import SamplingParams
 except ImportError: