fix: swap tensor bug (NVIDIA#683)

PerkzZheng · azahed98 · commit 42465ca5f4cd · 2023-07-20T10:49:26.000-07:00
diff --git a/src/fastertransformer/layers/TensorParallelGeluFfnLayer.cc b/src/fastertransformer/layers/TensorParallelGeluFfnLayer.cc
@@ -45,6 +45,7 @@ void TensorParallelGeluFfnLayer<T>::forward(TensorMap*          output_tensors,
     if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
         use_custom_all_reduce_kernel =
             custom_all_reduce_comm_->swapInternalBuffer(&swap_tensors, token_num * hidden_units);
+        output_tensors->at("ffn_output").data = swap_tensors[0].data;
     }
 
     GeluFfnLayer<T>::forward(output_tensors, input_tensors, ffn_weights);
@@ -57,6 +58,7 @@ void TensorParallelGeluFfnLayer<T>::forward(TensorMap*          output_tensors,
         }
         else {
             custom_all_reduce_comm_->customAllReduce(token_num * hidden_units, GeluFfnLayer<T>::stream_);
+            output_tensors->at("ffn_output").data = swap_tensors[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/layers/TensorParallelReluFfnLayer.cc b/src/fastertransformer/layers/TensorParallelReluFfnLayer.cc
@@ -45,6 +45,7 @@ void TensorParallelReluFfnLayer<T>::forward(TensorMap*          output_tensors,
     if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
         use_custom_all_reduce_kernel =
             custom_all_reduce_comm_->swapInternalBuffer(&swap_tensors, token_num * hidden_units);
+        output_tensors->at("ffn_output").data = swap_tensors[0].data;
     }
 
     ReluFfnLayer<T>::forward(output_tensors, input_tensors, ffn_weights);
@@ -57,6 +58,7 @@ void TensorParallelReluFfnLayer<T>::forward(TensorMap*          output_tensors,
         }
         else {
             custom_all_reduce_comm_->customAllReduce(token_num * hidden_units, ReluFfnLayer<T>::stream_);
+            output_tensors->at("ffn_output").data = swap_tensors[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/layers/TensorParallelSiluFfnLayer.cc b/src/fastertransformer/layers/TensorParallelSiluFfnLayer.cc
@@ -44,6 +44,7 @@ void TensorParallelSiluFfnLayer<T>::forward(TensorMap*          output_tensors,
     if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
         use_custom_all_reduce_kernel =
             custom_all_reduce_comm_->swapInternalBuffer(&swap_tensors, token_num * hidden_units);
+        output_tensors->at("ffn_output").data = swap_tensors[0].data;
     }
 
     SiluFfnLayer<T>::forward(output_tensors, input_tensors, ffn_weights);
@@ -55,6 +56,7 @@ void TensorParallelSiluFfnLayer<T>::forward(TensorMap*          output_tensors,
         }
         else {
             custom_all_reduce_comm_->customAllReduce(token_num * hidden_units, SiluFfnLayer<T>::stream_);
+            output_tensors->at("ffn_output").data = swap_tensors[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/layers/attention_layers/TensorParallelDecoderCrossAttentionLayer.cc b/src/fastertransformer/layers/attention_layers/TensorParallelDecoderCrossAttentionLayer.cc
@@ -104,11 +104,12 @@ void TensorParallelDecoderCrossAttentionLayer<T>::forward(TensorMap*
     //      value_cache [batch, head_num, max_seq_len, size_per_head]
 
     const size_t size = output_tensors->at("hidden_features").size();
+    std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};
 
     bool use_custom_all_reduce_kernel = false;
     if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
-        std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};
         use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&reduce_tensor, size);
+        output_tensors->at("hidden_features").data = reduce_tensor[0].data;
     }
 
     DecoderCrossAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);
@@ -121,6 +122,7 @@ void TensorParallelDecoderCrossAttentionLayer<T>::forward(TensorMap*
         }
         else {
             custom_all_reduce_comm_->customAllReduce(size, DecoderCrossAttentionLayer<T>::stream_);
+            output_tensors->at("hidden_features").data = reduce_tensor[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/layers/attention_layers/TensorParallelDecoderSelfAttentionLayer.cc b/src/fastertransformer/layers/attention_layers/TensorParallelDecoderSelfAttentionLayer.cc
@@ -200,11 +200,12 @@ void TensorParallelDecoderSelfAttentionLayer<T>::forward(TensorMap*
     //      value_cache [batch, head_num, max_seq_len, size_per_head]
 
     const size_t size = output_tensors->at("hidden_features").size();
+    std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};
 
     bool use_custom_all_reduce_kernel = false;
     if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr && do_all_reduce_) {
-        std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};
         use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&reduce_tensor, size);
+        output_tensors->at("hidden_features").data = reduce_tensor[0].data;
     }
 
     DecoderSelfAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);
@@ -217,6 +218,7 @@ void TensorParallelDecoderSelfAttentionLayer<T>::forward(TensorMap*
         }
         else {
             custom_all_reduce_comm_->customAllReduce(size, DecoderSelfAttentionLayer<T>::stream_);
+            output_tensors->at("hidden_features").data = reduce_tensor[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/layers/attention_layers/TensorParallelDisentangledAttentionLayer.cc b/src/fastertransformer/layers/attention_layers/TensorParallelDisentangledAttentionLayer.cc
@@ -35,11 +35,12 @@ void TensorParallelDisentangledAttentionLayer<T>::forward(TensorMap*
     // For more information, please refer to DisentangledAttentionLayer
 
     const size_t size = output_tensors->at("hidden_features").size();
+    std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};
 
     bool use_custom_all_reduce_kernel = false;
     if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
-        std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};
         use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&hidden_features_reduce, size);
+        output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;
     }
 
     DisentangledAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);
