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/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py
--- a/colossalai/shardformer/modeling/llama.py
+++ b/colossalai/shardformer/modeling/llama.py
@@ -16,6 +16,7 @@
 from colossalai.shardformer.shard import ShardConfig
 
 from ..layer import cross_entropy_1d
+from ..layer._operation import _gather
 
 try:
     from transformers.models.llama.modeling_llama import _prepare_4d_causal_attention_mask
@@ -288,6 +289,9 @@ def llama_for_causal_lm_forward(
                     shift_logits = shift_logits.view(-1, self.config.vocab_size)
                     loss = loss_fct(shift_logits, shift_labels)
 
+            if not shard_config.parallel_output:
+                logits = _gather(logits, -1, shard_config.tensor_parallel_process_group)
+
             if not return_dict:
                 output = (logits,) + outputs[1:]
                 return (loss,) + output if loss is not None else output
@@ -588,6 +592,9 @@ def forward(
                 shift_logits = shift_logits.view(-1, self.config.vocab_size)
                 loss = loss_fct(shift_logits, shift_labels)
 
+        if not shard_config.parallel_output:
+            logits = _gather(logits, -1, shard_config.tensor_parallel_process_group)
+
         if not return_dict:
             output = (logits,) + outputs[1:]
             return (loss,) + output if loss is not None else output
diff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py
--- a/colossalai/shardformer/shard/shard_config.py
+++ b/colossalai/shardformer/shard/shard_config.py
@@ -34,6 +34,7 @@ class ShardConfig:
     enable_all_optimization: bool = False
     enable_sequence_parallelism: bool = False
     enable_sequence_overlap: bool = False
+    parallel_output = True
     extra_kwargs: Dict[str, Any] = field(default_factory=dict)
     # pipeline_parallel_size: int
     # data_parallel_size: int


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -121,8 +121,9 @@ def get_version() -> str:
             "tests",
             "scripts",
             "requirements",
+            "extensions",
             "*.egg-info",
-        )
+        ),
     ),
     description="An integrated large-scale model training system with efficient parallelization techniques",
     long_description=fetch_readme(),
@@ -153,10 +154,7 @@ def get_version() -> str:
     ],
     package_data={
         "colossalai": [
-            "_C/*.pyi",
-            "kernel/cuda_native/csrc/*",
-            "kernel/cuda_native/csrc/kernel/*",
-            "kernel/cuda_native/csrc/kernels/include/*",
+            "kernel/extensions/csrc/**/*",
         ]
     },
 )


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
@@ -15,7 +15,7 @@ def load_model(model_path, device="cuda", **kwargs):
     model.to(device)
 
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side='left')
     except OSError:
         raise ImportError("Tokenizer not found. Please check if the tokenizer exists or the model path is correct.")
 
@@ -29,6 +29,7 @@ def generate(args):
     if args.prompt_style == "sft":
         conversation = default_conversation.copy()
         conversation.append_message("Human", args.input_txt)
+        conversation.append_message("Assistant", None)
         input_txt = conversation.get_prompt()
     else:
         BASE_INFERENCE_SUFFIX = "\n\n->\n\n"
@@ -46,7 +47,7 @@ def generate(args):
         num_return_sequences=1,
     )
     response = tokenizer.decode(output.cpu()[0, num_input_tokens:], skip_special_tokens=True)
-    logger.info(f"Question: {input_txt} \n\n Answer: \n{response}")
+    logger.info(f"\nHuman: {args.input_txt} \n\nAssistant: \n{response}")
     return response
 
 


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[BUG]: Import error 
### 🐛 Describe the bug

train a diffusion model from scatch ,bug showed up 
code
```
 python main.py --logdir /tmp_log  --train --base configs/train_colossalai.yaml --ckpt 512-base-ema.ckpt
```
error 
```
ImportError: cannot import name 'forward_vit' from 'ldm.modules.midas.midas.blocks' (/home/danyow/code/ColossalAI/examples/images/diffusion/ldm/modules/midas/midas/blocks.py)

```

### Environment

_No response_


Discussion:

Resolving patch:
diff --git a/examples/images/diffusion/ldm/modules/midas/midas/dpt_depth.py b/examples/images/diffusion/ldm/modules/midas/midas/dpt_depth.py
--- a/examples/images/diffusion/ldm/modules/midas/midas/dpt_depth.py
+++ b/examples/images/diffusion/ldm/modules/midas/midas/dpt_depth.py
@@ -2,7 +2,8 @@
 import torch.nn as nn
 
 from .base_model import BaseModel
