internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/ColossalQA/colossalqa/local/llm.py b/applications/ColossalQA/colossalqa/local/llm.py
--- a/applications/ColossalQA/colossalqa/local/llm.py
+++ b/applications/ColossalQA/colossalqa/local/llm.py
@@ -136,6 +136,19 @@ def _identifying_params(self) -> Mapping[str, int]:
         """Get the identifying parameters."""
         return {"n": self.n}
 
+    def get_token_ids(self, text: str) -> List[int]:
+        """Return the ordered ids of the tokens in a text.
+
+        Args:
+            text: The string input to tokenize.
+
+        Returns:
+            A list of ids corresponding to the tokens in the text, in order they occur
+                in the text.
+        """
+        # use the colossal llm's tokenizer instead of langchain's cached GPT2 tokenizer
+        return self.api.tokenizer.encode(text)
+
 
 class VllmLLM(LLM):
     """


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

Update DOC



Discussion:

Resolving patch:
diff --git a/colossalai/inference/core/engine.py b/colossalai/inference/core/engine.py
--- a/colossalai/inference/core/engine.py
+++ b/colossalai/inference/core/engine.py
@@ -242,6 +242,7 @@ def step(self) -> List[str]:
         finished_sequences = self.request_handler.update()
 
         # Decode completed sentences.
+        # TODO : update decoding step
         for seq in finished_sequences:
             output_str = self.tokenizer.decode(seq.input_token_id + seq.output_token_id, skip_special_tokens=True)
             output_list.append(output_str)


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/examples/language/llama2/pretrain.py b/examples/language/llama2/pretrain.py
--- a/examples/language/llama2/pretrain.py
+++ b/examples/language/llama2/pretrain.py
@@ -273,11 +273,10 @@ def main():
     dataloader.sampler.set_start_index(sampler_start_idx)
     for epoch in range(start_epoch, args.num_epochs):
         dataloader.sampler.set_epoch(epoch)
-        step_nums = num_steps_per_epoch - start_step
         dataloader_iter = iter(dataloader)
 
         with tqdm(
-            range(step_nums),
+            range(start_step, num_steps_per_epoch),
             desc=f"Epoch {epoch}",
             disable=not print_flag,
             total=num_steps_per_epoch,


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py b/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
--- a/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
+++ b/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
@@ -16,7 +16,10 @@
 
 
 def unwrap(model):
-    return model.unwrap().module
+    if hasattr(model, "module"):
+        return unwrap_model(model.module)
+    else:
+        return model
 
 
 def neftune_post_forward_hook(module, input, output):


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/Colossal-LLaMA-2/inference_example.py b/applications/Colossal-LLaMA-2/inference_example.py
--- a/applications/Colossal-LLaMA-2/inference_example.py
+++ b/applications/Colossal-LLaMA-2/inference_example.py
@@ -68,7 +68,7 @@ def generate(args):
     parser.add_argument("--do_sample", type=bool, default=True, help="Set whether or not to use sampling")
     parser.add_argument("--temperature", type=float, default=0.3, help="Set temperature value")
     parser.add_argument("--top_k", type=int, default=50, help="Set top_k value for top-k-filtering")
-    parser.add_argument("--top_p", type=int, default=0.95, help="Set top_p value for generation")
+    parser.add_argument("--top_p", type=float, default=0.95, help="Set top_p value for generation")
     parser.add_argument("--input_txt", type=str, default="明月松间照，", help="The prompt input to the model")
     parser.add_argument("--prompt_style", choices=["sft", "pretrained"], default="sft", help="The style of the prompt")
     args = parser.parse_args()


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/accelerator/npu_accelerator.py b/colossalai/accelerator/npu_accelerator.py
--- a/colossalai/accelerator/npu_accelerator.py
+++ b/colossalai/accelerator/npu_accelerator.py
@@ -31,7 +31,7 @@ def get_version(self) -> str:
         """
         Return the version of the accelerator which torch is built against.
         """
-        return torch.version.npu
+        return torch.version.cann
 
     def get_current_device(self) -> torch.device:
         """


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/shardformer/modeling/t5.py b/colossalai/shardformer/modeling/t5.py
--- a/colossalai/shardformer/modeling/t5.py
+++ b/colossalai/shardformer/modeling/t5.py
@@ -593,10 +593,6 @@ def t5_encoder_model_forward(
 
 
 def get_t5_flash_attention_forward():
-    try:
-        from xformers.ops import memory_efficient_attention as me_attention
-    except:
-        raise ImportError("Error: xformers module is not installed. Please install it to use flash attention.")
     from transformers.models.t5.modeling_t5 import T5Attention
 
     def forward(
@@ -632,11 +628,11 @@ def forward(
 
         def shape(states):
             """projection"""
-            return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim)
+            return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
 
         def unshape(states):
             """reshape"""
-            return states.view(batch_size, -1, self.inner_dim)
+            return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)
 
         def project(hidden_states, proj_layer, key_value_states, past_key_value):
             """projects hidden states correctly to key/query states"""
@@ -653,8 +649,8 @@ def project(hidden_states, proj_layer, key_value_states, past_key_value):
                 if key_value_states is None:
                     # self-attn
                     # (batch_size, n_heads, key_length, dim_per_head)
-                    hidden_states = torch.cat([past_key_value, hidden_states], dim=1)
-                elif past_key_value.shape[1] != key_value_states.shape[1]:
+                    hidden_states = torch.cat([past_key_value, hidden_states], dim=2)
+                elif past_key_value.shape[2] != key_value_states.shape[1]:
                     # checking that the `sequence_length` of the `past_key_value` is the same as
                     # the provided `key_value_states` to support prefix tuning
                     # cross-attn
@@ -701,10 +697,15 @@ def project(hidden_states, proj_layer, key_value_states, past_key_value):
         else:
             position_bias_masked = position_bias
 
-        position_bias_masked = position_bias_masked.contiguous()
-        attn_output = me_attention(
-            query_states, key_states, value_states, attn_bias=position_bias_masked, p=self.dropout, scale=1.0
-        )
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_mem_efficient=True):
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=position_bias_masked,
+                dropout_p=self.dropout,
+                scale=1.0,
+            )
         attn_output = unshape(attn_output)
         attn_output = self.o(attn_output)
 


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/extensions/base_extension.py b/extensions/base_extension.py
--- a/extensions/base_extension.py
+++ b/extensions/base_extension.py
@@ -1,7 +1,7 @@
 import hashlib
 import os
 from abc import ABC, abstractmethod