@@ -52,6 +53,7 @@ void TensorParallelDisentangledAttentionLayer<T>::forward(TensorMap*
         }
         else {
             custom_all_reduce_comm_->customAllReduce(size, DisentangledAttentionLayer<T>::stream_);
+            output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/layers/attention_layers/TensorParallelGptContextAttentionLayer.cc b/src/fastertransformer/layers/attention_layers/TensorParallelGptContextAttentionLayer.cc
@@ -35,11 +35,12 @@ void TensorParallelGptContextAttentionLayer<T>::forward(TensorMap*
     //      value_cache [batch, local_head_num, max_seq_len, size_per_head]
 
     const size_t size = output_tensors->at("hidden_features").size();
+    std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};
 
     bool use_custom_all_reduce_kernel = false;
     if (do_all_reduce_ && enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
-        std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};
         use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&reduce_tensor, size);
+        output_tensors->at("hidden_features").data = reduce_tensor[0].data;
     }
 
     GptContextAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);
@@ -52,6 +53,7 @@ void TensorParallelGptContextAttentionLayer<T>::forward(TensorMap*
         }
         else {
             custom_all_reduce_comm_->customAllReduce(size, GptContextAttentionLayer<T>::stream_);
+            output_tensors->at("hidden_features").data = reduce_tensor[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/layers/attention_layers/TensorParallelUnfusedAttentionLayer.cc b/src/fastertransformer/layers/attention_layers/TensorParallelUnfusedAttentionLayer.cc
@@ -37,11 +37,12 @@ void TensorParallelUnfusedAttentionLayer<T>::forward(TensorMap*                o
     // For more information, please refer to UnfusedAttentionLayer
 
     const size_t size = output_tensors->at("hidden_features").size();
+    std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};
 
     bool use_custom_all_reduce_kernel = false;
     if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
-        std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};
         use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&hidden_features_reduce, size);
+        output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;
     }
 
     UnfusedAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);
@@ -53,6 +54,7 @@ void TensorParallelUnfusedAttentionLayer<T>::forward(TensorMap*                o
         }
         else {
             custom_all_reduce_comm_->customAllReduce(size, UnfusedAttentionLayer<T>::stream_);
+            output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;
         }
         sync_check_cuda_error();
     }
diff --git a/src/fastertransformer/models/bert/Bert.cc b/src/fastertransformer/models/bert/Bert.cc
@@ -510,10 +510,11 @@ void Bert<T>::forward(TensorMap* output_tensors, TensorMap* input_tensors, const
                       Tensor{MEMORY_GPU, data_type, std::vector<size_t>{h_token_num, hidden_units_}, attn_out_buf_}}});
 
                 bool use_custom_all_reduce_kernel = false;
+                std::vector<Tensor> hidden_features{attn_output_tensors.at("hidden_features")};
                 if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {
-                    std::vector<Tensor> hidden_features{attn_output_tensors.at("hidden_features")};
                     use_custom_all_reduce_kernel =
                         custom_all_reduce_comm_->swapInternalBuffer(&hidden_features, h_token_num * hidden_units_);
+                    attn_output_tensors.at("hidden_features").data = hidden_features[0].data;
                 }
 
                 if (attention_type == AttentionType::FUSED_MHA || attention_type == AttentionType::FUSED_PADDED_MHA) {
@@ -535,6 +536,7 @@ void Bert<T>::forward(TensorMap* output_tensors, TensorMap* input_tensors, const
                     }
                     else {
                         custom_all_reduce_comm_->customAllReduce(h_token_num * hidden_units_, stream_);
+                        attn_output_tensors.at("hidden_features").data = hidden_features[0].data;
                     }
                     sync_check_cuda_error();
                 }

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,7 @@ void TensorParallelGeluFfnLayer<T>::forward(TensorMap* output_tensors,`
`45`	`45`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {`
`46`	`46`	`use_custom_all_reduce_kernel =`
`47`	`47`	`custom_all_reduce_comm_->swapInternalBuffer(&swap_tensors, token_num * hidden_units);`
	`48`	`+ output_tensors->at("ffn_output").data = swap_tensors[0].data;`
`48`	`49`	`}`
`49`	`50`
`50`	`51`	`GeluFfnLayer<T>::forward(output_tensors, input_tensors, ffn_weights);`
`@@ -57,6 +58,7 @@ void TensorParallelGeluFfnLayer<T>::forward(TensorMap* output_tensors,`
`57`	`58`	`}`
`58`	`59`	`else {`
`59`	`60`	`custom_all_reduce_comm_->customAllReduce(token_num * hidden_units, GeluFfnLayer<T>::stream_);`
	`61`	`+ output_tensors->at("ffn_output").data = swap_tensors[0].data;`
`60`	`62`	`}`
`61`	`63`	`sync_check_cuda_error();`
`62`	`64`	`}`
Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,7 @@ void TensorParallelReluFfnLayer<T>::forward(TensorMap* output_tensors,`
`45`	`45`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {`
`46`	`46`	`use_custom_all_reduce_kernel =`
`47`	`47`	`custom_all_reduce_comm_->swapInternalBuffer(&swap_tensors, token_num * hidden_units);`
	`48`	`+ output_tensors->at("ffn_output").data = swap_tensors[0].data;`
`48`	`49`	`}`
`49`	`50`
`50`	`51`	`ReluFfnLayer<T>::forward(output_tensors, input_tensors, ffn_weights);`
`@@ -57,6 +58,7 @@ void TensorParallelReluFfnLayer<T>::forward(TensorMap* output_tensors,`
`57`	`58`	`}`
`58`	`59`	`else {`
`59`	`60`	`custom_all_reduce_comm_->customAllReduce(token_num * hidden_units, ReluFfnLayer<T>::stream_);`
	`61`	`+ output_tensors->at("ffn_output").data = swap_tensors[0].data;`
`60`	`62`	`}`
`61`	`63`	`sync_check_cuda_error();`
`62`	`64`	`}`
Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,7 @@ void TensorParallelSiluFfnLayer<T>::forward(TensorMap* output_tensors,`
`44`	`44`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {`
`45`	`45`	`use_custom_all_reduce_kernel =`
`46`	`46`	`custom_all_reduce_comm_->swapInternalBuffer(&swap_tensors, token_num * hidden_units);`
	`47`	`+ output_tensors->at("ffn_output").data = swap_tensors[0].data;`
`47`	`48`	`}`
`48`	`49`
`49`	`50`	`SiluFfnLayer<T>::forward(output_tensors, input_tensors, ffn_weights);`
`@@ -55,6 +56,7 @@ void TensorParallelSiluFfnLayer<T>::forward(TensorMap* output_tensors,`
`55`	`56`	`}`
`56`	`57`	`else {`
`57`	`58`	`custom_all_reduce_comm_->customAllReduce(token_num * hidden_units, SiluFfnLayer<T>::stream_);`
	`59`	`+ output_tensors->at("ffn_output").data = swap_tensors[0].data;`
`58`	`60`	`}`
`59`	`61`	`sync_check_cuda_error();`
`60`	`62`	`}`
Original file line number	Diff line number	Diff line change
`@@ -104,11 +104,12 @@ void TensorParallelDecoderCrossAttentionLayer<T>::forward(TensorMap*`
`104`	`104`	`// value_cache [batch, head_num, max_seq_len, size_per_head]`
`105`	`105`
`106`	`106`	`const size_t size = output_tensors->at("hidden_features").size();`
	`107`	`+ std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};`
`107`	`108`
`108`	`109`	`bool use_custom_all_reduce_kernel = false;`
`109`	`110`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {`
`110`		`- std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};`
`111`	`111`	`use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&reduce_tensor, size);`
	`112`	`+ output_tensors->at("hidden_features").data = reduce_tensor[0].data;`
`112`	`113`	`}`
`113`	`114`
`114`	`115`	`DecoderCrossAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);`
`@@ -121,6 +122,7 @@ void TensorParallelDecoderCrossAttentionLayer<T>::forward(TensorMap*`
`121`	`122`	`}`
`122`	`123`	`else {`
`123`	`124`	`custom_all_reduce_comm_->customAllReduce(size, DecoderCrossAttentionLayer<T>::stream_);`
	`125`	`+ output_tensors->at("hidden_features").data = reduce_tensor[0].data;`
`124`	`126`	`}`
`125`	`127`	`sync_check_cuda_error();`
`126`	`128`	`}`
Original file line number	Diff line number	Diff line change
`@@ -200,11 +200,12 @@ void TensorParallelDecoderSelfAttentionLayer<T>::forward(TensorMap*`
`200`	`200`	`// value_cache [batch, head_num, max_seq_len, size_per_head]`
`201`	`201`
`202`	`202`	`const size_t size = output_tensors->at("hidden_features").size();`
	`203`	`+ std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};`
`203`	`204`
`204`	`205`	`bool use_custom_all_reduce_kernel = false;`
`205`	`206`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr && do_all_reduce_) {`
`206`		`- std::vector<Tensor> reduce_tensor{output_tensors->at("hidden_features")};`
`207`	`207`	`use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&reduce_tensor, size);`
	`208`	`+ output_tensors->at("hidden_features").data = reduce_tensor[0].data;`
`208`	`209`	`}`
`209`	`210`
`210`	`211`	`DecoderSelfAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);`
`@@ -217,6 +218,7 @@ void TensorParallelDecoderSelfAttentionLayer<T>::forward(TensorMap*`
`217`	`218`	`}`
`218`	`219`	`else {`
`219`	`220`	`custom_all_reduce_comm_->customAllReduce(size, DecoderSelfAttentionLayer<T>::stream_);`
	`221`	`+ output_tensors->at("hidden_features").data = reduce_tensor[0].data;`
`220`	`222`	`}`
`221`	`223`	`sync_check_cuda_error();`
`222`	`224`	`}`
Original file line number	Diff line number	Diff line change
`@@ -35,11 +35,12 @@ void TensorParallelDisentangledAttentionLayer<T>::forward(TensorMap*`
`35`	`35`	`// For more information, please refer to DisentangledAttentionLayer`
`36`	`36`
`37`	`37`	`const size_t size = output_tensors->at("hidden_features").size();`
	`38`	`+ std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};`
`38`	`39`
`39`	`40`	`bool use_custom_all_reduce_kernel = false;`
`40`	`41`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {`
`41`		`- std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};`
`42`	`42`	`use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&hidden_features_reduce, size);`
	`43`	`+ output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;`
`43`	`44`	`}`
`44`	`45`
`45`	`46`	`DisentangledAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);`
`@@ -52,6 +53,7 @@ void TensorParallelDisentangledAttentionLayer<T>::forward(TensorMap*`
`52`	`53`	`}`
`53`	`54`	`else {`
`54`	`55`	`custom_all_reduce_comm_->customAllReduce(size, DisentangledAttentionLayer<T>::stream_);`
	`56`	`+ output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;`
`55`	`57`	`}`
`56`	`58`	`sync_check_cuda_error();`
`57`	`59`	`}`
Original file line number	Diff line number	Diff line change
`@@ -37,11 +37,12 @@ void TensorParallelUnfusedAttentionLayer<T>::forward(TensorMap* o`
`37`	`37`	`// For more information, please refer to UnfusedAttentionLayer`
`38`	`38`
`39`	`39`	`const size_t size = output_tensors->at("hidden_features").size();`
	`40`	`+ std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};`
`40`	`41`
`41`	`42`	`bool use_custom_all_reduce_kernel = false;`
`42`	`43`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {`
`43`		`- std::vector<Tensor> hidden_features_reduce = {output_tensors->at("hidden_features")};`
`44`	`44`	`use_custom_all_reduce_kernel = custom_all_reduce_comm_->swapInternalBuffer(&hidden_features_reduce, size);`
	`45`	`+ output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;`
`45`	`46`	`}`
`46`	`47`
`47`	`48`	`UnfusedAttentionLayer<T>::forward(output_tensors, input_tensors, attention_weights);`
`@@ -53,6 +54,7 @@ void TensorParallelUnfusedAttentionLayer<T>::forward(TensorMap* o`
`53`	`54`	`}`
`54`	`55`	`else {`
`55`	`56`	`custom_all_reduce_comm_->customAllReduce(size, UnfusedAttentionLayer<T>::stream_);`
	`57`	`+ output_tensors->at("hidden_features").data = hidden_features_reduce[0].data;`
`56`	`58`	`}`
`57`	`59`	`sync_check_cuda_error();`
`58`	`60`	`}`
Original file line number	Diff line number	Diff line change
`@@ -510,10 +510,11 @@ void Bert<T>::forward(TensorMap* output_tensors, TensorMap* input_tensors, const`
`510`	`510`	`Tensor{MEMORY_GPU, data_type, std::vector<size_t>{h_token_num, hidden_units_}, attn_out_buf_}}});`
`511`	`511`
`512`	`512`	`bool use_custom_all_reduce_kernel = false;`
	`513`	`+ std::vector<Tensor> hidden_features{attn_output_tensors.at("hidden_features")};`
`513`	`514`	`if (enable_custom_all_reduce_ && custom_all_reduce_comm_ != nullptr) {`
`514`		`- std::vector<Tensor> hidden_features{attn_output_tensors.at("hidden_features")};`
`515`	`515`	`use_custom_all_reduce_kernel =`
`516`	`516`	`custom_all_reduce_comm_->swapInternalBuffer(&hidden_features, h_token_num * hidden_units_);`
	`517`	`+ attn_output_tensors.at("hidden_features").data = hidden_features[0].data;`
`517`	`518`	`}`
`518`	`519`
`519`	`520`	`if (attention_type == AttentionType::FUSED_MHA \|\| attention_type == AttentionType::FUSED_PADDED_MHA) {`
`@@ -535,6 +536,7 @@ void Bert<T>::forward(TensorMap* output_tensors, TensorMap* input_tensors, const`
`535`	`536`	`}`
`536`	`537`	`else {`
`537`	`538`	`custom_all_reduce_comm_->customAllReduce(h_token_num * hidden_units_, stream_);`
	`539`	`+ attn_output_tensors.at("hidden_features").data = hidden_features[0].data;`
`538`	`540`	`}`
`539`	`541`	`sync_check_cuda_error();`
`540`	`542`	`}`