-from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder, forward_vit
+from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder
+from .vit import forward_vit
 
 
 def _make_fusion_block(features, use_bn):


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/inference/inference_ops_cuda.py b/extensions/inference/inference_ops_cuda.py
--- a/extensions/inference/inference_ops_cuda.py
+++ b/extensions/inference/inference_ops_cuda.py
@@ -12,6 +12,7 @@ def sources_files(self):
             for fname in [
                 "cuda/colossal_inference_C_frontend.cpp",
                 "cuda/decode_kv_cache_memcpy_kernel.cu",
+                "cuda/activation_kernel.cu",
             ]
         ]
         return ret


Test patch:
diff --git a/tests/test_infer/test_ops/cuda/test_silu_and_mul.py b/tests/test_infer/test_ops/cuda/test_silu_and_mul.py
new file mode 100644
--- /dev/null
+++ b/tests/test_infer/test_ops/cuda/test_silu_and_mul.py
@@ -0,0 +1,33 @@
+import pytest
+import torch
+
+from colossalai.kernel.kernel_loader import InferenceOpsLoader
+from colossalai.utils import get_current_device
+
+inference_ops = InferenceOpsLoader().load()
+
+
+@pytest.mark.parametrize("SHAPE_X", [2])
+@pytest.mark.parametrize("SHAPE_Y", [64])
+@pytest.mark.parametrize("SHAPE_Z", [11008])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+def test_silu_and_mul(SHAPE_X, SHAPE_Y, SHAPE_Z, dtype):
+    torch.manual_seed(5)
+    device = get_current_device()
+    ref_input = torch.randn(SHAPE_X, SHAPE_Y, SHAPE_Z, dtype=dtype, device=device)
+    origin_input = ref_input.clone()
+
+    act_out = torch.nn.functional.silu(ref_input[0], inplace=True)
+    ref_out = act_out * ref_input[1]
+
+    origin_out = inference_ops.silu_and_mul(origin_input)
+
+    if dtype == torch.float32:
+        assert torch.allclose(origin_out, ref_out, atol=1e-5, rtol=1e-5)
+    else:
+        assert torch.allclose(origin_out, ref_out, atol=1e-3, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    test_silu_and_mul(2, 64, 11008, torch.float32)
+    test_silu_and_mul(2, 64, 11008, torch.float16)
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/legacy/tensor/tensor_spec.py b/colossalai/legacy/tensor/tensor_spec.py
--- a/colossalai/legacy/tensor/tensor_spec.py
+++ b/colossalai/legacy/tensor/tensor_spec.py
@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 from colossalai.legacy.tensor.distspec import DistPlacementPattern, _DistSpec
@@ -17,5 +17,5 @@ class ColoTensorSpec:
     """
 
     pg: ProcessGroup
-    dist_attr: Optional[_DistSpec] = _DistSpec(DistPlacementPattern.REPLICATE)
+    dist_attr: Optional[_DistSpec] = field(default_factory=lambda: _DistSpec(DistPlacementPattern.REPLICATE))
     compute_attr: Optional[ComputeSpec] = None


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
@@ -56,6 +56,7 @@ def format_numel_str(numel: int) -> str:
 
 def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
     dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
+    tensor = tensor.data
     tensor.div_(dist.get_world_size())
     return tensor
 


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/moe/moe_cuda.py b/extensions/moe/moe_cuda.py
--- a/extensions/moe/moe_cuda.py
+++ b/extensions/moe/moe_cuda.py
@@ -11,7 +11,7 @@ def include_dirs(self):
         return ret
 
     def sources_files(self):
-        ret = [self.csrc_abs_path(fname) for fname in ["cuda/moe.cpp", "cuda/moe_kernel.cu"]]
+        ret = [self.csrc_abs_path(fname) for fname in ["cuda/pybind/moe.cpp", "cuda/moe_kernel.cu"]]
         return ret
 
     def cxx_flags(self):


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/inference/core/request_handler.py b/colossalai/inference/core/request_handler.py
--- a/colossalai/inference/core/request_handler.py
+++ b/colossalai/inference/core/request_handler.py
@@ -298,8 +298,8 @@ def search_tokens(self, generation_config: GenerationConfig, logits):
         """
         # do logit processor
         # NOTE: need to decide the granularity to process logits (sequence or batch)
+        config_dict = generation_config.to_dict()
         for type in ["top_k", "top_p", "min_p"]:
-            config_dict = generation_config.to_dict()
             if type in config_dict and config_dict[type] is not None:
                 logits = logit_processor(type, logits, config_dict[type])
 
diff --git a/colossalai/inference/logit_processors.py b/colossalai/inference/logit_processors.py
--- a/colossalai/inference/logit_processors.py
+++ b/colossalai/inference/logit_processors.py
@@ -36,21 +36,23 @@ def top_p_logit_processor(logits, top_p: float):
     cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
 
     sorted_indices_to_remove = cumulative_probs > top_p
-    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+
+    sorted_indices_to_remove = torch.roll(sorted_indices_to_remove, 1, -1)
     sorted_indices_to_remove[..., 0] = 0
 
     indices_to_remove = sorted_indices_to_remove.scatter(dim=1, index=sorted_indices, src=sorted_indices_to_remove)
     logits[indices_to_remove] = -float("inf")
     return logits
 
-def logit_processor(processor:str, logits , attrs):
+
+def logit_processor(processor: str, logits, attrs):
     """
     do logit process for given logits.
 
     Args:
-        processor(str): the type of logit processor 
+        processor(str): the type of logit processor
         logits(torch.Tensor): input logits
-        attrs(dict): attrs of the logit processor 
+        attrs(dict): attrs of the logit processor
 
     Returns:
         logits after process
@@ -61,6 +63,6 @@ def logit_processor(processor:str, logits , attrs):
         func = _LOGIT_PROCESSOR_MAP[processor]
         try:
             logits = func(logits, attrs)
-        except Exception as e:
+        except Exception:
             return logits
-        return logits
\ No newline at end of file
+        return logits


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/grok-1/inference.py b/examples/language/grok-1/inference.py
--- a/examples/language/grok-1/inference.py
+++ b/examples/language/grok-1/inference.py
@@ -1,7 +1,7 @@
 import time
 
 import torch
-from transformers import AutoModelForCausalLM, LlamaTokenizerFast
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from utils import get_defualt_parser, inference, print_output
 
 if __name__ == "__main__":
@@ -9,6 +9,9 @@
     args = parser.parse_args()
     start = time.time()
     torch.set_default_dtype(torch.bfloat16)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained, trust_remote_code=True)
+
     model = AutoModelForCausalLM.from_pretrained(
         args.pretrained,
         trust_remote_code=True,
@@ -18,10 +21,6 @@
     model.eval()
     init_time = time.time() - start
 
-    # A transformers-compatible version of the grok-1 tokenizer by Xenova
-    # https://huggingface.co/Xenova/grok-1-tokenizer
-    tokenizer = LlamaTokenizerFast.from_pretrained("Xenova/grok-1-tokenizer")
-
     for text in args.text:
         output = inference(
             model,
diff --git a/examples/language/grok-1/inference_tp.py b/examples/language/grok-1/inference_tp.py
--- a/examples/language/grok-1/inference_tp.py
+++ b/examples/language/grok-1/inference_tp.py
@@ -2,7 +2,7 @@
 
 import torch
 from grok1_policy import Grok1ForCausalLMPolicy
-from transformers import AutoModelForCausalLM, LlamaTokenizerFast
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from utils import get_defualt_parser, inference, print_output
 
 import colossalai
@@ -27,6 +27,9 @@
     )
     booster = Booster(plugin=plugin)
     torch.set_default_dtype(torch.bfloat16)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained, trust_remote_code=True)
+
     with LazyInitContext(default_device=get_current_device()):
         model = AutoModelForCausalLM.from_pretrained(
             args.pretrained, trust_remote_code=True, torch_dtype=torch.bfloat16
@@ -35,10 +38,6 @@
     model.eval()
     init_time = time.time() - start
 
-    # A transformers-compatible version of the grok-1 tokenizer by Xenova
-    # https://huggingface.co/Xenova/grok-1-tokenizer
-    tokenizer = LlamaTokenizerFast.from_pretrained("Xenova/grok-1-tokenizer")
-
     for text in args.text:
         output = inference(
             model.unwrap(),


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/inference/benchmark_ops/benchmark_rmsnorm.py b/examples/inference/benchmark_ops/benchmark_rmsnorm.py
--- a/examples/inference/benchmark_ops/benchmark_rmsnorm.py
+++ b/examples/inference/benchmark_ops/benchmark_rmsnorm.py
@@ -35,7 +35,7 @@
         styles=[("red", "-"), ("blue", "-"), ("yellow", "-"), ("red", "--"), ("blue", "--"), ("yellow", "--")],
         ylabel="ms",
         plot_name=f"RMSNorm benchmarking results",
-        args={"HIDDEN_SIZE": 1024},
+        args={"HIDDEN_SIZE": 5120},
     )
 ]
 


Test patch:
diff --git a/tests/test_infer/test_ops/cuda/test_rms_layernorm.py b/tests/test_infer/test_ops/cuda/test_rms_layernorm.py
--- a/tests/test_infer/test_ops/cuda/test_rms_layernorm.py
+++ b/tests/test_infer/test_ops/cuda/test_rms_layernorm.py
@@ -9,7 +9,7 @@
 
 
 @pytest.mark.parametrize("M", [2, 4, 8, 16])
-@pytest.mark.parametrize("N", [64, 128, 512])
+@pytest.mark.parametrize("N", [64, 128, 512, 5120])
 def test_rms_layernorm(M: int, N: int):
     torch.manual_seed(123)
     torch.cuda.empty_cache()
@@ -48,4 +48,4 @@ def test_rms_layernorm(M: int, N: int):
 
 
 if __name__ == "__main__":
-    test_rms_layernorm(16, 512)
+    test_rms_layernorm(16, 5120)
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/benchmark.py b/examples/language/llama/benchmark.py
similarity index 100%
rename from examples/language/llama2/benchmark.py
rename to examples/language/llama/benchmark.py


Test patch:
diff --git a/examples/language/llama2/test_ci.sh b/examples/language/llama/test_ci.sh
similarity index 100%
rename from examples/language/llama2/test_ci.sh
rename to examples/language/llama/test_ci.sh
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/opt.py b/colossalai/shardformer/modeling/opt.py
--- a/colossalai/shardformer/modeling/opt.py
+++ b/colossalai/shardformer/modeling/opt.py
@@ -58,20 +58,6 @@ class OPTPipelineForwards:
     under pipeline setting.
     """
 
-    @staticmethod
-    def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
-        """
-        Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
-        """
-        bsz, src_len = mask.size()
-        tgt_len = tgt_len if tgt_len is not None else src_len
-
-        expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
-
-        inverted_mask = 1.0 - expanded_mask
-
-        return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
-
     @staticmethod
     def opt_model_forward(
         self: OPTModel,


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/policies/gptj.py b/colossalai/shardformer/policies/gptj.py
--- a/colossalai/shardformer/policies/gptj.py
+++ b/colossalai/shardformer/policies/gptj.py
@@ -54,7 +54,6 @@ def module_policy(self):
         if self.shard_config.enable_sequence_parallelism:
             self.shard_config.enable_sequence_parallelism = False
             warnings.warn("GPTJ doesn't support sequence parallelism now, will ignore the sequence parallelism flag.")
-        use_sequence_parallel = self.shard_config.enable_sequence_parallelism
 
         overlap = self.shard_config.enable_sequence_overlap
         if self.shard_config.enable_tensor_parallelism:
@@ -78,7 +77,6 @@ def module_policy(self):
                         suffix="attn.k_proj",
                         target_module=col_nn.Linear1D_Col,
                         kwargs={
-                            "seq_parallel": use_sequence_parallel,
                             "overlap": overlap,
                         },
                     ),
@@ -86,7 +84,6 @@ def module_policy(self):
                         suffix="attn.q_proj",
                         target_module=col_nn.Linear1D_Col,
                         kwargs={
-                            "seq_parallel": use_sequence_parallel,
                             "overlap": overlap,
                         },
                     ),
@@ -94,24 +91,20 @@ def module_policy(self):
                         suffix="attn.v_proj",
                         target_module=col_nn.Linear1D_Col,
                         kwargs={
-                            "seq_parallel": use_sequence_parallel,
                             "overlap": overlap,
                         },
                     ),
                     SubModuleReplacementDescription(
                         suffix="attn.out_proj",
                         target_module=col_nn.Linear1D_Row,
-                        kwargs={"seq_parallel": use_sequence_parallel},
                     ),
                     SubModuleReplacementDescription(
                         suffix="mlp.fc_in",
                         target_module=col_nn.Linear1D_Col,
-                        kwargs={"seq_parallel": use_sequence_parallel},
                     ),
                     SubModuleReplacementDescription(
                         suffix="mlp.fc_out",
                         target_module=col_nn.Linear1D_Row,
-                        kwargs={"seq_parallel": use_sequence_parallel},
                     ),
                     SubModuleReplacementDescription(
                         suffix="attn.attn_dropout",


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/openmoe/model/modeling_openmoe.py b/examples/language/openmoe/model/modeling_openmoe.py
--- a/examples/language/openmoe/model/modeling_openmoe.py
+++ b/examples/language/openmoe/model/modeling_openmoe.py
@@ -35,7 +35,20 @@
     replace_return_docstrings,
 )
 
-from colossalai.kernel.extensions.flash_attention import HAS_FLASH_ATTN
+try:
+    # TODO: remove this after updating openmoe example
+    # NOTE(yuanheng-zhao): This is a temporary fix for the issue that
+    # the flash_attention module is not imported correctly for different CI tests.
+    # We replace the import path `colossalai.kernel.extensions.flash_attention`
+    # because in the current example test, colossalai version <= 0.3.6 is installed,
+    # where `colossalai.kernel.extensions.flash_attention` is still valid;
+    # however in unit test `test_moe_checkpoint`, the lastest version of colossalai is installed,
+    # where extension has been refactored and the path is not valid.
+    import flash_attention  # noqa
+
+    HAS_FLASH_ATTN = True
+except:
+    HAS_FLASH_ATTN = False
 from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
 from colossalai.moe.layers import SparseMLP
 from colossalai.moe.manager import MOE_MANAGER


Test patch: