[gptq] inference change replace linear to shardformer

[Xu-Kai]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs

🚨 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

📝 What does this PR do?

• Change gptq tensor parallel replace linear to Shardformer.
• Change colossalai/gptq to colossalai/inference/quant/gptq.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

[github-actions]

The code coverage for the changed files is 58%.

Click me to view the complete report

Name                                                                                                      Stmts   Miss  Cover
-----------------------------------------------------------------------------------------------------------------------------
colossalai/amp/naive_amp/mixed_precision_optimizer.py                                                        98     21    79%
colossalai/auto_parallel/offload/base_offload_module.py                                                      71     71     0%
colossalai/auto_parallel/passes/runtime_preparation_pass.py                                                 297    297     0%
colossalai/auto_parallel/tensor_shard/node_handler/registry.py                                               19     19     0%
colossalai/booster/booster.py                                                                                69     12    83%
colossalai/booster/plugin/__init__.py                                                                        11      0   100%
colossalai/booster/plugin/gemini_plugin.py                                                                  123     12    90%
colossalai/booster/plugin/hybrid_parallel_plugin.py                                                         210     14    93%
colossalai/booster/plugin/low_level_zero_plugin.py                                                          152     11    93%
colossalai/booster/plugin/pp_plugin_base.py                                                                   9      1    89%
colossalai/checkpoint_io/__init__.py                                                                          5      0   100%
colossalai/checkpoint_io/general_checkpoint_io.py                                                            91      8    91%
colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py                                                   336     32    90%
colossalai/checkpoint_io/utils.py                                                                           321     46    86%
colossalai/cli/benchmark/models.py                                                                           12     12     0%
colossalai/cli/launcher/run.py                                                                              139    139     0%
colossalai/cluster/__init__.py                                                                                5      0   100%
colossalai/cluster/process_group_mesh.py                                                                     73      1    99%
colossalai/context/parallel_context.py                                                                      227     19    92%
colossalai/context/process_group_initializer/initializer_1d.py                                               29      0   100%
colossalai/context/process_group_initializer/initializer_2d.py                                               72      1    99%
colossalai/context/process_group_initializer/initializer_2p5d.py                                            136      2    99%
colossalai/context/process_group_initializer/initializer_3d.py                                              155      1    99%
colossalai/context/process_group_initializer/initializer_data.py                                             27      0   100%
colossalai/context/process_group_initializer/initializer_model.py                                            28      0   100%
colossalai/context/process_group_initializer/initializer_pipeline.py                                         26      0   100%
colossalai/context/process_group_initializer/initializer_sequence.py                                         42     29    31%
colossalai/context/process_group_initializer/initializer_tensor.py                                           27      0   100%
colossalai/inference/__init__.py                                                                              0      0   100%
colossalai/inference/tensor_parallel/__init__.py                                                              3      0   100%
colossalai/inference/tensor_parallel/batch_infer_state.py                                                    35      9    74%
colossalai/inference/tensor_parallel/engine.py                                                              173    141    18%
colossalai/inference/tensor_parallel/kvcache_manager.py                                                      56     42    25%
colossalai/inference/tensor_parallel/modeling/__init__.py                                                     3      3     0%
colossalai/inference/tensor_parallel/modeling/bloom.py                                                      214    214     0%
colossalai/inference/tensor_parallel/modeling/llama.py                                                      166    166     0%
colossalai/inference/tensor_parallel/policies/__init__.py                                                     3      3     0%
colossalai/inference/tensor_parallel/policies/bloom.py                                                       43     43     0%
colossalai/inference/tensor_parallel/policies/llama.py                                                       46     46     0%
colossalai/initialize.py                                                                                    188    139    26%
colossalai/interface/__init__.py                                                                              3      0   100%
colossalai/interface/model.py                                                                                14      2    86%
colossalai/interface/optimizer.py                                                                            45      5    89%
colossalai/kernel/__init__.py                                                                                 5      0   100%
colossalai/kernel/cuda_native/__init__.py                                                                     5      0   100%
colossalai/kernel/cuda_native/mha/mem_eff_attn.py                                                            36      6    83%
colossalai/kernel/jit/option.py                                                                              48     40    17%
colossalai/kernel/triton/__init__.py                                                                          5      0   100%
colossalai/kernel/triton/context_attention.py                                                                92     79    14%
colossalai/kernel/triton/copy_kv_cache_dest.py                                                               32     22    31%
colossalai/kernel/triton/fused_layernorm.py                                                                  50     40    20%
colossalai/kernel/triton/rms_norm.py                                                                         42     33    21%
colossalai/kernel/triton/rotary_embedding_kernel.py                                                          37     30    19%
colossalai/kernel/triton/self_attention_nofusion.py                                                          58     47    19%
colossalai/kernel/triton/softmax.py                                                                          57     48    16%
colossalai/kernel/triton/token_attention_kernel.py                                                          152    129    15%
colossalai/lazy/lazy_init.py                                                                                315     44    86%
colossalai/legacy/__init__.py                                                                                 0      0   100%
colossalai/legacy/builder/__init__.py                                                                         2      0   100%
colossalai/legacy/builder/builder.py                                                                         21     16    24%
colossalai/legacy/communication/__init__.py                                                                   5      0   100%
colossalai/legacy/communication/collective.py                                                                92     78    15%
colossalai/legacy/communication/p2p.py                                                                      131    107    18%
colossalai/legacy/communication/p2p_v2.py                                                                   113    113     0%
colossalai/legacy/communication/ring.py                                                                      20     15    25%
colossalai/legacy/communication/utils.py                                                                     64     51    20%
colossalai/legacy/engine/__init__.py                                                                          3      0   100%
colossalai/legacy/engine/_base_engine.py                                                                     90     56    38%
colossalai/legacy/engine/gradient_accumulation/__init__.py                                                   15      7    53%
colossalai/legacy/engine/gradient_accumulation/_gradient_accumulation.py                                    107     70    35%
colossalai/legacy/engine/gradient_handler/__init__.py                                                         7      0   100%
colossalai/legacy/engine/gradient_handler/_base_gradient_handler.py                                           8      3    62%
colossalai/legacy/engine/gradient_handler/_data_parallel_gradient_handler.py                                 10      2    80%
colossalai/legacy/engine/gradient_handler/_moe_gradient_handler.py                                           20      9    55%
colossalai/legacy/engine/gradient_handler/_pipeline_parallel_gradient_handler.py                             24     14    42%
colossalai/legacy/engine/gradient_handler/_sequence_parallel_gradient_handler.py                             10      2    80%
colossalai/legacy/engine/gradient_handler/_zero_gradient_handler.py                                           6      1    83%
colossalai/legacy/engine/gradient_handler/utils.py                                                           21     16    24%
colossalai/legacy/engine/schedule/__init__.py                                                                 4      0   100%
colossalai/legacy/engine/schedule/_base_schedule.py                                                          76     57    25%
colossalai/legacy/engine/schedule/_non_pipeline_schedule.py                                                  30     22    27%
colossalai/legacy/engine/schedule/_pipeline_schedule.py                                                     430    397     8%
colossalai/legacy/engine/schedule/_pipeline_schedule_v2.py                                                   78     78     0%
colossalai/legacy/nn/__init__.py                                                                              4      0   100%
colossalai/legacy/nn/_ops/__init__.py                                                                         9      0   100%
colossalai/legacy/nn/_ops/_utils.py                                                                         155     93    40%
colossalai/legacy/nn/_ops/addmm.py                                                                           45     36    20%
colossalai/legacy/nn/_ops/batch_norm.py                                                                      17     10    41%
colossalai/legacy/nn/_ops/element_wise.py                                                                   217     10    95%
colossalai/legacy/nn/_ops/embedding.py                                                                       47     37    21%
colossalai/legacy/nn/_ops/embedding_bag.py                                                                   32     22    31%
colossalai/legacy/nn/_ops/layernorm.py                                                                       14      7    50%
colossalai/legacy/nn/_ops/linear.py                                                                          94     79    16%
colossalai/legacy/nn/_ops/loss.py                                                                            26     17    35%
colossalai/legacy/nn/_ops/view.py                                                                            63     50    21%
colossalai/legacy/nn/layer/__init__.py                                                                        9      0   100%
colossalai/legacy/nn/layer/base_layer.py                                                                     36     22    39%
colossalai/legacy/nn/layer/colossalai_layer/__init__.py                                                       6      0   100%
colossalai/legacy/nn/layer/colossalai_layer/_utils.py                                                        30     18    40%
colossalai/legacy/nn/layer/colossalai_layer/dropout.py                                                       17      9    47%
colossalai/legacy/nn/layer/colossalai_layer/embedding.py                                                     30     11    63%
colossalai/legacy/nn/layer/colossalai_layer/linear.py                                                        32     12    62%
colossalai/legacy/nn/layer/colossalai_layer/normalization.py                                                 17      5    71%
colossalai/legacy/nn/layer/parallel_1d/__init__.py                                                            2      0   100%
colossalai/legacy/nn/layer/parallel_1d/_operation.py                                                         53     35    34%
colossalai/legacy/nn/layer/parallel_1d/_utils.py                                                             96     52    46%
colossalai/legacy/nn/layer/parallel_1d/layers.py                                                            474    378    20%
colossalai/legacy/nn/layer/parallel_2d/__init__.py                                                            3      0   100%
colossalai/legacy/nn/layer/parallel_2d/_operation.py                                                        396    309    22%
colossalai/legacy/nn/layer/parallel_2d/_utils.py                                                             12      7    42%
colossalai/legacy/nn/layer/parallel_2d/layers.py                                                            484    407    16%
colossalai/legacy/nn/layer/parallel_2p5d/__init__.py                                                          3      0   100%
colossalai/legacy/nn/layer/parallel_2p5d/_operation.py                                                      431    336    22%
colossalai/legacy/nn/layer/parallel_2p5d/_utils.py                                                           14      9    36%
colossalai/legacy/nn/layer/parallel_2p5d/layers.py                                                          477    401    16%
colossalai/legacy/nn/layer/parallel_3d/__init__.py                                                            3      0   100%
colossalai/legacy/nn/layer/parallel_3d/_operation.py                                                        237    164    31%
colossalai/legacy/nn/layer/parallel_3d/_utils.py                                                             63     39    38%
colossalai/legacy/nn/layer/parallel_3d/layers.py                                                            513    431    16%
colossalai/legacy/nn/layer/parallel_sequence/__init__.py                                                      3      0   100%
colossalai/legacy/nn/layer/parallel_sequence/_operation.py                                                   82     60    27%
colossalai/legacy/nn/layer/parallel_sequence/_utils.py                                                        9      7    22%
colossalai/legacy/nn/layer/parallel_sequence/layers.py                                                       92     70    24%
colossalai/legacy/nn/layer/utils/__init__.py                                                                  2      0   100%
colossalai/legacy/nn/layer/utils/common.py                                                                   48     18    62%
colossalai/legacy/nn/layer/vanilla/__init__.py                                                                2      0   100%
colossalai/legacy/nn/layer/vanilla/layers.py                                                                147    103    30%
colossalai/legacy/nn/layer/wrapper/__init__.py                                                                2      0   100%
colossalai/legacy/nn/layer/wrapper/pipeline_wrapper.py                                                       38     28    26%
colossalai/legacy/nn/loss/__init__.py                                                                        23      9    61%
colossalai/legacy/nn/loss/loss_1d.py                                                                         59     41    31%
colossalai/legacy/nn/loss/loss_2d.py                                                                         74     48    35%
colossalai/legacy/nn/loss/loss_2p5d.py                                                                       74     48    35%
colossalai/legacy/nn/loss/loss_3d.py                                                                         81     55    32%
colossalai/legacy/nn/metric/__init__.py                                                                      16      6    62%
colossalai/legacy/nn/metric/_utils.py                                                                         5      3    40%
colossalai/legacy/nn/metric/accuracy_2d.py                                                                   13      6    54%
colossalai/legacy/nn/metric/accuracy_2p5d.py                                                                 13      6    54%
colossalai/legacy/nn/metric/accuracy_3d.py                                                                   18      9    50%
colossalai/legacy/nn/parallel/__init__.py                                                                     2      2     0%
colossalai/legacy/nn/parallel/data_parallel.py                                                               99     99     0%
colossalai/legacy/nn/parallel/layers/__init__.py                                                              6      6     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/__init__.py                                              8      8     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/base_embedding.py                                       20     20     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/cache_mgr.py                                           294    294     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/cached_embedding.py                                     64     64     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/copyer.py                                               25     25     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/embedding_config.py                                     10     10     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/parallel_cached_embedding.py                            55     55     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise.py                  97     97     0%
colossalai/legacy/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise_split_cache.py      73     73     0%
colossalai/legacy/nn/parallel/layers/colo_module.py                                                          30     30     0%
colossalai/legacy/nn/parallel/layers/embedding.py                                                            15     15     0%
colossalai/legacy/nn/parallel/layers/linear.py                                                               15     15     0%
colossalai/legacy/nn/parallel/layers/module_utils.py                                                         82     82     0%
colossalai/legacy/nn/parallel/reducer.py                                                                     77     77     0%
colossalai/legacy/registry/__init__.py                                                                       17      0   100%
colossalai/legacy/registry/registry.py                                                                       31     13    58%
colossalai/legacy/trainer/__init__.py                                                                         2      2     0%
colossalai/legacy/trainer/_trainer.py                                                                       170    170     0%
colossalai/legacy/trainer/hooks/__init__.py                                                                   6      6     0%
colossalai/legacy/trainer/hooks/_base_hook.py                                                                34     34     0%
colossalai/legacy/trainer/hooks/_checkpoint_hook.py                                                          31     31     0%
colossalai/legacy/trainer/hooks/_commons_.py                                                                  7      7     0%
colossalai/legacy/trainer/hooks/_log_hook.py                                                                145    145     0%
colossalai/legacy/trainer/hooks/_lr_scheduler_hook.py                                                        21     21     0%
colossalai/legacy/trainer/hooks/_metric_hook.py                                                             227    227     0%
colossalai/logging/logger.py                                                                                 82     27    67%
colossalai/nn/__init__.py                                                                                     5      0   100%
colossalai/nn/layer/__init__.py                                                                               2      0   100%
colossalai/nn/layer/utils.py                                                                                  4      0   100%
colossalai/nn/loss/__init__.py                                                                                1      0   100%
colossalai/nn/loss/loss_moe.py                                                                               24     12    50%
colossalai/nn/lr_scheduler/cosine.py                                                                         25     15    40%
colossalai/nn/lr_scheduler/linear.py                                                                         10      6    40%
colossalai/nn/lr_scheduler/multistep.py                                                                      13      6    54%
colossalai/nn/lr_scheduler/onecycle.py                                                                        5      2    60%
colossalai/nn/lr_scheduler/poly.py                                                                           18     10    44%
colossalai/nn/lr_scheduler/torch.py                                                                          16      4    75%
colossalai/nn/optimizer/cpu_adam.py                                                                          66      4    94%
colossalai/nn/optimizer/fused_adam.py                                                                        53     10    81%
colossalai/nn/optimizer/fused_lamb.py                                                                        80     74     8%
colossalai/nn/optimizer/fused_sgd.py                                                                         58     50    14%
colossalai/nn/optimizer/hybrid_adam.py                                                                       61      4    93%
colossalai/nn/optimizer/lamb.py                                                                              52     47    10%
colossalai/nn/optimizer/lars.py                                                                              53     46    13%
colossalai/pipeline/p2p.py                                                                                   96      7    93%
colossalai/pipeline/pipelinable.py                                                                          162    135    17%
colossalai/pipeline/schedule/__init__.py                                                                      3      0   100%
colossalai/pipeline/schedule/_utils.py                                                                       75      7    91%
colossalai/pipeline/schedule/base.py                                                                         10      1    90%
colossalai/pipeline/schedule/interleaved_pp.py                                                              172     11    94%
colossalai/pipeline/schedule/one_f_one_b.py                                                                 141      7    95%
colossalai/pipeline/stage_manager.py                                                                         50      0   100%
colossalai/pipeline/utils.py                                                                                183    164    10%
colossalai/shardformer/_utils.py                                                                             54     15    72%
colossalai/shardformer/layer/__init__.py                                                                      8      0   100%
colossalai/shardformer/layer/_operation.py                                                                  298    142    52%
colossalai/shardformer/layer/embedding.py                                                                   130     23    82%
colossalai/shardformer/layer/linear.py                                                                      190     53    72%
colossalai/shardformer/layer/normalization.py                                                                51     10    80%
colossalai/shardformer/layer/parallel_module.py                                                              72     20    72%
colossalai/shardformer/layer/qkv_fused_linear.py                                                            300     75    75%
colossalai/shardformer/layer/utils.py                                                                        84     17    80%
colossalai/shardformer/modeling/bert.py                                                                     486    147    70%
colossalai/shardformer/modeling/blip2.py                                                                     53      1    98%
colossalai/shardformer/modeling/bloom.py                                                                    458    141    69%
colossalai/shardformer/modeling/chatglm2.py                                                                 183     41    78%
colossalai/shardformer/modeling/chatglm2_6b/configuration_chatglm.py                                         30      0   100%
colossalai/shardformer/modeling/chatglm2_6b/modeling_chatglm.py                                             571    239    58%
colossalai/shardformer/modeling/gpt2.py                                                                     398    121    70%
colossalai/shardformer/modeling/jit.py                                                                       19      3    84%
colossalai/shardformer/modeling/llama.py                                                                    218     68    69%
colossalai/shardformer/modeling/opt.py                                                                      284     90    68%
colossalai/shardformer/modeling/sam.py                                                                       94      6    94%
colossalai/shardformer/modeling/t5.py                                                                       297     74    75%
colossalai/shardformer/modeling/vit.py                                                                      148     28    81%
colossalai/shardformer/modeling/whisper.py                                                                  319     69    78%
colossalai/shardformer/policies/auto_policy.py                                                               33      4    88%
colossalai/shardformer/policies/base_policy.py                                                               85      7    92%
colossalai/shardformer/policies/bert.py                                                                     261      0   100%
colossalai/shardformer/policies/blip2.py                                                                     49      2    96%
colossalai/shardformer/policies/bloom.py                                                                    155      2    99%
colossalai/shardformer/policies/chatglm2.py                                                                 107      6    94%
colossalai/shardformer/policies/gpt2.py                                                                     185      1    99%
colossalai/shardformer/policies/llama.py                                                                    121      3    98%
colossalai/shardformer/policies/opt.py                                                                      144     13    91%
colossalai/shardformer/policies/sam.py                                                                       32      0   100%
colossalai/shardformer/policies/t5.py                                                                       181      5    97%
colossalai/shardformer/policies/vit.py                                                                      112      1    99%
colossalai/shardformer/policies/whisper.py                                                                  200      9    96%
colossalai/shardformer/shard/shard_config.py                                                                 42      3    93%
colossalai/shardformer/shard/sharder.py                                                                      96      3    97%
colossalai/shardformer/shard/shardformer.py                                                                  15      0   100%
colossalai/shardformer/shard/utils.py                                                                        11      0   100%
colossalai/tensor/colo_parameter.py                                                                          59     10    83%
colossalai/tensor/colo_tensor.py                                                                             52      9    83%
colossalai/tensor/d_tensor/api.py                                                                           149     24    84%
colossalai/tensor/dist_spec_mgr.py                                                                          118     36    69%
colossalai/tensor/param_op_hook.py                                                                           98     12    88%
colossalai/utils/__init__.py                                                                                  9      0   100%
colossalai/utils/common.py                                                                                  303    171    44%
colossalai/utils/data_sampler/data_parallel_sampler.py                                                       56     13    77%
colossalai/utils/profiler/profiler.py                                                                        62     62     0%
colossalai/utils/profiler/stateful_tensor_mem_extention.py                                                   92     92     0%
colossalai/zero/__init__.py                                                                                   4      0   100%
colossalai/zero/gemini/__init__.py                                                                            7      0   100%
colossalai/zero/gemini/chunk/chunk.py                                                                       314     44    86%
colossalai/zero/gemini/chunk/manager.py                                                                     132     16    88%
colossalai/zero/gemini/chunk/search_utils.py                                                                 86      2    98%
colossalai/zero/gemini/colo_init_context.py                                                                 100     86    14%
colossalai/zero/gemini/gemini_ddp.py                                                                        399     83    79%
colossalai/zero/gemini/gemini_mgr.py                                                                         97      8    92%
colossalai/zero/gemini/gemini_optimizer.py                                                                  392     39    90%
colossalai/zero/gemini/memory_tracer/memory_stats.py                                                         74     35    53%
colossalai/zero/gemini/memory_tracer/runtime_mem_tracer.py                                                   64     45    30%
colossalai/zero/gemini/placement_policy.py                                                                  119     24    80%
colossalai/zero/gemini/utils.py                                                                              58     37    36%
colossalai/zero/legacy/gemini/ophooks/_shard_grad_ophook.py                                                  19     19     0%
colossalai/zero/legacy/gemini/ophooks/_shard_param_ophook.py                                                 33     33     0%
colossalai/zero/legacy/sharded_model/zero_hook.py                                                            73     50    32%
colossalai/zero/low_level/bookkeeping/bucket_store.py                                                        59      0   100%
colossalai/zero/low_level/bookkeeping/gradient_store.py                                                      31      0   100%
colossalai/zero/low_level/low_level_optim.py                                                                340     24    93%
colossalai/zero/wrapper.py                                                                                   36     29    19%
op_builder/gptq.py                                                                                           28      0   100%
op_builder/utils.py                                                                                          94     52    45%
tests/components_to_test/hanging_param_model.py                                                              31      0   100%
tests/components_to_test/inline_op_model.py                                                                  32     19    41%
tests/components_to_test/nested_model.py                                                                     37      0   100%
tests/components_to_test/repeated_computed_layers.py                                                         29      0   100%
tests/components_to_test/simple_net.py                                                                       37      0   100%
tests/kit/model_zoo/transformers/__init__.py                                                                 12      0   100%
tests/kit/model_zoo/transformers/albert.py                                                                   37      0   100%
tests/kit/model_zoo/transformers/bert.py                                                                     50      0   100%
tests/kit/model_zoo/transformers/blip2.py                                                                    21      0   100%
tests/kit/model_zoo/transformers/bloom.py                                                                    36      0   100%
tests/kit/model_zoo/transformers/chatglm2.py                                                                 20      0   100%
tests/kit/model_zoo/transformers/gpt.py                                                                      51      0   100%
tests/kit/model_zoo/transformers/llama.py                                                                    28      2    93%
tests/kit/model_zoo/transformers/opt.py                                                                      31      4    87%
tests/kit/model_zoo/transformers/sam.py                                                                      14      0   100%
tests/kit/model_zoo/transformers/t5.py                                                                       25      0   100%
tests/kit/model_zoo/transformers/vit.py                                                                      24      0   100%
tests/kit/model_zoo/transformers/whisper.py                                                                  23      0   100%
tests/test_booster/test_plugin/test_3d_plugin.py                                                             64      7    89%
tests/test_booster/test_plugin/test_gemini_plugin.py                                                         72      9    88%
tests/test_checkpoint_io/test_gemini_checkpoint_io.py                                                        86      0   100%
tests/test_checkpoint_io/test_gemini_torch_compability.py                                                   116      0   100%
tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py                                        87      0   100%
tests/test_checkpoint_io/test_low_level_zero_checkpoint_io.py                                                58      1    98%
tests/test_checkpoint_io/test_plugins_huggingface_compatibility.py                                           56      1    98%
tests/test_cluster/test_process_group_mesh.py                                                                86      1    99%
tests/test_config/test_load_config.py                                                                         9      0   100%
tests/test_context/test_hybrid_parallel.py                                                                  105     25    76%
tests/test_data/test_cifar10_dataset.py                                                                      14      1    93%
tests/test_data/test_data_parallel_sampler.py                                                                35      1    97%
tests/test_data/test_deterministic_dataloader.py                                                             34      1    97%
tests/test_gptq/test_gptq_linear.py                                                                          93     60    35%
tests/test_infer/test_bloom_infer.py                                                                         39     15    62%
tests/test_infer/test_infer_engine.py                                                                        61     34    44%
tests/test_infer/test_kvcache_manager.py                                                                     44     24    45%
tests/test_infer/test_llama_infer.py                                                                         59     32    46%
tests/test_infer_ops/cuda/test_vllm_rmsnorm.py                                                               42     22    48%
tests/test_infer_ops/cuda/test_vllm_rotary_embedding.py                                                      75     54    28%
tests/test_infer_ops/triton/kernel_utils.py                                                                  21     16    24%
tests/test_infer_ops/triton/test_bloom_context_attention.py                                                  38     22    42%
tests/test_infer_ops/triton/test_copy_kv_dest.py                                                             25     12    52%
tests/test_infer_ops/triton/test_layernorm_triton.py                                                         33     17    48%
tests/test_infer_ops/triton/test_llama_context_attention.py                                                  37     21    43%
tests/test_infer_ops/triton/test_rotary_embedding.py                                                         39     25    36%
tests/test_infer_ops/triton/test_self_attention_nonfusion.py                                                 74     56    24%
tests/test_infer_ops/triton/test_softmax.py                                                                  24     11    54%
tests/test_infer_ops/triton/test_token_attn_1.py                                                             50     35    30%
tests/test_infer_ops/triton/test_token_attn_2.py                                                             41     27    34%
tests/test_infer_ops/triton/test_token_attn_fwd.py                                                           48     34    29%
tests/test_infer_ops/triton/test_token_softmax.py                                                            32     20    38%
tests/test_lazy/test_models.py                                                                               14      1    93%
tests/test_pipeline/test_p2p_communication.py                                                                44      1    98%
tests/test_pipeline/test_pipelinable.py                                                                      42     28    33%
tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py                                            17      0   100%
tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py                                       20      2    90%
tests/test_pipeline/test_schedule/test_interleaved.py                                                        99      1    99%
tests/test_pipeline/test_schedule/test_oneF_oneB.py                                                          80      2    98%
tests/test_pipeline/test_schedule/test_pipeline_schedule_utils.py                                            40      0   100%
tests/test_pipeline/test_stage_manager.py                                                                    45      1    98%
tests/test_shardformer/test_layer/test_embedding.py                                                          37      1    97%
tests/test_shardformer/test_layer/test_gpt2_qkv_fused_linear_1d.py                                           94      1    99%
tests/test_shardformer/test_layer/test_layernorm.py                                                          35      1    97%
tests/test_shardformer/test_layer/test_linear_1d.py                                                         116      1    99%
tests/test_shardformer/test_layer/test_qkv_fused_linear_1d.py                                                89      1    99%
tests/test_shardformer/test_layer/test_vocab_parallel_embedding_1d.py                                        39      1    97%
tests/test_shardformer/test_model/_utils.py                                                                 182     30    84%
tests/test_shardformer/test_model/test_shard_bert.py                                                         81     12    85%
tests/test_shardformer/test_model/test_shard_blip2.py                                                        40      1    98%
tests/test_shardformer/test_model/test_shard_bloom.py                                                        80     12    85%
tests/test_shardformer/test_model/test_shard_chatglm2.py                                                     80     11    86%
tests/test_shardformer/test_model/test_shard_gpt2.py                                                         82     11    87%
tests/test_shardformer/test_model/test_shard_llama.py                                                        83     12    86%
tests/test_shardformer/test_model/test_shard_opt.py                                                          82     11    87%
tests/test_shardformer/test_model/test_shard_sam.py                                                          39      1    97%
tests/test_shardformer/test_model/test_shard_t5.py                                                           80     11    86%
tests/test_shardformer/test_model/test_shard_vit.py                                                          79     11    86%
tests/test_shardformer/test_model/test_shard_whisper.py                                                      88     14    84%
tests/test_shardformer/test_shard_utils.py                                                                   21      0   100%
tests/test_shardformer/test_with_torch_ddp.py                                                                52      1    98%
tests/test_utils/test_activation_checkpointing.py                                                            81      1    99%
tests/test_utils/test_checkpoint/test_checkpoint_1d.py                                                       55     34    38%
tests/test_utils/test_checkpoint/test_checkpoint_2d.py                                                       55     34    38%
tests/test_utils/test_checkpoint/test_checkpoint_2p5d.py                                                     55     34    38%
tests/test_utils/test_checkpoint/test_checkpoint_3d.py                                                       55     34    38%
tests/test_utils/test_flash_attention.py                                                                    119      0   100%
tests/test_utils/test_norm_gradient_clipping.py                                                              60     35    42%
tests/test_zero/test_gemini/test_chunk_mgrv2.py                                                              49      1    98%
tests/test_zero/test_gemini/test_chunkv2.py                                                                  87      1    99%
tests/test_zero/test_gemini/test_fwd_bwd.py                                                                  74      1    99%
tests/test_zero/test_gemini/test_gemini_use_rmt.py                                                           69     46    33%
tests/test_zero/test_gemini/test_grad_clip.py                                                                76      2    97%
tests/test_zero/test_gemini/test_inference.py                                                                92      1    99%
tests/test_zero/test_gemini/test_optim.py                                                                   119      1    99%
tests/test_zero/test_gemini/test_runtime_mem_tracer.py                                                       40     28    30%
tests/test_zero/test_gemini/test_search.py                                                                   38      2    95%
tests/test_zero/test_gemini/test_zeroddp_state_dict.py                                                       91      5    95%
tests/test_zero/test_gemini/test_zerooptim_state_dict.py                                                     59      2    97%
tests/test_zero/test_low_level/test_grad_acc.py                                                              86      1    99%
tests/test_zero/test_low_level/test_zero1_2.py                                                              102      1    99%
tests/test_zero/test_low_level/test_zero_ckpt.py                                                             69      5    93%
tests/test_zero/test_low_level/test_zero_tp.py                                                               67     43    36%
-----------------------------------------------------------------------------------------------------------------------------
TOTAL                                                                                                     30112  12591    58%