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/op_builder/gptq.py b/op_builder/gptq.py
--- a/op_builder/gptq.py
+++ b/op_builder/gptq.py
@@ -1,19 +1,17 @@
-import os
-import torch
 import re
 
+import torch
+
 from .builder import Builder
-from .utils import append_nvcc_threads, get_cuda_cc_flag
+from .utils import append_nvcc_threads
 
-class GPTQBuilder(Builder):
 
+class GPTQBuilder(Builder):
     NAME = "cu_gptq"
     PREBUILT_IMPORT_PATH = "colossalai._C.cu_gptq"
 
     def __init__(self):
-        super().__init__(name=GPTQBuilder.NAME,
-        prebuilt_import_path=GPTQBuilder.PREBUILT_IMPORT_PATH)
-
+        super().__init__(name=GPTQBuilder.NAME, prebuilt_import_path=GPTQBuilder.PREBUILT_IMPORT_PATH)
 
     def include_dirs(self):
         ret = [self.csrc_abs_path("gptq"), self.get_cuda_home_include()]
@@ -21,32 +19,38 @@ def include_dirs(self):
 
     def sources_files(self):
         ret = [
-            self.csrc_abs_path(fname) for fname in [
-                'gptq/linear_gptq.cpp',
-                'gptq/column_remap.cu',
-                'gptq/cuda_buffers.cu',
-                'gptq/q4_matmul.cu',
-                'gptq/q4_matrix.cu'
+            self.csrc_abs_path(fname)
+            for fname in [
+                "gptq/linear_gptq.cpp",
+                "gptq/column_remap.cu",
+                "gptq/cuda_buffers.cu",
+                "gptq/q4_matmul.cu",
+                "gptq/q4_matrix.cu",
             ]
         ]
         return ret
 
     def cxx_flags(self):
-        return ['-O3'] + self.version_dependent_macros
+        return ["-O3"] + self.version_dependent_macros
 
     def nvcc_flags(self):
-        extra_cuda_flags = ['-v',
-            '-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__',
-            '-U__CUDA_NO_HALF2_OPERATORS__', '-DTHRUST_IGNORE_CUB_VERSION_CHECK', "-lcublas", "-std=c++17"
+        extra_cuda_flags = [
+            "-v",
+            "-std=c++14",
+            "-U__CUDA_NO_HALF_OPERATORS__",
+            "-U__CUDA_NO_HALF_CONVERSIONS__",
+            "-U__CUDA_NO_HALF2_OPERATORS__",
+            "-DTHRUST_IGNORE_CUB_VERSION_CHECK",
+            "-lcublas",
+            "-std=c++17",
         ]
 
-
         for arch in torch.cuda.get_arch_list():
-            res = re.search(r'sm_(\d+)', arch)
+            res = re.search(r"sm_(\d+)", arch)
             if res:
                 arch_cap = res[1]
                 if int(arch_cap) >= 80:
-                    extra_cuda_flags.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])
+                    extra_cuda_flags.extend(["-gencode", f"arch=compute_{arch_cap},code={arch}"])
 
-        ret = ['-O3', '--use_fast_math'] + self.version_dependent_macros + extra_cuda_flags
-        return append_nvcc_threads(ret)
\ No newline at end of file
+        ret = ["-O3", "--use_fast_math"] + self.version_dependent_macros + extra_cuda_flags
+        return append_nvcc_threads(ret)


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[DOC]: A typo in document of booster low_level_zero plugin
### 📚 The doc issue

There is a typo in the document of low_level zero plugin. [English document website](https://colossalai.org/docs/basics/booster_plugins/#low-level-zero-plugin)

![image](https://github.com/hpcaitech/ColossalAI/assets/34569304/c97a0fbf-3bd5-4501-85c1-dc92d1fd0327)

And this typo results from the comment of relavant code, [colossalai/booster/plugin/low_level_zero_plugin.py#L229](https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/booster/plugin/low_level_zero_plugin.py#L229) more exactly.


[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/booster/plugin/low_level_zero_plugin.py b/colossalai/booster/plugin/low_level_zero_plugin.py
--- a/colossalai/booster/plugin/low_level_zero_plugin.py
+++ b/colossalai/booster/plugin/low_level_zero_plugin.py
@@ -226,7 +226,7 @@ class LowLevelZeroPlugin(DPPluginBase):
     ```
 
     Args:
-        strage (int, optional): ZeRO stage. Defaults to 1.
+        stage (int, optional): ZeRO stage. Defaults to 1.
         precision (str, optional): precision. Support 'fp16', 'bf16' and 'fp32'. Defaults to 'fp16'.
         initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32.
         min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1.


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/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py
--- a/colossalai/nn/optimizer/cpu_adam.py
+++ b/colossalai/nn/optimizer/cpu_adam.py
@@ -77,6 +77,7 @@ def __init__(
         super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)
         self.adamw_mode = adamw_mode
         cpu_adam = CPUAdamBuilder().load()
+        # if you find yourself stuck here, make sure that you install colossalai with CUDA_EXT=1 specification
         self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)
 
     def torch_adam_update(


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/flash_attention_patch.py b/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
--- a/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
+++ b/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
@@ -65,6 +65,7 @@ def attention_forward(
     past_key_value: Optional[Tuple[torch.Tensor]] = None,
     output_attentions: bool = False,
     use_cache: bool = False,
+    **kwargs
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
     """
     Re-define LLaMA-2 `LlamaAttention` forward method using flash-attention.


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/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py
--- a/colossalai/nn/optimizer/cpu_adam.py
+++ b/colossalai/nn/optimizer/cpu_adam.py
@@ -9,7 +9,8 @@
 
 
 class CPUAdam(NVMeOptimizer):
-    """Implements Adam algorithm.
+    """
+    Implements Adam algorithm.
 
     Supports parameters updating on both GPU and CPU, depending on the device of parameters.
     But the parameters and gradients should on the same device:


Test patch:
diff --git a/tests/test_zero/test_low_level/test_zero1_2.py b/tests/test_zero/test_low_level/test_zero1_2.py
--- a/tests/test_zero/test_low_level/test_zero1_2.py
+++ b/tests/test_zero/test_low_level/test_zero1_2.py
@@ -106,7 +106,8 @@ def exam_zero_1_2():
 
 
 @parameterize("dtype", [torch.float16, torch.bfloat16])
-def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype):
+@parameterize("master_weights", [True, False])
+def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype, master_weights: bool):
     """
     In this test, two pairs of model and optimizers are created.
     1. zero: use sharded optimizer and fp16 parameters
@@ -131,7 +132,11 @@ def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype):
     # in `check_sharded_param_consistency.py`, we will test whether
     # level 1 and 2 will produce exactly the same results
     zero_optimizer = LowLevelZeroOptimizer(
-        zero_optimizer, overlap_communication=True, initial_scale=1, reduce_bucket_size=1024 * 1024
+        zero_optimizer,
+        overlap_communication=True,
+        initial_scale=1,
+        reduce_bucket_size=1024 * 1024,
+        master_weights=master_weights,
     )
 
     torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)
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/gemini_plugin.py b/colossalai/booster/plugin/gemini_plugin.py
--- a/colossalai/booster/plugin/gemini_plugin.py
+++ b/colossalai/booster/plugin/gemini_plugin.py
@@ -150,24 +150,24 @@ def save_sharded_optimizer(
         # Preparing file paths and index file.
         states_name, save_index_file, param_group_file = get_optimizer_base_filenames(prefix)
         index_file = CheckpointIndexFile(checkpoint)
+        index_file.append_meta_data("param_groups", param_group_file)
 
         # Store the information of param groups to param_group_file.
-        index_file.append_meta_data("param_groups", param_group_file)
-        group_file_path = os.path.join(checkpoint, param_group_file)
-        param_groups = optimizer.get_param_groups_for_saving()
-        torch.save(param_groups, group_file_path)
+        if self.coordinator.is_master():
+            group_file_path = os.path.join(checkpoint, param_group_file)
+            param_groups = optimizer.get_param_groups_for_saving()
+            torch.save(param_groups, group_file_path)
 
         # States are broken into shards within max_shard_size.
         state_dict_shard = optimizer.state_shard(prefix=prefix, max_shard_size=size_per_shard, only_rank_0=True)
 
         # Save shards of optimizer states.
-        is_master = self.coordinator.is_master()
         total_size = save_state_dict_shards(
             sharded_state_dict=state_dict_shard,
             checkpoint=checkpoint,
             index_file=index_file,
             base_filename=states_name,
-            is_master=is_master,
+            is_master=self.coordinator.is_master(),
             use_safetensors=False,
         )
 
diff --git a/colossalai/booster/plugin/low_level_zero_plugin.py b/colossalai/booster/plugin/low_level_zero_plugin.py
--- a/colossalai/booster/plugin/low_level_zero_plugin.py
+++ b/colossalai/booster/plugin/low_level_zero_plugin.py
@@ -119,11 +119,12 @@ def save_sharded_optimizer(
         # Preparing file paths and index file.
         states_name, save_index_file, param_group_file = get_optimizer_base_filenames(prefix)
         index_file = CheckpointIndexFile(checkpoint)
+        index_file.append_meta_data("param_groups", param_group_file)
 
         # Store the information of param groups to param_group_file.
-        index_file.append_meta_data("param_groups", param_group_file)
-        group_file_path = os.path.join(checkpoint, param_group_file)
-        save_param_groups(state_dict, group_file_path)
+        if self.coordinator.is_master():
+            group_file_path = os.path.join(checkpoint, param_group_file)
+            save_param_groups(state_dict, group_file_path)
 
         # Save shards of optimizer states.
         total_size = 0


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/bert/finetune.py b/examples/language/bert/finetune.py
--- a/examples/language/bert/finetune.py
+++ b/examples/language/bert/finetune.py
@@ -248,7 +248,7 @@ def main():
     cfg = AutoConfig.from_pretrained(model_name, num_labels=data_builder.num_labels)
 
     if model_name == "bert-base-uncased":
-        model = BertForSequenceClassification.from_pretrained(model_name, config=cfg).cuda()
+        model = BertForSequenceClassification.from_pretrained(model_name, config=cfg, cache_dir="./").cuda()
     elif model_name == "albert-xxlarge-v2":
         model = AlbertForSequenceClassification.from_pretrained(model_name, config=cfg)
     else:


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/ColossalEval/colossal_eval/evaluate/dataset_evaluator/dataset_evaluator.py b/applications/ColossalEval/colossal_eval/evaluate/dataset_evaluator/dataset_evaluator.py
--- a/applications/ColossalEval/colossal_eval/evaluate/dataset_evaluator/dataset_evaluator.py
+++ b/applications/ColossalEval/colossal_eval/evaluate/dataset_evaluator/dataset_evaluator.py
@@ -60,6 +60,11 @@ def _calculate_label_metrics(self, metric: str, category: str):
                             sample["output"], ref, all_classes=self.data[category]["inference_kwargs"]["all_classes"]
                         ),
                     )
+
+                    score = max(
+                        score,
+                        metric_helper.accuracy_by_options(sample["input"], sample["output"], ref),
+                    )
                 softmaxs.append(references[i] if score == 1 else -1)
             else:
                 softmaxs.append(np.argmax(np.array(list(sample["softmax_over_choices"].values()))))
diff --git a/applications/ColossalEval/colossal_eval/evaluate/dataset_evaluator/metrics.py b/applications/ColossalEval/colossal_eval/evaluate/dataset_evaluator/metrics.py
--- a/applications/ColossalEval/colossal_eval/evaluate/dataset_evaluator/metrics.py
+++ b/applications/ColossalEval/colossal_eval/evaluate/dataset_evaluator/metrics.py
@@ -443,6 +443,20 @@ def multi_choice_accuracy(prediction, reference, **kwargs):
     return score
 
 
+def accuracy_by_options(question, prediction, reference):
+    pattern = r"[A-Z]\. [^\n]+"
+    options = re.findall(pattern, question)
+    answer = prediction.split("\n\n")[0]
+
+    for option in options:
+        choice, content = option.split(". ", 1)
+
+        if choice == reference and content == answer:
+            return 1
+
+    return 0
+
+
 def combined_single_choice_accuracy(prediction, reference, **kwargs):
     return single_choice_accuracy(prediction, reference, **kwargs)
 
diff --git a/applications/ColossalEval/colossal_eval/models/huggingface.py b/applications/ColossalEval/colossal_eval/models/huggingface.py
--- a/applications/ColossalEval/colossal_eval/models/huggingface.py
+++ b/applications/ColossalEval/colossal_eval/models/huggingface.py
@@ -96,7 +96,7 @@ def _load_tokenizer(self, path: str, tokenizer_path: Optional[str], tokenizer_kw
             self.logger.warning("pad_token_id is not set for the tokenizer. " "Using eos_token_id as pad_token_id.")
             if self.tokenizer.eos_token:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
-            elif self.tokenizer.eod_id:
+            elif hasattr(self.tokenizer, "eod_id"):
                 # Qwen has an eod token "
".
                 self.tokenizer.pad_token_id = self.tokenizer.eod_id
 


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/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py
--- a/colossalai/checkpoint_io/utils.py
+++ b/colossalai/checkpoint_io/utils.py
@@ -11,6 +11,7 @@
 import torch.nn as nn
 from packaging.version import Version
 from torch.optim import Optimizer
+from torch.utils._pytree import tree_map
 
 from colossalai.tensor.d_tensor import (
     is_customized_distributed_tensor,
@@ -293,7 +294,6 @@ def shard_optimizer_checkpoint(state_dict: dict, max_shard_size: int = 1024) ->
 # Helper functions for saving state dict
 # ======================================
 
-
 def save_state_dict(state_dict: dict, checkpoint_file_path: str, use_safetensors: bool) -> None:
     """
     Save state dict to checkpoint.
@@ -303,6 +303,9 @@ def save_state_dict(state_dict: dict, checkpoint_file_path: str, use_safetensors
         checkpoint_file_path (str): path to the checkpoint file.
         use_safetensors (bool): whether to use safetensors to save the checkpoint.
     """
+    # Move all tensors in the state_dict to CPU before saving to avoid serialization issues
+    state_dict_cpu = tree_map(lambda x: x.cpu() if torch.is_tensor(x) else x, state_dict)
+    
     if use_safetensors:
         assert is_safetensors_available(), "safetensors is not available."
         assert checkpoint_file_path.endswith(
@@ -310,9 +313,9 @@ def save_state_dict(state_dict: dict, checkpoint_file_path: str, use_safetensors
         ), "safetensors only supports .safetensors suffix for checkpoint file."
         from safetensors.torch import save_file as safe_save_file
 
-        safe_save_file(state_dict, checkpoint_file_path, metadata={"format": "pt"})
+        safe_save_file(state_dict_cpu, checkpoint_file_path, metadata={"format": "pt"})
     else:
-        torch.save(state_dict, checkpoint_file_path)
+        torch.save(state_dict_cpu, checkpoint_file_path)
 
 
 def save_param_groups(state_dict: dict, group_file_path: str) -> 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/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py
--- a/colossalai/shardformer/modeling/llama.py
+++ b/colossalai/shardformer/modeling/llama.py
@@ -1,5 +1,5 @@
 import warnings
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 import torch
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
@@ -13,6 +13,11 @@
 
 from colossalai.pipeline.stage_manager import PipelineStageManager
 
+try:
+    from transformers.models.llama.modeling_llama import _prepare_4d_causal_attention_mask
+    LATEST_VERSION = True
+except ImportError:
+    LATEST_VERSION = False
 
 class LlamaPipelineForwards:
     """
@@ -97,9 +102,14 @@ def llama_model_forward(
             attention_mask = torch.ones(
                 (batch_size, seq_length_with_past), dtype=torch.bool, device=hidden_states.device
             )
-        attention_mask = self._prepare_decoder_attention_mask(
-            attention_mask, (batch_size, seq_length), hidden_states, past_key_values_length
-        )
+        if LATEST_VERSION:
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, (batch_size, seq_length), hidden_states, past_key_values_length
+            )
+        else:
+            attention_mask = self._prepare_decoder_attention_mask(
+                attention_mask, (batch_size, seq_length), hidden_states, past_key_values_length
+            )
 
         if self.gradient_checkpointing and self.training:
             if use_cache:


Test patch:
Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

[BUG]: gemini 插件微调时，疑似不能正确加载模型权重
### 🐛 Describe the bug

# 环境
```
------------ Environment ------------
Colossal-AI version: 0.3.4
PyTorch version: 2.0.1
System CUDA version: 11.7
CUDA version required by PyTorch: 11.7
```

# BUG 细节
微调代码修改自：https://github.com/hpcaitech/ColossalAI/blob/main/examples/language/llama2/finetune.py

加载模型：https://github.com/hpcaitech/ColossalAI/blob/main/examples/language/llama2/finetune.py#L237


除了plugin的类型其余变量都保持一致，发现zero2时，loss的表现正常，而使用gemini时，更像是从一个随机初始化的weight进行优化

zero2，loss 正常从比较低的水平开始下降:
![image](https://github.com/hpcaitech/ColossalAI/assets/26140960/5ce96b4c-b85b-432e-8ed5-2dc58c7982ed)

gemini，loss 从特别高的水平下降：
![image](https://github.com/hpcaitech/ColossalAI/assets/26140960/e14c9afb-a112-44fd-a3ef-ca721d1eefa4)





### Environment

_No response_


Discussion:



 
> Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿      
----  
**Title:** [BUG]: When fine-tuning the gemini plug-in, it is suspected that the model weights cannot be loaded correctly.    
      
感谢您的反馈，我们会尽快看一下这个问题。
 
> Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿         
----    
Thanks for your feedback, we'll take a look at this issue as soon as possible.  
      
When fine-tuning ChatGLM2 / Llama2, I also observed a similar phenomenon, under PyTorch version: 1.13.1. It seems that this phenomenon occurs when using Colossal-AI version >= 0.3.3.

Resolving patch:
diff --git a/examples/language/llama2/finetune.py b/examples/language/llama2/finetune.py
--- a/examples/language/llama2/finetune.py
+++ b/examples/language/llama2/finetune.py
@@ -58,6 +58,7 @@ def tokenize_batch_for_finetune(batch, tokenizer: Optional[LlamaTokenizer] = Non
 
 def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
     dist.all_reduce(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/colossalai/shardformer/modeling/chatglm2_6b/__init__.py b/colossalai/shardformer/modeling/chatglm2_6b/__init__.py
new file mode 100644
diff --git a/examples/inference/benchmark.py b/examples/inference/benchmark.py
--- a/examples/inference/benchmark.py
+++ b/examples/inference/benchmark.py
@@ -95,11 +95,27 @@ def print_details_info(timestamps, model_config, args, whole_end2end):
 
 def benchmark_inference(args):
     if args.model == "toy":
-        model = transformers.LlamaForCausalLM(transformers.LlamaConfig(num_hidden_layers=8))
+        model = transformers.LlamaForCausalLM(transformers.LlamaConfig(num_hidden_layers=4))
     elif args.model == "7b":
-        model = transformers.LlamaForCausalLM(transformers.AutoConfig.from_pretrained("decapoda-research/llama-7b-hf"))
+        model = transformers.LlamaForCausalLM(
+            transformers.LlamaConfig(
+                hidden_size=4096,
+                intermediate_size=11008,
+                num_attention_heads=32,
+                num_hidden_layers=32,
+                num_key_value_heads=32,
+            )
+        )
     elif args.model == "13b":
-        model = transformers.LlamaForCausalLM(transformers.AutoConfig.from_pretrained("decapoda-research/llama-13b-hf"))
+        model = transformers.LlamaForCausalLM(
+            transformers.LlamaConfig(
+                hidden_size=5120,
+                intermediate_size=13824,
+                num_attention_heads=40,
+                num_hidden_layers=40,
+                num_key_value_heads=40,
+            )
+        )
     else:
         raise NotImplementedError
 


Test patch: