ra

.coveragerc

@@ -0,0 +1,4 @@
+[run]
+concurrency = multiprocessing
+parallel = true
+sigterm = true

.github/workflows/build_on_pr.yml

@@ -68,9 +68,9 @@ jobs:
     needs: detect
     runs-on: [self-hosted, gpu]
     container:
-      image: hpcaitech/pytorch-cuda:1.11.0-11.3.0
+      image: hpcaitech/pytorch-cuda:1.12.0-11.3.0
       options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10
-    timeout-minutes: 40
+    timeout-minutes: 60
     defaults:
       run:
         shell: bash
@@ -120,15 +120,26 @@ jobs:
           # -p flag is required to preserve the file timestamp to avoid ninja rebuild
           cp -p -r /__w/ColossalAI/ColossalAI/build /github/home/cuda_ext_cache/
 
+      - name: Restore Testmon Cache
+        run: |
+          if [ -d /github/home/testmon_cache ]; then
+            [ ! -z "$(ls -A /github/home/testmon_cache)" ] && cp -p -r /github/home/testmon_cache/.testmondata /__w/ColossalAI/ColossalAI/
+          fi
+
       - name: Execute Unit Testing
         if: needs.detect.outputs.anyLibraryFileChanged == 'true'
         run: |
-          CURL_CA_BUNDLE="" PYTHONPATH=$PWD pytest --cov=. --cov-report xml tests/
+          CURL_CA_BUNDLE="" PYTHONPATH=$PWD pytest --testmon --testmon-cov=. tests/
         env:
           DATA: /data/scratch/cifar-10
           NCCL_SHM_DISABLE: 1
           LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
 
+      - name: Store Testmon Cache
+        run: |
+          [ -d /github/home/testmon_cache ] || mkdir /github/home/testmon_cache
+          cp -p -r /__w/ColossalAI/ColossalAI/.testmondata /github/home/testmon_cache/
+
       - name: Collate artifact
         env:
           PR_NUMBER: ${{ github.event.number }}
@@ -140,7 +151,7 @@ jobs:
           echo $PR_NUMBER > ./report/pr_number
 
           # generate coverage.xml if any
-          if [ "$anyLibraryFileChanged" == "true" ]; then
+          if [ "$anyLibraryFileChanged" == "true" ] && [ -e .coverage ]; then
             allFiles=""
             for file in $changedLibraryFiles; do
               if [ "$allFiles" == "" ]; then

.github/workflows/build_on_schedule.yml

@@ -12,7 +12,7 @@ jobs:
     if: github.repository == 'hpcaitech/ColossalAI'
     runs-on: [self-hosted, 8-gpu]
     container:
-      image: hpcaitech/pytorch-cuda:1.11.0-11.3.0
+      image: hpcaitech/pytorch-cuda:1.12.0-11.3.0
       options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10
     timeout-minutes: 40
     steps:

applications/Chat/examples/community/peft/README.md

@@ -18,7 +18,7 @@ For SFT training, just call train_peft_sft.py
 Its arguments are almost identical to train_sft.py instead adding a new eval_dataset if you have a eval_dataset file. The data file is just a plain datafile, please check the format in the easy_dataset.py.
 
 For stage-3 rlhf training, call train_peft_prompts.py.
-Its arguments are almost idential to train_prompts.py. The only difference is that I use text files to indicate the prompt and pretrained data file. The models are included in easy_models.py. Currently only bloom models are tested, but technically gpt2/opt/llama should be supported.
+Its arguments are almost identical to train_prompts.py. The only difference is that I use text files to indicate the prompt and pretrained data file. The models are included in easy_models.py. Currently only bloom models are tested, but technically gpt2/opt/llama should be supported.
 
 # Dataformat
 Please refer the formats in test_sft.txt, test_prompts.txt, test_pretrained.txt.

applications/Chat/inference/README.md

@@ -75,7 +75,7 @@ E.g. you can set `export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH`.
 
 Please ensure you have downloaded HF-format model weights of LLaMA models first.
 
-Then you can follow [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). This lib provides efficient CUDA kernels and weight convertion script.
+Then you can follow [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). This lib provides efficient CUDA kernels and weight conversion script.
 
 After installing this lib, we may convert the original HF-format LLaMA model weights to 4-bit version.
 

applications/Chat/inference/benchmark.py

@@ -123,7 +123,7 @@ def evaluate(
     start = time()
     for instruction in instructions:
         print(f"Instruction: {instruction}")
-        resp, tokens = evaluate(model, tokenizer, instruction, temparature=0.2, num_beams=1)
+        resp, tokens = evaluate(model, tokenizer, instruction, temperature=0.2, num_beams=1)
         total_tokens += tokens
         print(f"Response: {resp}")
         print('\n----------------------------\n')

colossalai/auto_parallel/README.md

@@ -16,8 +16,8 @@ A *symbolic profiler* for collecting computing and memory overhead related to st
 ### Solver
 
 **Solver** is designed to find the optimal execution plan for a given computation graph and cluster in two stages:
-1) *Intra-op parallelism stage* is to find the plan with the minimum total execution time of all nodes with respect to the constraint of the memory budget. The optimaztion goal of intra-op parallelism solver is modified from <a href="https://arxiv.org/abs/2201.12023"> Alpa </a>'s intra-op parallelsim ILP solver.
-2) *Activation checkpoint stage* is to search for the fastest execution plan that meets the memory budget on the computation graph after inserting the communication nodes by the intra-op parallelism stage. The algorithm to find optimial activation checkpoint is modified from <a href="https://hal.inria.fr/hal-02352969"> Rotor </a>. The reason we use two-stage optimization is that if the two tasks are formulated together, the solving time will be significantly increased, which will greatly affect the user experience of the system. On the contrary, solving in two hierarchical levels has many advantages. Firstly, compared with the computation graph with activation checkpointing, the original graph has fewer nodes, which can reduce the solving cost of intra-op parallelism solver. In addition, a more optimal solution can be found by adding the communication overhead into the activation checkpoint modeling.
+1) *Intra-op parallelism stage* is to find the plan with the minimum total execution time of all nodes with respect to the constraint of the memory budget. The optimization goal of intra-op parallelism solver is modified from <a href="https://arxiv.org/abs/2201.12023"> Alpa </a>'s intra-op parallelism ILP solver.
+2) *Activation checkpoint stage* is to search for the fastest execution plan that meets the memory budget on the computation graph after inserting the communication nodes by the intra-op parallelism stage. The algorithm to find optimal activation checkpoint is modified from <a href="https://hal.inria.fr/hal-02352969"> Rotor </a>. The reason we use two-stage optimization is that if the two tasks are formulated together, the solving time will be significantly increased, which will greatly affect the user experience of the system. On the contrary, solving in two hierarchical levels has many advantages. Firstly, compared with the computation graph with activation checkpointing, the original graph has fewer nodes, which can reduce the solving cost of intra-op parallelism solver. In addition, a more optimal solution can be found by adding the communication overhead into the activation checkpoint modeling.
 
 ### Generator
 **Generator** applies the searched execution plan to the computation graph and recompiles the computation graph to optimized PyTorch code. It has *a series compile pass* to insert a communication node or do the kernel substitution as the intra-op parallelism solver required. Additionally, we implement a *code generation* feature to recognize the annotation from the activation checkpoint solver and inject the activation checkpoint block following annotation instructions.

colossalai/auto_parallel/passes/runtime_preparation_pass.py

@@ -169,7 +169,7 @@ def _post_processing(node, size_processing_node):
         This function is used to process the dependency between the size node and its users after
         inserting the size_process_node.
         '''
-        # store original node and processing node pair in node_pairs dictioanry
+        # store original node and processing node pair in node_pairs dictionary
         # It will be used to replace the original node with processing node in slice object
         node_pairs[node] = size_processing_node
         size_processing_node._meta_data = node._meta_data
@@ -388,7 +388,7 @@ def module_params_sharding_pass(gm: torch.fx.GraphModule, device_mesh: DeviceMes
     """
     mod_graph = gm.graph
     nodes = tuple(mod_graph.nodes)
-    # This stream is created for overlaping the communication and computation.
+    # This stream is created for overlapping the communication and computation.
     reduction_stream = torch.cuda.Stream()
 
     def _add_hook_for_grad_communication(node, param, name=None):

colossalai/autochunk/autochunk_codegen.py

@@ -40,7 +40,7 @@ def _gen_chunk_slice_dim(chunk_dim: int, chunk_indice_name: str, shape: List) ->
     return new_shape
 
 
-def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_ouput_dim: int, chunk_size=2) -> str:
+def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_output_dim: int, chunk_size=2) -> str:
     """
     Generate chunk loop start
 
@@ -52,7 +52,7 @@ def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_oup
     Args:
         chunk_input (List[Node]): chunk input node
         chunk_output (Node): chunk output node
-        chunk_ouput_dim (int): chunk output node chunk dim
+        chunk_output_dim (int): chunk output node chunk dim
         chunk_size (int): chunk size. Defaults to 2.
 
     Returns:
@@ -74,7 +74,7 @@ def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_oup
                                                                                       input_node.name, input_node.name)
 
     out_shape = get_node_shape(chunk_output[0])
-    chunk_shape = out_shape[chunk_ouput_dim[0]]
+    chunk_shape = out_shape[chunk_output_dim[0]]
     context += "chunk_size = %d\nfor chunk_idx in range(0, %d, chunk_size):\n" % (chunk_size, chunk_shape)
     return context
 

colossalai/autochunk/trace_indice.py

@@ -18,7 +18,7 @@ class TraceIndice(object):
         dim(x1)=dim(x2)=dim(x3)=[a, b, c]
     This class will record every node's dims' indice, compute and source.
 
-    Attibutes:
+    Attributes:
         node_list (List)
         indice_trace_list (List): [{"indice": [...], "compute": [...], "source": [...]}, {...}]
         indice_view_list (Dict): not used for now

colossalai/booster/plugin/__init__.py

@@ -4,3 +4,10 @@
 from .torch_ddp_plugin import TorchDDPPlugin
 
 __all__ = ['Plugin', 'TorchDDPPlugin', 'GeminiPlugin', 'LowLevelZeroPlugin']
+
+import torch
+from packaging import version
+
+if version.parse(torch.__version__) >= version.parse('1.12.0'):
+    from .torch_fsdp_plugin import TorchFSDPPlugin
+    __all__.append('TorchFSDPPlugin')