-from typing import Union
+from typing import Callable, Union
 
 __all__ = ["_Extension"]
 
@@ -64,7 +64,7 @@ def is_hardware_available(self) -> bool:
         """
 
     @abstractmethod
-    def assert_hardware_compatible(self) -> bool:
+    def assert_hardware_compatible(self) -> None:
         """
         Check if the hardware required by the kernel is compatible.
         """
@@ -74,9 +74,9 @@ def build_aot(self) -> Union["CppExtension", "CUDAExtension"]:
         pass
 
     @abstractmethod
-    def build_jit(self) -> None:
+    def build_jit(self) -> Callable:
         pass
 
     @abstractmethod
-    def load(self):
+    def load(self) -> Callable:
         pass


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/extensions/cpp_extension.py b/extensions/cpp_extension.py
--- a/extensions/cpp_extension.py
+++ b/extensions/cpp_extension.py
@@ -126,7 +126,7 @@ def cxx_flags(self) -> List[str]:
     def load(self):
         try:
             op_kernel = self.import_op()
-        except ImportError:
+        except (ImportError, ModuleNotFoundError):
             # if import error occurs, it means that the kernel is not pre-built
             # so we build it jit
             op_kernel = self.build_jit()


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/Chat/coati/dataset/sft_dataset.py b/applications/Chat/coati/dataset/sft_dataset.py
--- a/applications/Chat/coati/dataset/sft_dataset.py
+++ b/applications/Chat/coati/dataset/sft_dataset.py
@@ -49,12 +49,13 @@ def _preprocess(
     max_length: int,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Preprocess the data by tokenizing."""
-    sequences = [s + t for s, t in zip(sources, targets)]
+    sequences = [s + t + tokenizer.eos_token for s, t in zip(sources, targets)]
     sequences_token = tokenizer(
-        sequences, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
+        sequences, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt", add_special_tokens=False
     )
+
     sources_token = tokenizer(
-        sources, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
+        sources, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt", add_special_tokens=False
     )
 
     assert sequences_token["attention_mask"].dim() == 2, "seq2seq model should be preprocessed differently"
@@ -65,7 +66,8 @@ def _preprocess(
         if tokenizer.padding_side == "right":
             # |prompt|completion|eos|pad|
             labels[i][:source_len] = IGNORE_INDEX
-            labels[i][-pad_len:] = IGNORE_INDEX
+            if pad_len>0:
+                labels[i][-pad_len:] = IGNORE_INDEX
         elif tokenizer.padding_side == "left":
             # |pad|prompt|completion|eos|
             labels[i][: pad_len + source_len] = IGNORE_INDEX


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py
--- a/colossalai/booster/plugin/hybrid_parallel_plugin.py
+++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py
@@ -1,5 +1,6 @@
 import ctypes
 import random
+import warnings
 from contextlib import contextmanager
 from functools import partial
 from types import MethodType
@@ -1134,7 +1135,12 @@ def configure(
                         tp_process_group=self.tp_group,
                     )
             else:
-                assert self.dp_size > 1, "Please use Zero when data parallel size is greater than 1."
+                if self.dp_size == 1:
+                    warnings.warn(
+                        "Use Zero Optimizer when data parallel size is 1 may introduce unnecessary overhead. "
+                        "If you are not intended to use cpu_offload, please consider set zero_stage=0."
+                    )
+
                 assert self.precision != "fp32", "Please set precision to 'fp16' or 'bf16' when using ZeRO."
                 optimizer = HybridParallelZeroOptimizer(
                     optimizer,


Test patch:
diff --git a/tests/test_booster/test_plugin/test_3d_plugin.py b/tests/test_booster/test_plugin/test_3d_plugin.py
--- a/tests/test_booster/test_plugin/test_3d_plugin.py
+++ b/tests/test_booster/test_plugin/test_3d_plugin.py
@@ -118,6 +118,20 @@ def check_3d_plugin(init_method: str = "none", early_stop: bool = True):
 @parameterize(
     "test_args",
     [
+        {
+            "batch_size": 8,
+            "num_steps": 4,
+            "tp": 2,
+            "pp": 2,
+            "pp_style": "1f1b",
+            "num_model_chunks": 1,
+            "num_microbatches": 4,
+            "zero": 1,
+            "precision": "fp16",
+            "initial_scale": 1,
+            "max_length": 512,
+            "gradient_accumulation_step": 2,
+        },
         {
             "batch_size": 8,
             "num_steps": 4,
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/tensor/colo_parameter.py b/colossalai/tensor/colo_parameter.py
--- a/colossalai/tensor/colo_parameter.py
+++ b/colossalai/tensor/colo_parameter.py
@@ -7,11 +7,12 @@
 
 from .colo_tensor import _convert_output
 
-WHITE_LIST_FUNCS = {torch.Tensor.__getitem__, torch.Tensor.is_floating_point}
+WHITE_LIST_FUNCS = {torch.Tensor.__getitem__}
+NO_HOOK_FUNCS = {torch.Tensor.is_floating_point}
 
 
 def is_no_hook_op(func) -> bool:
-    return func.__name__.startswith("__") and func not in WHITE_LIST_FUNCS
+    return (func.__name__.startswith("__") and func not in WHITE_LIST_FUNCS) or func in NO_HOOK_FUNCS
 
 
 def filter_colo_parameters(*args, **kwargs):
diff --git a/colossalai/tensor/param_op_hook.py b/colossalai/tensor/param_op_hook.py
--- a/colossalai/tensor/param_op_hook.py
+++ b/colossalai/tensor/param_op_hook.py
@@ -92,7 +92,10 @@ def pre_op(params: List[torch.Tensor], *args: Any) -> list:
     @staticmethod
     def post_op(params: List[torch.Tensor], arg: Any) -> Any:
         ColoParamOpHookManager._trigger_post_forward(params)
-        return PostFwdPreBwd.apply(params, arg)
+        # incase the output is a tuple, we have to flatten it
+        grad_args, other_args, grad_flags, spec = _flatten_grad_args(arg)
+        new_grad_args = PostFwdPreBwd.apply(params, *grad_args)
+        return _merge_args(new_grad_args, other_args, grad_flags, spec)
 
     @staticmethod
     def has_hook() -> bool:
@@ -113,7 +116,7 @@ def backward(ctx, *grads):
 
 class PostFwdPreBwd(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, params, args):
+    def forward(ctx, params, *args):
         ctx.params = params
         return args
 
@@ -142,7 +145,6 @@ def _flatten_grad_args(args) -> Tuple[list, list, List[bool], TreeSpec]:
             grad_args.append(arg)
         else:
             other_args.append(arg)
-    assert len(grad_args) > 0
     return grad_args, other_args, grad_flags, spec
 
 


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py b/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
--- a/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
+++ b/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
@@ -17,7 +17,7 @@
 
 def unwrap(model):
     if hasattr(model, "module"):
-        return unwrap_model(model.module)
+        return model.unwrap()
     else:
         return model
 
diff --git a/applications/Colossal-LLaMA-2/train.py b/applications/Colossal-LLaMA-2/train.py
--- a/applications/Colossal-LLaMA-2/train.py
+++ b/applications/Colossal-LLaMA-2/train.py
@@ -329,9 +329,9 @@ def main() -> None:
 
     for epoch in range(start_epoch, args.num_epochs):
         dataloader.sampler.set_epoch(epoch=epoch)
-        pbar = tqdm(desc=f"Epoch {epoch}", disable=not coordinator.is_master(), total=num_steps_per_epoch)
+        pbar = tqdm(desc=f"Epoch {epoch}", disable=not coordinator.is_master(), total=num_steps_per_epoch, initial=start_step // args.accumulation_steps)
         total_loss = torch.tensor(0.0, device=get_current_device())
-        for step, batch in enumerate(dataloader):
+        for step, batch in enumerate(dataloader, start=start_step):
             batch = {k: v.to(get_current_device()) for k, v in batch.items() if isinstance(v, torch.Tensor)}
 
             batch_output = model(**batch)


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/applications/Colossal-LLaMA-2/train.py b/applications/Colossal-LLaMA-2/train.py
--- a/applications/Colossal-LLaMA-2/train.py
+++ b/applications/Colossal-LLaMA-2/train.py
@@ -23,7 +23,7 @@
 from colossal_llama2.utils.neftune_patch import activate_neftune, deactivate_neftune
 from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
-from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer
+from transformers import LlamaForCausalLM, LlamaTokenizer
 
 import colossalai
 from colossalai.accelerator import get_accelerator
@@ -232,10 +232,12 @@ def main() -> None:
         else nullcontext()
     )
     with init_ctx:
-        model = LlamaForCausalLM(LlamaConfig.from_pretrained(args.pretrained))
+        model = LlamaForCausalLM.from_pretrained(args.pretrained)
         # Freeze part of parameters.
         if args.freeze_non_embeds_params:
             freeze_non_embeds_parameters(model=model)
+    # this is essential, otherwise the grad checkpoint will not work.
+    model.train()
 
     if args.use_grad_checkpoint:
         model.gradient_checkpointing_enable()
@@ -277,8 +279,6 @@ def main() -> None:
         lr_scheduler=lr_scheduler,
         dataloader=dataloader,
     )
-    if args.load_checkpoint is None:
-        booster.load_model(model, args.pretrained)
 
     torch.set_default_dtype(torch.float)
 


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/.github/workflows/scripts/update_setup_for_nightly.py b/.github/workflows/scripts/update_setup_for_nightly.py
new file mode 100644
--- /dev/null
+++ b/.github/workflows/scripts/update_setup_for_nightly.py
@@ -0,0 +1,34 @@
+from datetime import datetime
+
+
+def open_setup_file():
+    with open("setup.py", "r") as f:
+        file_lines = f.readlines()
+    return file_lines
+
+
+def replace_nightly_package_info(file_lines):
+    version = datetime.today().strftime("%Y.%m.%d")
+    package_name = "colossalai-nightly"
+
+    for idx, line in enumerate(file_lines):
+        if "version = get_version()" in line:
+            file_lines[idx] = f'version = "{version}"\n'
+        if 'package_name = "colossalai"' in line:
+            file_lines[idx] = f'package_name = "{package_name}"\n'
+    return file_lines
+
+
+def write_setup_file(file_lines):
+    with open("setup.py", "w") as f:
+        f.writelines(file_lines)
+
+
+def main():
+    file_lines = open_setup_file()
+    file_lines = replace_nightly_package_info(file_lines)
+    write_setup_file(file_lines)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,5 @@
 import os
 import sys
-from datetime import datetime
 from typing import List
 
 from setuptools import find_packages, setup
@@ -15,7 +14,6 @@
 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 BUILD_EXT = int(os.environ.get("BUILD_EXT", "0")) == 1
-IS_NIGHTLY = int(os.environ.get("NIGHTLY", "0")) == 1
 
 # we do not support windows currently
 if sys.platform == "win32":
@@ -96,23 +94,15 @@ def get_version() -> str:
 else:
     ext_modules = []
 
-# always put not nightly branch as the if branch
-# otherwise github will treat colossalai-nightly as the project name
-# and it will mess up with the dependency graph insights
-if not IS_NIGHTLY:
-    version = get_version()
-    package_name = "colossalai"
-else:
-    # use date as the nightly version
-    version = datetime.today().strftime("%Y.%m.%d")
-    package_name = "colossalai-nightly"
+version = get_version()
+package_name = "colossalai"
 
 setup(
     name=package_name,
     version=version,
     packages=find_packages(
         exclude=(
-            "op_builder",
+            "extensions",
             "benchmark",
             "docker",
             "tests",


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Wrong import in ColossalAuto's meta_registry/binary_elementwise_ops.py
### 🐛 Describe the bug

# Problem description

The file `colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py` contains the following line: 

```python
from ..constants import BCAST_FUNC_OP
```

However, the file `colossalai/auto_parallel/meta_profiler/constants.py` which this import refers to does not contain any `BCAST_FUNC_OP`. This leads to an `ImportError` when running ColossalAuto in release 0.3.3 and newer. 

This constant can be found in the file `colossalai/auto_parallel/tensor_shard/constants.py`. The last commit to `colossalai/auto_parallel/meta_profiler/constants.py` (commit ID `079bf3cb`) removes the import of tensor_shard's `constants.py` from meta_profiler's `constants.py` (seemingly due to an automated refactoring).

# Solution

Since no other file in the `meta_registry` module uses constants from the `tensor_shard/constants.py` and to avoid automated removal of "unused" imports in the future, the import statement in question in above-mentioned `binary_elementwise_ops.py` could be changed to: 

```python
from colossalai.auto_parallel.tensor_shard.constants import BCAST_FUNC_OP
```

### Environment

- Python 3.8
- Torch 1.12.0
- no CUDA


Discussion:

Resolving patch:
diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py b/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py
--- a/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py
+++ b/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py
@@ -4,9 +4,9 @@
 
 from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
 from colossalai._analyzer.fx.node_util import compute_size_in_bytes as activation_size
+from colossalai.auto_parallel.tensor_shard.constants import BCAST_FUNC_OP
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
 
-from ..constants import BCAST_FUNC_OP
 from ..registry import meta_register
 
 __all__ = ["binary_elementwise_meta_info"]


Test patch: