fix potential hang during TP model initialization (#3033)

InternLM · Jan 15, 2025 · 3c660f1 · 3c660f1
1 parent cba13b6
commit 3c660f1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/turbomind/triton_backend/llama/LlamaTritonModel.cc b/src/turbomind/triton_backend/llama/LlamaTritonModel.cc
@@ -347,6 +347,8 @@ LlamaTritonModel<T>::createSharedModelInstance(int
     FT_CHECK(tensor_para.world_size_ == tensor_para_size_);
     FT_CHECK(pipeline_para.world_size_ == pipeline_para_size_);
 
+    shared_state_->barrier->wait();
+
     auto model = std::make_unique<LlamaV2<T>>(model_param_,  //
                                               attn_param_,
                                               moe_param_,
@@ -356,6 +358,8 @@ LlamaTritonModel<T>::createSharedModelInstance(int
                                               engine_param_.max_batch_size,
                                               weights_[device_id]);
 
+    shared_state_->barrier->wait();
+
     auto engine = std::make_unique<Engine<T>>(engine_param_,  //
                                               std::move(model),
                                               std::move(ctx),