fix gather output

Author: flybird11111

colossalai/shardformer/modeling/gpt2.py

@@ -25,7 +25,7 @@
 from colossalai.shardformer.shard import ShardConfig
 
 from ..layer import cross_entropy_1d
-from ..layer._operation import _gather
+from ..layer._operation import gather_forward_split_backward
 
 
 class GPT2PipelineForwards:
@@ -339,7 +339,7 @@ def gpt2_lmhead_model_forward(
                 loss = loss_fct(shift_logits, shift_labels)
 
             if not shard_config.parallel_output:
-                lm_logits = _gather(lm_logits, -1, shard_config.tensor_parallel_process_group)
+                lm_logits = gather_forward_split_backward(lm_logits, -1, shard_config.tensor_parallel_process_group)
 
         if not return_dict:
             output = (lm_logits,) + outputs[1:]
@@ -1089,7 +1089,7 @@ def forward(
                 loss = loss_fct(shift_logits, shift_labels)
 
         if not shard_config.parallel_output:
-            lm_logits = _gather(lm_logits, -1, shard_config.tensor_parallel_process_group)
+            lm_logits = gather_forward_split_backward(lm_logits, -1, shard_config.tensor_parallel_process_group)
 
         if not return_dict:
             output = (lm_logits,) + transformer_outputs[1:]

colossalai/shardformer/modeling/llama.py

@@ -16,7 +16,7 @@
 from colossalai.shardformer.shard import ShardConfig
 
 from ..layer import cross_entropy_1d
-from ..layer._operation import _gather
+from ..layer._operation import gather_forward_split_backward
 
 try:
     from transformers.models.llama.modeling_llama import _prepare_4d_causal_attention_mask
@@ -290,7 +290,7 @@ def llama_for_causal_lm_forward(
                     loss = loss_fct(shift_logits, shift_labels)
 
             if not shard_config.parallel_output:
-                logits = _gather(logits, -1, shard_config.tensor_parallel_process_group)
+                logits = gather_forward_split_backward(logits, -1, shard_config.tensor_parallel_process_group)
 
             if not return_dict:
                 output = (logits,) + outputs[1:]
@@ -594,7 +594,7 @@ def forward(
                 loss = loss_fct(shift_logits, shift_labels)
 
         if not shard_config.parallel_output:
-            logits = _gather(logits, -1, shard_config.tensor_parallel_process_group)
+            logits = gather_forward_split_backward(logits, -1, shard_config.tensor_parallel_process_group)
 
         if not return_dict:
             output = (logits,) + outputs[1:]