[misc] update pre-commit and run all files (#4752)

* [misc] update pre-commit

* [misc] run pre-commit

* [misc] remove useless configuration files

* [misc] ignore cuda for clang-format

tests/test_infer/_utils.py

@@ -1,20 +1,6 @@
 import copy
 
-import torch
-import torch.distributed as dist
-from torch import Tensor
-from torch import distributed as dist
-from torch.distributed import ProcessGroup
-from torch.nn import Module
-from torch.optim import Adam, Optimizer
-
-from colossalai.booster import Booster
-from colossalai.booster.plugin import HybridParallelPlugin
-from colossalai.booster.plugin.hybrid_parallel_plugin import HybridParallelModule
 from colossalai.shardformer import ShardConfig, ShardFormer
-from colossalai.shardformer._utils import getattr_
-from colossalai.shardformer.policies.auto_policy import Policy
-from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor
 
 
 def build_model(
@@ -28,11 +14,13 @@ def build_model(
     org_model = model_fn()
 
     # shard model
-    shard_config = ShardConfig(enable_fused_normalization=enable_fused_normalization,
-                               enable_tensor_parallelism=enable_tensor_parallelism,
-                               enable_flash_attention=enable_flash_attention,
-                               enable_jit_fused=enable_jit_fused,
-                               inference_only=True)
+    shard_config = ShardConfig(
+        enable_fused_normalization=enable_fused_normalization,
+        enable_tensor_parallelism=enable_tensor_parallelism,
+        enable_flash_attention=enable_flash_attention,
+        enable_jit_fused=enable_jit_fused,
+        inference_only=True,
+    )
     model_copy = copy.deepcopy(org_model)
     shard_former = ShardFormer(shard_config=shard_config)
     sharded_model, shared_params = shard_former.optimize(model_copy)

tests/test_infer/test_bloom_infer.py

@@ -1,5 +1,3 @@
-import os
-
 import pytest
 import torch
 from packaging import version
@@ -16,22 +14,27 @@ MAX_BATCH_SIZE = 4
 MAX_INPUT_LEN = 16
 MAX_OUTPUT_LEN = 32
 
-CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse('11.5')
+CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.5")
 
 
-@parameterize('test_config', [{
-    'tp_size': TP_SIZE,
-}])
+@parameterize(
+    "test_config",
+    [
+        {
+            "tp_size": TP_SIZE,
+        }
+    ],
+)
 def run(test_config):
-
-    sub_model_zoo = model_zoo.get_sub_registry('transformers_bloom_for_causal_lm')
+    sub_model_zoo = model_zoo.get_sub_registry("transformers_bloom_for_causal_lm")
     for name, (model_fn, data_gen_fn, _, _, _) in sub_model_zoo.items():
         orig_model = model_fn()
         orig_model = orig_model.half()
         data = data_gen_fn()
 
-        shard_config = ShardConfig(enable_tensor_parallelism=True if test_config['tp_size'] > 1 else False,
-                                   inference_only=True)
+        shard_config = ShardConfig(
+            enable_tensor_parallelism=True if test_config["tp_size"] > 1 else False, inference_only=True
+        )
         infer_engine = TPInferEngine(orig_model, shard_config, MAX_BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
 
         generate_kwargs = dict(do_sample=False)
@@ -42,7 +45,7 @@ def run(test_config):
 
 def check_bloom(rank, world_size, port):
     disable_existing_loggers()
-    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
     run()
 
 
@@ -54,5 +57,5 @@ def test_bloom_infer():
     spawn(check_bloom, TP_SIZE)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     test_bloom_infer()

tests/test_infer/test_infer_engine.py

@@ -2,14 +2,12 @@ from itertools import accumulate
 
 import pytest
 import torch
-import torch.nn as nn
 from packaging import version
-from transformers import BloomConfig, BloomForCausalLM, LlamaConfig, LlamaForCausalLM
+from transformers import BloomConfig, BloomForCausalLM
 from transformers.tokenization_utils_base import BatchEncoding
 
 import colossalai
 from colossalai.inference.tensor_parallel import TPInferEngine
-from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
 from colossalai.logging import disable_existing_loggers
 from colossalai.shardformer import ShardConfig
 from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
@@ -19,12 +17,17 @@ MAX_BATCH_SIZE = 4
 MAX_INPUT_LEN = 16
 MAX_OUTPUT_LEN = 8
 
-CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse('11.5')
+CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.5")
 
 
-@parameterize('test_config', [{
-    'tp_size': TP_SIZE,
-}])
+@parameterize(
+    "test_config",
+    [
+        {
+            "tp_size": TP_SIZE,
+        }
+    ],
+)
 def run(test_config):
     model_config = BloomConfig(num_hidden_layers=4, hidden_size=128, intermediate_size=256, num_attention_heads=4)
     model = BloomForCausalLM(model_config)
@@ -32,8 +35,9 @@ def run(test_config):
     model.to(torch.cuda.current_device())
 
     # 1. check TPInferEngine init and model optimization
-    shard_config = ShardConfig(enable_tensor_parallelism=True if test_config['tp_size'] > 1 else False,
-                               inference_only=True)
+    shard_config = ShardConfig(
+        enable_tensor_parallelism=True if test_config["tp_size"] > 1 else False, inference_only=True
+    )
     infer_engine = TPInferEngine(model, shard_config, MAX_BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
 
     assert infer_engine.cache_manager is not None
@@ -41,13 +45,17 @@ def run(test_config):
     assert infer_engine.head_num == model_config.num_attention_heads // TP_SIZE
 
     # 2. check data preparation
-    input_ids_list = [[80540, 15473, 3331, 11970, 90472, 361, 61335], [80540, 15473, 3331, 11970],
-                      [80540, 15473, 3331, 11970], [80540, 15473]]
+    input_ids_list = [
+        [80540, 15473, 3331, 11970, 90472, 361, 61335],
+        [80540, 15473, 3331, 11970],
+        [80540, 15473, 3331, 11970],
+        [80540, 15473],
+    ]
     batch_size = len(input_ids_list)
     max_seq_len = max(len(li) for li in input_ids_list)
     attention_mask = [[0] * max_seq_len for _ in range(batch_size)]
     for i, li in enumerate(input_ids_list):
-        attention_mask[i][max_seq_len - len(li):] = [1 for _ in range(len(li))]
+        attention_mask[i][max_seq_len - len(li) :] = [1 for _ in range(len(li))]
     data = dict(input_ids=input_ids_list, attention_mask=attention_mask)
     inputs_batch_encoding = BatchEncoding(data=data)
     seq_lengths = [len(li) for li in input_ids_list]
@@ -78,7 +86,7 @@ def run(test_config):
 
 def check_engine(rank, world_size, port):
     disable_existing_loggers()
-    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
     run()
 
 
@@ -90,5 +98,5 @@ def test_engine():
     spawn(check_engine, TP_SIZE)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     test_engine()

tests/test_infer/test_kvcache_manager.py

@@ -1,7 +1,8 @@
 import os
-from packaging import version
+
 import pytest
 import torch
+from packaging import version
 
 from colossalai.inference.tensor_parallel import MemoryManager
 from colossalai.logging import disable_existing_loggers
@@ -14,14 +15,15 @@ LAYER_NUM = 4
 HEAD_NUM = 32
 HEAD_DIM = 128
 
-CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse('11.5')
+CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.5")
+
 
 def create_cache_manager(rank, world_size, port, batch_size, input_len, output_len, layer_num, head_num, head_dim):
-    os.environ['RANK'] = str(rank)
-    os.environ['LOCAL_RANK'] = str(rank)
-    os.environ['WORLD_SIZE'] = str(world_size)
-    os.environ['MASTER_ADDR'] = 'localhost'
-    os.environ['MASTER_PORT'] = str(port)
+    os.environ["RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(port)
     disable_existing_loggers()
 
     size = batch_size * (input_len + output_len)
@@ -41,21 +43,24 @@ def create_cache_manager(rank, world_size, port, batch_size, input_len, output_l
     assert torch.equal(prefill_locs, prefill_locs_contiguous)
     assert torch.sum(kvcache_manager.mem_state).item() == size - total_token_prefill
     kvcache_manager.alloc_contiguous(batch_size)
-    assert torch.all(kvcache_manager.mem_state[:total_token_prefill + batch_size] == False)
+    assert torch.all(kvcache_manager.mem_state[: total_token_prefill + batch_size] == False)
+
 
 @pytest.mark.skipif(not CUDA_SUPPORT, reason="kv-cache manager engine requires cuda version to be higher than 11.5")
 @pytest.mark.dist
 @rerun_if_address_is_in_use()
 def test_cache_manager_dist():
-    spawn(create_cache_manager,
-          4,
-          batch_size=BATCH_SIZE,
-          input_len=INPUT_LEN,
-          output_len=OUTPUT_LEN,
-          layer_num=LAYER_NUM,
-          head_num=HEAD_NUM,
-          head_dim=HEAD_DIM)
+    spawn(
+        create_cache_manager,
+        4,
+        batch_size=BATCH_SIZE,
+        input_len=INPUT_LEN,
+        output_len=OUTPUT_LEN,
+        layer_num=LAYER_NUM,
+        head_num=HEAD_NUM,
+        head_dim=HEAD_DIM,
+    )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     test_cache_manager_dist()

tests/test_infer/test_llama_infer.py

@@ -1,5 +1,4 @@
 import os
-import warnings
 
 import pytest
 import torch
@@ -12,13 +11,13 @@ from colossalai.shardformer import ShardConfig
 from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
 from tests.kit.model_zoo import model_zoo
 
-os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true'
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
 TPSIZE = 2
 BATCH_SIZE = 8
 MAX_INPUT_LEN = 12
 MAX_OUTPUT_LEN = 100
 
-CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse('11.5')
+CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.5")
 
 
 def init_to_get_rotary(self, base=10000):
@@ -34,8 +33,9 @@ def init_to_get_rotary(self, base=10000):
     else:
         max_seq_len = 2048 * rope_scaling_factor
     base = float(base)
-    inv_freq = 1.0 / (base**(torch.arange(0, self.config.head_dim_, 2, device="cpu", dtype=torch.float32) /
-                             self.config.head_dim_))
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, self.config.head_dim_, 2, device="cpu", dtype=torch.float32) / self.config.head_dim_)
+    )
     t = torch.arange(max_seq_len + 1024 * 64, device="cpu", dtype=torch.float32) / rope_scaling_factor
     freqs = torch.outer(t, inv_freq)
 
@@ -44,20 +44,25 @@ def init_to_get_rotary(self, base=10000):
     return
 
 
-@parameterize('test_config', [{
-    'tp_size': TPSIZE,
-}])
+@parameterize(
+    "test_config",
+    [
+        {
+            "tp_size": TPSIZE,
+        }
+    ],
+)
 def run_llama_test(test_config):
-
-    sub_model_zoo = model_zoo.get_sub_registry('transformers_llama_for_casual_lm')
+    sub_model_zoo = model_zoo.get_sub_registry("transformers_llama_for_casual_lm")
     for name, (model_fn, data_gen_fn, _, _, _) in sub_model_zoo.items():
         orig_model = model_fn()
         init_to_get_rotary(orig_model.model, base=10000)
         orig_model = orig_model.half()
         data = data_gen_fn()
 
-        shard_config = ShardConfig(enable_tensor_parallelism=True if test_config['tp_size'] > 1 else False,
-                                   inference_only=True)
+        shard_config = ShardConfig(
+            enable_tensor_parallelism=True if test_config["tp_size"] > 1 else False, inference_only=True
+        )
         infer_engine = TPInferEngine(orig_model, shard_config, BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
 
         generate_kwargs = dict(do_sample=False)
@@ -68,7 +73,7 @@ def run_llama_test(test_config):
 
 def check_llama(rank, world_size, port):
     disable_existing_loggers()
-    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
     run_llama_test()