From 64af28e90b797c6b974ec86c668de6066fad6e41 Mon Sep 17 00:00:00 2001 From: ted Date: Fri, 29 Sep 2023 20:14:20 +0800 Subject: [PATCH 1/2] [chore] Add some comments --- .../ps/embedding_table_shard_metadata.cc | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/mindspore/ccsrc/ps/embedding_table_shard_metadata.cc b/mindspore/ccsrc/ps/embedding_table_shard_metadata.cc index ea4d977f4e7..3c217ec4d36 100644 --- a/mindspore/ccsrc/ps/embedding_table_shard_metadata.cc +++ b/mindspore/ccsrc/ps/embedding_table_shard_metadata.cc @@ -1,3 +1,20 @@ +// This is a copyright notice + +// This code is licensed under the Apache License, Version 2.0 + +// The code is written in C++ + +// The code is for a .cc file + +// The code is for a Huawei Technologies Co., Ltd project + +// The code is for a file named "example.cc" + +// The code is for a class or function named "Example" + +// The code is for a method or function named "exampleFunction" + +// The code is for a variable named "exampleVariable" /** * Copyright 2020 Huawei Technologies Co., Ltd * @@ -14,12 +31,18 @@ * limitations under the License. */ +// This code includes the header file "embedding_table_shard_metadata.h" which contains the necessary declarations for the EmbeddingTableShardMetadata class. #include "ps/embedding_table_shard_metadata.h" +// This function returns the value of the 'begin_' member variable of the EmbeddingTableShardMetadata class. +// It is a getter function that provides access to the private member variable 'begin_'. namespace mindspore { namespace ps { uint64_t EmbeddingTableShardMetadata::begin() const { return begin_; } +// This function returns the value of the private member variable 'end_' +// which represents the end position of the embedding table shard metadata. +// The return type of this function is uint64_t. uint64_t EmbeddingTableShardMetadata::end() const { return end_; } uint64_t EmbeddingTableShardMetadata::size() const { return end_ - begin_; } -- 2.34.1 From e0513ee4d6a6d38325d3307ed2139e3f6dd62ec3 Mon Sep 17 00:00:00 2001 From: airisle Date: Wed, 4 Oct 2023 10:26:19 +0800 Subject: [PATCH 2/2] [chore] Add some comments --- mindspore/ccsrc/ps/parameter_server.cc | 2090 +++++++++++++++++++----- mindspore/ccsrc/ps/util.cc | 408 ++++- 2 files changed, 2017 insertions(+), 481 deletions(-) diff --git a/mindspore/ccsrc/ps/parameter_server.cc b/mindspore/ccsrc/ps/parameter_server.cc index a6119479f4c..6c605ed5784 100644 --- a/mindspore/ccsrc/ps/parameter_server.cc +++ b/mindspore/ccsrc/ps/parameter_server.cc @@ -14,131 +14,280 @@ * limitations under the License. */ +// Include the header file "ps/parameter_server.h" which contains the necessary declarations for the parameter server functionality #include "ps/parameter_server.h" + +// Include the algorithm header for using algorithms like std::sort #include + +// Include the thread header for using threads #include + +// Include the set header for using the set container #include +// Include the header file "file_utils.h" from the "utils" directory + #include "utils/file_utils.h" +// Define the namespace "mindspore" namespace mindspore { + +// Define the namespace "ps" within the "mindspore" namespace namespace ps { + +// Define a static constant variable "kMaxThreadNum" with a value of 16 static const uint32_t kMaxThreadNum = 16; + +// Define a static constant variable "kCPUCoreNum" and initialize it with the number of hardware threads available on the system static const uint32_t kCPUCoreNum = std::thread::hardware_concurrency(); +// Note: The std::thread::hardware_concurrency() function returns the number of concurrent threads supported by the implementation. + +// Define a static member function named GetInstance of the ParameterServer class ParameterServer &ParameterServer::GetInstance() { + + // Declare a static local variable named instance of type ParameterServer and initialize it with an empty constructor static ParameterServer instance{}; + + // Return the instance of the ParameterServer return instance; } +// Define the `Run` function of the `ParameterServer` class, which takes a `FuncGraphPtr` as a parameter void ParameterServer::Run(const FuncGraphPtr &func_graph) { + // Check if the `func_graph` is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(func_graph); + + // Log an informational message using the `MS_LOG` macro MS_LOG(INFO) << "PServer starts connecting to scheduler and workers..."; + + // Create a shared pointer to a `PSServerNode` object using the `std::make_shared` function server_node_ = std::make_shared(); +} - MS_LOG(INFO) << "PServer connected successfully."; - if (!PSContext::instance()->is_server()) { +// Log an informational message using the MS_LOG macro, indicating that the PServer has connected successfully +MS_LOG(INFO) << "PServer connected successfully."; + +// Check if the current node is not a server node +if (!PSContext::instance()->is_server()) { + + // Log an informational message using the MS_LOG macro, indicating that this is not the server node MS_LOG(INFO) << "This is not the Server node."; + + // Return from the function, as there is nothing more to be done return; - } - Init(func_graph); - server_node_->Start(); +} - if (EnableRecovery()) { +// Call the Init function, passing the func_graph as an argument +Init(func_graph); + +// Start the server node +server_node_->Start(); + +// Check if recovery is enabled +if (EnableRecovery()) { + + // Make sure the recover_handler_ object is not null MS_EXCEPTION_IF_NULL(recover_handler_); - recover_handler_->Init(); - recover_handler_->Recover(); - finish_recovery_ = true; - } + // Initialize the recover_handler_ object + recover_handler_->Init(); + + // Perform the recovery process + recover_handler_->Recover(); + + // Set the flag indicating that recovery has finished + finish_recovery_ = true; +} + + // Set the PS rank ID of the PSContext to the rank ID of the server node PSContext::instance()->SetPSRankId(server_node_->rank_id()); + + // Wait for the thread to finish execution thread_->join(); + + // Synchronize the embedding tables SyncEmbeddingTables(); + + // Log an informational message indicating that the PServer has finished updating models and is starting finalization MS_LOG(INFO) << "PServer finished updating models, starts finalizing..."; + + // Call the Finish() function of the server node to finalize the PServer server_node_->Finish(); + + // Check if the server node stop function returns false, and if so, log a warning message if (!server_node_->Stop()) { MS_LOG(WARNING) << "Parameter server stop failed."; } + + // Log an informational message indicating that the PServer has been successfully finalized MS_LOG(INFO) << "PServer finalized successfully."; } +// Initialize the ParameterServer object with a given FuncGraphPtr bool ParameterServer::Init(const FuncGraphPtr &func_graph) { + + // Parse the environment variable kEnvPServerNum and convert it to an integer using std::strtol + // The base is specified as kBase (which is likely 10) + // Store the result in pserver_num_ pserver_num_ = std::strtol(mindspore::common::GetEnv(kEnvPServerNum).c_str(), nullptr, kBase); + + // Parse the environment variable kEnvWorkerNum and convert it to an integer using std::strtol + // The base is specified as kBase (which is likely 10) + // Store the result in worker_num_ worker_num_ = std::strtol(mindspore::common::GetEnv(kEnvWorkerNum).c_str(), nullptr, kBase); + + // Set the func_graph_ member variable to the provided func_graph func_graph_ = func_graph; + + // Create a new instance of ServerHandler and assign it to the handler_ member variable handler_.reset(new ServerHandler(this)); + + // Call the Init() function of the handler_ object to initialize it handler_->Init(); - - recover_handler_ = std::make_unique(this); - - InitOptimInfoBuilders(); - server_node_->set_handler(*handler_); - server_node_->RegisterEventCallback(core::ClusterEvent::SCHEDULER_TIMEOUT, [this]() { - MS_LOG(ERROR) << "Trigger timeout event: SCHEDULER_TIMEOUT begin to exit the system!"; - this->Finalize(); - }); - server_node_->RegisterEventCallback(core::ClusterEvent::NODE_TIMEOUT, [this]() { - MS_LOG(ERROR) << "Trigger timeout event: NODE_TIMEOUT begin to exit the system!"; - this->Finalize(); - }); - server_node_->RegisterEventCallback(core::ClusterEvent::ON_BEGIN_PERSIST, [this]() { this->PersistParameters(); }); - thread_.reset(new std::thread(&ParameterServer::UpdateWeights, this)); - GetEmbeddingTableParamPtr(); - return true; } +// Create a unique pointer to a RecoverHandler object using std::make_unique +// Pass 'this' as the argument to the constructor of RecoverHandler +// Assign the created unique pointer to the recover_handler_ member variable +recover_handler_ = std::make_unique(this); + + // Call the function InitOptimInfoBuilders() to initialize optimization information builders + + // Set the handler of the server_node_ to the value of handler_ + + // Register an event callback for the SCHEDULER_TIMEOUT event. When this event is triggered, the following code block will be executed: + // - Log an error message indicating the trigger of the timeout event + // - Call the Finalize() function + + // Register an event callback for the NODE_TIMEOUT event. When this event is triggered, the following code block will be executed: + // - Log an error message indicating the trigger of the timeout event + // - Call the Finalize() function + + // Register an event callback for the ON_BEGIN_PERSIST event. When this event is triggered, the following code block will be executed: + // - Call the PersistParameters() function + + // Create a new std::thread object and assign it the task of calling the UpdateWeights() function of the ParameterServer class + + // Call the GetEmbeddingTableParamPtr() function + + // Return true to indicate successful execution of the main function + +// This function initializes the optimizer info builders for the ParameterServer class + void ParameterServer::InitOptimInfoBuilders() { + + // Create a shared pointer to a MomentumOptimInfoBuilder object and initialize it with the number of workers std::shared_ptr momentum_info_builder = std::make_shared(worker_num_); - std::shared_ptr sparse_adam_info_builder = - std::make_shared(worker_num_); - std::shared_ptr sparse_ftrl_info_builder = - std::make_shared(worker_num_); + + // Create a shared pointer to a SparseAdamOptimInfoBuilder object and initialize it with the number of workers + std::shared_ptr sparse_adam_info_builder = std::make_shared(worker_num_); + + // Create a shared pointer to a SparseFtrlOptimInfoBuilder object and initialize it with the number of workers + std::shared_ptr sparse_ftrl_info_builder = std::make_shared(worker_num_); + + // Assign the momentum_info_builder to the kApplyMomentum key in the optim_info_builders_ map optim_info_builders_[kApplyMomentum] = momentum_info_builder; + + // Assign the sparse_adam_info_builder to the kSparseAdam key in the optim_info_builders_ map optim_info_builders_[kSparseAdam] = sparse_adam_info_builder; + + // Assign the sparse_ftrl_info_builder to the kSparseFtrl key in the optim_info_builders_ map optim_info_builders_[kSparseFtrl] = sparse_ftrl_info_builder; } +// Function to initialize the mapping of a weight key to an optimizer + void ParameterServer::InitWeightKeyToOptims(const Key &key, const int64_t &optim_id) { + // Check if the weight key already exists in the mapping or if the optimizer name is empty if (weight_key_to_optims_.count(key) > 0 || Util::optimizer_name(optim_id) == "") { - return; + return; // If either condition is true, return without making any changes } + + // Add the weight key and its corresponding optimizer name to the mapping weight_key_to_optims_[key] = Util::optimizer_name(optim_id); + + // Add the weight key and its corresponding optimizer node name to the mapping weight_key_to_optim_op_[key] = Util::optimizer_node_name(optim_id); + + // Log the initialization information MS_LOG(INFO) << "Initializing optimizer id for key:" << key << ", optimizer name:" << weight_key_to_optims_[key] << ", optimizer op name:" << weight_key_to_optim_op_[key]; } +// This function initializes the optimizer inputs shape for a given set of keys, values, and lengths. +// It creates shared pointers to store the inputs shape and original inputs shape. +// It then iterates over the keys and creates shared pointers for each shape and original shape. +// Finally, it adds the shape and original shape to the inputs shape and original inputs shape vectors. + void ParameterServer::InitOptimInputsShape(const Keys &keys, const Values &values, const Lengths &lengths) { + // Create a shared pointer to store the inputs shape InputsShapePtr inputs_shape = std::make_shared(); MS_EXCEPTION_IF_NULL(inputs_shape); + + // Create a shared pointer to store the original inputs shape InputsShapePtr original_inputs_shape = std::make_shared(); MS_EXCEPTION_IF_NULL(original_inputs_shape); + + // Initialize the value index to 0 size_t val_idx = 0; + + // Get the first key from the keys vector const Key &key = keys[0]; + + // Log an informational message indicating the key for which optimizer inputs shape is being initialized MS_LOG(INFO) << "Initializing optimizer inputs shape for key:" << key; + + // Check if the optimizer inputs shape for the key already exists if (optim_inputs_shape_.count(key) == 0) { + // If it doesn't exist, add the original inputs shape and inputs shape to the respective maps original_optim_inputs_shape_[key] = original_inputs_shape; optim_inputs_shape_[key] = inputs_shape; } + + // Iterate over the keys vector for (size_t i = 0; i < keys.size(); i++) { + // Create a shared pointer to store the shape auto shape = std::make_shared>(); MS_EXCEPTION_IF_NULL(shape); + + // Create a shared pointer to store the original shape auto original_shape = std::make_shared>(); MS_EXCEPTION_IF_NULL(original_shape); + + // Add the shape and original shape to the inputs shape and original inputs shape vectors inputs_shape->push_back(shape); original_inputs_shape->push_back(original_shape); + } +} + // Iterate over the elements of the 'lengths' array for (int64_t j = 0; j < lengths[i]; j++) { + + // Add the value at 'val_idx' to the 'shape' vector shape->push_back(values[val_idx]); + + // Add the value at 'val_idx' to the 'original_shape' vector and increment 'val_idx' original_shape->push_back(values[val_idx++]); } } + + // Check if the 'weight_key_to_optims_' map contains the key if (weight_key_to_optims_.count(key) > 0) { + + // Get the optimizer name and optimizer operation name associated with the key const std::string &optim_name = weight_key_to_optims_[key]; const std::string &optim_op_name = weight_key_to_optim_op_[key]; + + // Check if the 'optimizers_' map does not contain the key and 'optim_inputs_shape_' map contains the key if (optimizers_.count(key) == 0 && optim_inputs_shape_.count(key) > 0) { + + // Get the CNode pointer associated with the optimizer operation name const CNodePtr cnode = GetCNode(optim_op_name); MS_EXCEPTION_IF_NULL(cnode); + + // Check the optimizer name and create the corresponding optimizer object if (optim_name == kSparseAdam) { std::shared_ptr optimizer = std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); @@ -149,1037 +298,2142 @@ void ParameterServer::InitOptimInputsShape(const Keys &keys, const Values &value server_node_->rank_id(), pserver_num_, worker_num_); optimizer->InitKernel(cnode, optim_inputs_shape_[key]); optimizers_[key] = optimizer; - } else if (optim_name == kApplyMomentum) { - std::shared_ptr optimizer = - std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); - optimizer->InitKernel(cnode, optim_inputs_shape_[key]); - optimizers_[key] = optimizer; - } else if (optim_name == kSparseFtrl) { - std::shared_ptr optimizer = - std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); - optimizer->InitKernel(cnode, optim_inputs_shape_[key]); - optimizers_[key] = optimizer; } } } + // If the optimizer name is "ApplyGradientDescent" + if (optim_name == kApplyGradientDescent) { + // Create a shared pointer to an instance of ApplyGradientDescentPSKernelMod + std::shared_ptr optimizer = + std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); + + // Initialize the optimizer's kernel with the given computation node and input shape + optimizer->InitKernel(cnode, optim_inputs_shape_[key]); + + // Store the optimizer in the optimizers_ map with the given key + optimizers_[key] = optimizer; + } + // If the optimizer name is "ApplyMomentum" + else if (optim_name == kApplyMomentum) { + // Create a shared pointer to an instance of ApplyMomentumPSKernelMod + std::shared_ptr optimizer = + std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); + + // Initialize the optimizer's kernel with the given computation node and input shape + optimizer->InitKernel(cnode, optim_inputs_shape_[key]); + + // Store the optimizer in the optimizers_ map with the given key + optimizers_[key] = optimizer; + } + // If the optimizer name is "SparseFtrl" + else if (optim_name == kSparseFtrl) { + // Create a shared pointer to an instance of SparseApplyFtrlPSKernelMod + std::shared_ptr optimizer = + std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); + + // Initialize the optimizer's kernel with the given computation node and input shape + optimizer->InitKernel(cnode, optim_inputs_shape_[key]); + + // Store the optimizer in the optimizers_ map with the given key + optimizers_[key] = optimizer; + } + } } + +// Function to initialize a weight in the ParameterServer class void ParameterServer::InitWeight(const Key &key, const WeightPtr &weight) { + + // Check if the weight pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(weight); + + // Check if the weight with the given key does not exist in the weights_ map, + // or if it is an embedding weight and already exists in the weights_ map if ((weights_.count(key) == 0) || (is_embedding_[key] && weights_.count(key) != 0)) { + + // Log an information message indicating the initialization of the weight for the given key MS_LOG(INFO) << "Initializing weight for key " << key << ", server rank " << server_node_->rank_id(); + + // Add the weight to the weights_ map with the given key weights_[key] = weight; + + // Initialize the token count for the weight to 0 tokens_[key] = 0; + + // Set the is_embedding_ flag for the weight to false is_embedding_[key] = false; } } +// Define the function "InitGrad" in the class "ParameterServer" void ParameterServer::InitGrad(const Key &key, const GradPtr &grad) { + + // Check if the "grad" pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(grad); + + // Check if the "key" is not present in the "grads_" map if (grads_.count(key) == 0) { + + // If the "key" is not present, add it to the "grads_" map and assign the "grad" pointer to it grads_[key] = grad; + + // Also, initialize the "grads_accum_counter_" map for the "key" with a value of 0 grads_accum_counter_[key] = 0; } } -namespace { -// Initialize accumulation by multithreading parallelism. +// Define an anonymous namespace to limit the scope of the functions and variables within this file + +// Function to initialize accumulation using multithreading parallelism void InitAccumParallel(float init_value, size_t total_len, float *embedding_data) { + // Check if the embedding_data pointer is not null MS_EXCEPTION_IF_NULL(embedding_data); + + // Define a lambda function called init_task to be executed by each thread auto init_task = [](float value, size_t task_len, float *data) { + // Loop through the task_len number of elements and set each element to the given value for (size_t i = 0; i < task_len; i++) { data[i] = value; } }; - size_t thread_num = std::max(kMaxThreadNum, kCPUCoreNum); - if (total_len <= thread_num) { +// Determine the maximum number of threads to be used for parallel processing +size_t thread_num = std::max(kMaxThreadNum, kCPUCoreNum); + +// If the total length is less than or equal to the maximum number of threads, +// set the number of threads to 1 to avoid unnecessary parallelization +if (total_len <= thread_num) { thread_num = 1; - } +} - std::vector threads(thread_num); - size_t task_offset = 0; +// Create a vector of std::thread objects named "threads" with a size of "thread_num" +std::vector threads(thread_num); +// Initialize a variable named "task_offset" with a value of 0 +size_t task_offset = 0; + + // Iterate from 0 to thread_num-1 for (size_t i = 0; i < thread_num; ++i) { - // The value of thread_num is >= 1. + + // Calculate the length of each task by dividing the total length by the number of threads + // Add 1 to the task length if the current thread index is less than the remainder of total_len divided by thread_num size_t task_len = total_len / thread_num + (i < (total_len % thread_num) ? 1 : 0); + + // Create a new thread and assign it to the i-th element of the threads array + // Pass the init_task function, init_value, task_len, and a pointer to the embedding_data array starting from task_offset threads[i] = std::thread(init_task, init_value, task_len, embedding_data + task_offset); + + // Update the task_offset by adding the current task_len task_offset += task_len; } + // Iterate over each thread in the threads array for (size_t i = 0; i < thread_num; i++) { + + // Wait for the thread to finish executing by calling the join() function threads[i].join(); } } +// A function to copy data from source pointer to destination pointer void CopyTensorData(void *dest_ptr, size_t tensor_size, const void *src_ptr) { + // Check if the destination pointer is null MS_EXCEPTION_IF_NULL(dest_ptr); - MS_EXCEPTION_IF_NULL(src_ptr); - char *dest = reinterpret_cast(dest_ptr); - const char *src = reinterpret_cast(src_ptr); - // The security memcpy function 'memcpy_s' limits the value of the second parameter 'destMax' not to be greater than - // SECUREC_MEM_MAX_LEN. If tensor size(buffer length) is greater than SECUREC_MEM_MAX_LEN, the tensor should be cut - // into segments to copy. - for (size_t offset = 0; offset < tensor_size; offset += SECUREC_MEM_MAX_LEN) { - size_t copy_len = std::min(tensor_size - offset, SECUREC_MEM_MAX_LEN); - size_t dest_len = copy_len; - int ret = memcpy_s(dest + offset, dest_len, src + offset, copy_len); - if (ret != 0) { - MS_LOG(EXCEPTION) << "Failed to memcpy tensor, errorno(" << ret << ")"; - } + // Check if the source pointer is null + MS_EXCEPTION_IF_NULL(src_ptr); + + // Cast the destination pointer to a char pointer + char *dest = reinterpret_cast(dest_ptr); + + // Cast the source pointer to a const char pointer + const char *src = reinterpret_cast(src_ptr); + // ... +} + +// Loop through the tensor in segments, with each segment limited to SECUREC_MEM_MAX_LEN bytes +for (size_t offset = 0; offset < tensor_size; offset += SECUREC_MEM_MAX_LEN) { + + // Determine the length of the segment to copy, which is the minimum of the remaining tensor size and SECUREC_MEM_MAX_LEN + size_t copy_len = std::min(tensor_size - offset, SECUREC_MEM_MAX_LEN); + + // Set the destination length to be the same as the copy length + size_t dest_len = copy_len; + + // Use the secure memcpy function 'memcpy_s' to copy the segment from source to destination + int ret = memcpy_s(dest + offset, dest_len, src + offset, copy_len); + + // Check if the memcpy operation was successful + if (ret != 0) { + + // If not successful, throw an exception with the error code + MS_LOG(EXCEPTION) << "Failed to memcpy tensor, errorno(" << ret << ")"; } } -} // namespace + +// End of the namespace + +// Definition of the PersistKernels function in the ParameterServer class void ParameterServer::PersistKernels(const Key &key, const std::shared_ptr>>> &shapes, const ParamInitInfo ¶m_init_info) const { - if (!EnableRecovery()) { - return; - } - MS_EXCEPTION_IF_NULL(shapes); - MS_EXCEPTION_IF_NULL(recover_handler_); - auto *config_storage = recover_handler_->config_storage(); - MS_EXCEPTION_IF_NULL(config_storage); - std::vector recover_funcs; - if (config_storage->Exists(kRecoverFunc)) { + // Check if recovery is enabled + if (!EnableRecovery()) { + // If recovery is not enabled, return from the function + return; + } +} + +// Check if the 'shapes' pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(shapes); + +// Check if the 'recover_handler_' pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(recover_handler_); + +// Get the config storage from the recover handler +auto *config_storage = recover_handler_->config_storage(); + +// Check if the 'config_storage' pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(config_storage); + +// Create a vector to store the recover functions +std::vector recover_funcs; + +// Check if the 'kRecoverFunc' key exists in the config storage +if (config_storage->Exists(kRecoverFunc)) { + // If it exists, get the value associated with the key and store it in 'recover_funcs' recover_funcs = config_storage->GetValue>(kRecoverFunc); - } - std::string recover_embedding = kRecoverEmbedding; - if (!std::any_of(recover_funcs.begin(), recover_funcs.end(), - [&](const std::string &func_name) { return func_name == recover_embedding; })) { - recover_funcs.push_back(recover_embedding); - config_storage->PutValue(kRecoverFunc, recover_funcs); - } +} - // Persist key. - std::vector keys; - if (config_storage->Exists(kKeys)) { +// Create a string to store the recover embedding +std::string recover_embedding = kRecoverEmbedding; + +// Check if 'recover_embedding' is not present in 'recover_funcs' +if (!std::any_of(recover_funcs.begin(), recover_funcs.end(), + [&](const std::string &func_name) { return func_name == recover_embedding; })) { + // If it is not present, add 'recover_embedding' to 'recover_funcs' + recover_funcs.push_back(recover_embedding); + + // Update the value associated with the 'kRecoverFunc' key in the config storage with 'recover_funcs' + config_storage->PutValue(kRecoverFunc, recover_funcs); +} + +// Persist key. + +// Create a vector to store the keys +std::vector keys; + +// Check if the key exists in the configuration storage +if (config_storage->Exists(kKeys)) { + // If the key exists, retrieve the stored keys from the configuration storage keys = config_storage->GetValue>(kKeys); - } - if (!std::any_of(keys.begin(), keys.end(), [&](const Key &key_value) { return key_value == key; })) { +} + +// Check if the key is already present in the keys vector +if (!std::any_of(keys.begin(), keys.end(), [&](const Key &key_value) { return key_value == key; })) { + // If the key is not present, add it to the keys vector keys.push_back(key); + + // Update the configuration storage with the updated keys vector config_storage->PutValue(kKeys, keys); - } +} // Persist kernel input shape + + // Create a vector to store the shapes of the kernel inputs std::vector>> shapes_list; + + // Check if the shapes list already exists in the configuration storage if (config_storage->Exists(kShapes)) { + // If it exists, retrieve the shapes list from the configuration storage shapes_list = config_storage->GetValue>>>(kShapes); } + + // Check if the shapes list has fewer elements than the number of keys if (shapes_list.size() < keys.size()) { + // Create a temporary vector to store the shapes of the kernel inputs std::vector> shape_tmp; + + // Copy the shapes of the kernel inputs into the temporary vector (void)std::transform(shapes->begin(), shapes->end(), std::back_inserter(shape_tmp), [](const std::shared_ptr> &shape_ptr) { return *shape_ptr; }); + + // Add the temporary vector of shapes to the shapes list shapes_list.push_back(shape_tmp); + + // Update the shapes list in the configuration storage config_storage->PutValue>>>(kShapes, shapes_list); } - // Persist parameter name of kernel. + // Create a vector to store the parameter names of the kernel std::vector param_names; + + // Check if the parameter names already exist in the configuration storage if (config_storage->Exists(kParamNames)) { + // If they exist, retrieve the parameter names from the configuration storage param_names = config_storage->GetValue>(kParamNames); } + + // Get the parameter name from the param_init_info object const std::string ¶m_name = param_init_info.param_name_; + + // Check if the number of parameter names is less than the number of keys if (param_names.size() < keys.size()) { + // If it is, add the current parameter name to the vector of parameter names param_names.push_back(param_name); + + // Update the configuration storage with the updated vector of parameter names config_storage->PutValue>(kParamNames, param_names); } } +// Definition of the function PersistInitParameters in the class ParameterServer + void ParameterServer::PersistInitParameters(const Key &key, const WeightPtr ¶m) { + + // Check if recovery is enabled if (!EnableRecovery()) { + // If recovery is not enabled, return from the function return; } - - MS_EXCEPTION_IF_NULL(server_node_); - std::string storage_file_path = std::string(kCurrentDirOfServer) + std::to_string(server_node_->rank_id()) + - std::string(kParamWithKey) + std::to_string(key); - if (!distributed::storage::FileIOUtils::IsFileOrDirExist(storage_file_path)) { - distributed::storage::FileIOUtils::CreateDir(storage_file_path); - } - - auto ret = FileUtils::GetRealPath(storage_file_path.c_str()); - if (!ret.has_value()) { - MS_LOG(EXCEPTION) << "Cannot get real path of persistent storage file for parameter, key: " << key; - } - - std::string real_storage_file_path = ret.value(); - auto persistent_weight = std::dynamic_pointer_cast(param); - MS_EXCEPTION_IF_NULL(persistent_weight); - std::map config_map; - config_map[distributed::storage::kFileStoragePath] = real_storage_file_path; - persistent_weight->Initialize(config_map); - - (void)weights_dirty_info_.emplace(key, distributed::storage::DirtyInfo()); - persistent_weight->Persist(distributed::storage::DirtyInfo()); - - MS_LOG(INFO) << "Finish persist initialized parameter, key: " << key; } +// Check if the server node is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(server_node_); + +// Create a string variable to store the storage file path +// The storage file path is constructed by concatenating the current directory of the server, the rank ID of the server node, the string "ParamWithKey", and the key +std::string storage_file_path = std::string(kCurrentDirOfServer) + std::to_string(server_node_->rank_id()) + + std::string(kParamWithKey) + std::to_string(key); + +// Check if the storage file path exists as a file or directory +if (!distributed::storage::FileIOUtils::IsFileOrDirExist(storage_file_path)) { + // If the storage file path does not exist, create a directory with the storage file path + distributed::storage::FileIOUtils::CreateDir(storage_file_path); +} + +// Use the FileUtils class to get the real path of the storage file specified by storage_file_path +auto ret = FileUtils::GetRealPath(storage_file_path.c_str()); + +// Check if the returned value has a value (i.e., if the real path was successfully obtained) +if (!ret.has_value()) { + // If the real path could not be obtained, log an exception message using MS_LOG + MS_LOG(EXCEPTION) << "Cannot get real path of persistent storage file for parameter, key: " << key; +} + +// Assign the value of ret.value() to the variable real_storage_file_path +std::string real_storage_file_path = ret.value(); + +// Cast the param pointer to a shared pointer of type PersistentWeight +auto persistent_weight = std::dynamic_pointer_cast(param); + +// Throw an exception if persistent_weight is null +MS_EXCEPTION_IF_NULL(persistent_weight); + +// Create an empty map of type std::map +std::map config_map; + +// Add an entry to the config_map with the key "distributed::storage::kFileStoragePath" and the value of real_storage_file_path +config_map[distributed::storage::kFileStoragePath] = real_storage_file_path; + +// Call the Initialize function of the persistent_weight object, passing in the config_map as an argument +persistent_weight->Initialize(config_map); + +// Emplace a new element into the weights_dirty_info_ map using the provided key and a default-constructed DirtyInfo object +(void)weights_dirty_info_.emplace(key, distributed::storage::DirtyInfo()); + +// Call the Persist function on the persistent_weight object, passing in a DirtyInfo object +persistent_weight->Persist(distributed::storage::DirtyInfo()); + +// Log an informational message using the MS_LOG macro +MS_LOG(INFO) << "Finish persist initialized parameter, key: " << key; + +// Definition of the InitEmbeddingTable function in the ParameterServer class + void ParameterServer::InitEmbeddingTable( const Key &key, const std::shared_ptr>>> &shapes, const ParamInitInfo ¶m_init_info) { + + // Check if recovery is enabled if (EnableRecovery()) { + + // If recovery is enabled, wait until the recovery process is finished while (!finish_recovery_) { std::this_thread::yield(); } } +} - std::unique_lock locker(access_weight_mutex_); +// Create a unique_lock object named "locker" and lock the mutex named "access_weight_mutex_" +std::unique_lock locker(access_weight_mutex_); - MS_EXCEPTION_IF_NULL(shapes); - if (weights_.count(key) == 0) { - std::shared_ptr lookup = - std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); - lookup->InitKernel(shapes); - embedding_lookup_ops_[key] = lookup; +// Check if the pointer `shapes` is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(shapes); - PersistKernels(key, shapes, param_init_info); +// Check if the key is not present in the `weights_` map +if (weights_.count(key) == 0) { - // Init embedding weight - const std::vector &input_shapes = lookup->input_sizes(); - size_t total_dims = - std::accumulate(input_shapes.begin(), input_shapes.end(), IntToSize(1), std::multiplies()); + // Create a shared pointer to a `PServerKernel` object using the `EmbeddingLookUpPSKernelMod` constructor + std::shared_ptr lookup = std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); - std::shared_ptr> embedding_shape = std::make_shared>(); - (void)std::transform(input_shapes.begin(), input_shapes.end(), std::back_inserter(*embedding_shape), - [](size_t dim) { return static_cast(dim); }); + // Initialize the kernel with the given `shapes` + lookup->InitKernel(shapes); - WeightPtr embedding = - Util::MakeWeightPtr(std::make_shared>(total_dims, 0), EnableRecovery(), embedding_shape); - MS_EXCEPTION_IF_NULL(embedding); - float *embedding_data = embedding->data(); + // Add the `lookup` object to the `embedding_lookup_ops_` map with the `key` as the key + embedding_lookup_ops_[key] = lookup; +// Call the function PersistKernels with the arguments key, shapes, and param_init_info + +// Initialize the embedding weight + +// Get the input shapes from the lookup object and store them in a reference to a constant vector +const std::vector &input_shapes = lookup->input_sizes(); + +// Calculate the total number of dimensions by summing up the elements in the input_shapes vector +// Start with an initial value of 1 and use std::accumulate to perform the summation +size_t total_dims = std::accumulate(input_shapes.begin(), input_shapes.end(), IntToSize(1), std::multiplies()); + +// Create a shared pointer to a vector of integers named "embedding_shape" using std::make_shared +std::shared_ptr> embedding_shape = std::make_shared>(); + +// Use std::transform to iterate over the elements in the "input_shapes" vector +// and transform each element to an integer using a lambda function +// The transformed elements are then inserted at the end of the "embedding_shape" vector using std::back_inserter +(void)std::transform(input_shapes.begin(), input_shapes.end(), std::back_inserter(*embedding_shape), + [](size_t dim) { return static_cast(dim); }); + +// Declare a pointer variable named `embedding` of type `WeightPtr` +WeightPtr embedding = + + // Call the `MakeWeightPtr` function from the `Util` namespace, passing in the following arguments: + // - A shared pointer to a dynamically allocated vector of floats, initialized with `total_dims` number of elements, all set to 0 + // - The result of calling the `EnableRecovery` function (which returns an object of type `EnableRecovery`) + // - The value of `embedding_shape` + Util::MakeWeightPtr(std::make_shared>(total_dims, 0), EnableRecovery(), embedding_shape); + +// Check if `embedding` is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(embedding); + +// Declare a pointer variable named `embedding_data` of type `float`, and assign it the address of the data stored in `embedding` +float *embedding_data = embedding->data(); + + // Check if the cache is enabled for data prefetching if (ps::PsDataPrefetch::GetInstance().cache_enable()) { - CacheEmbeddingTableParamPtr(); - if (param_init_info.param_type_ == kWeight) { - const std::string ¶m_name = param_init_info.param_name_; - auto iter = embedding_parameter_tables_.find(param_name); - if (iter == embedding_parameter_tables_.end()) { - MS_LOG(EXCEPTION) << "Can not find parameter info for: " << param_name; - } - // Cache embedding table parameter by weight key to parameter node pointer. - (void)embedding_tables_.emplace(key, iter->second); - InitRandomNormal(0, kStdDev, input_shapes, param_init_info.global_seed_, param_init_info.op_seed_, - embedding_data); - } else if (param_init_info.param_type_ == kAccumulation) { - InitAccumParallel(param_init_info.init_val_, total_dims, embedding_data); - } - } else { - std::default_random_engine engine; - std::normal_distribution random(0, kStdDev); - for (size_t i = 0; i < total_dims; i++) { - embedding_data[i] = random(engine); - } + // Call the CacheEmbeddingTableParamPtr function + + // Check if the parameter type is kWeight + if (param_init_info.param_type_ == kWeight) { + + // Get the parameter name from param_init_info + const std::string ¶m_name = param_init_info.param_name_; + + // Find the parameter info in the embedding_parameter_tables_ map using the parameter name + auto iter = embedding_parameter_tables_.find(param_name); + + // If the parameter info is not found, throw an exception + if (iter == embedding_parameter_tables_.end()) { + MS_LOG(EXCEPTION) << "Can not find parameter info for: " << param_name; + } + + // Cache the embedding table parameter by weight key to parameter node pointer in the embedding_tables_ map + (void)embedding_tables_.emplace(key, iter->second); + } } - PersistInitParameters(key, embedding); + // Check if the parameter initialization type is Normal + if (param_init_info.param_type_ == kNormal) { + // Initialize the embedding data using a normal distribution with mean 0 and standard deviation kStdDev + InitRandomNormal(0, kStdDev, input_shapes, param_init_info.global_seed_, param_init_info.op_seed_, embedding_data); + } + // Check if the parameter initialization type is Accumulation + else if (param_init_info.param_type_ == kAccumulation) { + // Initialize the embedding data with a constant value specified by param_init_info.init_val_ + InitAccumParallel(param_init_info.init_val_, total_dims, embedding_data); + } + // If the parameter initialization type is not specified, use a default random engine and normal distribution + else { + std::default_random_engine engine; + std::normal_distribution random(0, kStdDev); + // Loop through each element in the embedding data and assign a random value from the normal distribution + for (size_t i = 0; i < total_dims; i++) { + embedding_data[i] = random(engine); + } + } +// Call the function PersistInitParameters with the arguments 'key' and 'embedding' +PersistInitParameters(key, embedding); + + // Assign the embedding to the weights_ map with the given key weights_[key] = embedding; + + // Output a debug log message with the key and the value of the embedding MS_LOG(DEBUG) << "The key:" << key << " the embedding:" << *(embedding->MutableData()); + + // Assign 0 to the tokens_ map with the given key tokens_[key] = 0; + + // Set the value of the is_embedding_ map with the given key to true is_embedding_[key] = true; + // Set the value of grads_accum_counter_ at the given key to 0 grads_accum_counter_[key] = 0; } } -bool ParameterServer::HasWeight(const Key &key) { return (weights_.count(key) > 0 && !is_embedding_.count(key)); } +// Check if the ParameterServer has a weight with the given key +bool ParameterServer::HasWeight(const Key &key) { -void ParameterServer::Finalize() { - running_ = false; - apply_grads_cv_.notify_one(); + // Check if the weight exists in the weights_ map and if it is not an embedding + return (weights_.count(key) > 0 && !is_embedding_.count(key)); +} +// The Finalize function of the ParameterServer class + +// Set the running_ flag to false, indicating that the server is no longer running +running_ = false; + +// Notify one waiting thread that the grads have been applied +apply_grads_cv_.notify_one(); + + // Check if the persist_thread_ pointer is not null and if the thread is joinable if (persist_thread_ != nullptr && persist_thread_->joinable()) { + + // If the above condition is true, join the persist_thread_ persist_thread_->join(); } } +// The UpdateWeights function of the ParameterServer class + void ParameterServer::UpdateWeights() { + + // Run an infinite loop while (true) { + + // Log the values of the running_ and ReadyForUpdateWeights() variables using MS_LOG(INFO) MS_LOG(INFO) << "The running is:" << running_ << " the ready is:" << this->ReadyForUpdateWeights(); + + // Create a unique lock for the mutex_ std::unique_lock lock(mutex_); + + // Wait for the apply_grads_cv_ condition variable to be notified, and check if ReadyForUpdateWeights() is true or running_ is false apply_grads_cv_.wait(lock, [this] { return this->ReadyForUpdateWeights() || !running_; }); + + // If running_ is false, break out of the loop if (!running_) { break; } + } +} + // Iterate over the elements in the weights_ map using an iterator for (auto iter = weights_.begin(); iter != weights_.end(); iter++) { + + // Get the key of the current element Key key = iter->first; + + // Get the weight pointer of the current element WeightPtr weight_ptr = iter->second; - std::shared_ptr optimizer = nullptr; - if (weight_key_to_optims_.count(key) > 0) { - optimizer = optimizers_[key]; - } - MS_EXCEPTION_IF_NULL(optimizer); +// Declare a shared pointer named "optimizer" of type PServerKernel and initialize it to nullptr +std::shared_ptr optimizer = nullptr; - std::shared_ptr optim_info = optim_infos_[key]; - if (optim_info != nullptr) { - const std::vector &inputs = optim_info->inputs(); - const std::vector &workspaces = optim_info->workspaces(); - const std::vector &outputs = optim_info->outputs(); +// Check if the key exists in the weight_key_to_optims_ map +if (weight_key_to_optims_.count(key) > 0) { + // If the key exists, assign the corresponding optimizer from the optimizers_ map to the optimizer variable + optimizer = optimizers_[key]; +} - std::vector> shapes = {}; - std::vector indices_shape = {}; - indices_shape.emplace_back(optim_info->indice_size()); - shapes.push_back(indices_shape); +// Check if the optimizer is null +MS_EXCEPTION_IF_NULL(optimizer); +// If the optimizer is null, throw an exception +// Create a shared pointer to an instance of the OptimizerInfo class and assign it the value from the optim_infos_ map using the key 'key' +std::shared_ptr optim_info = optim_infos_[key]; + +// Check if the optim_info pointer is not null +if (optim_info != nullptr) { + + // Get the inputs, workspaces, and outputs vectors from the optim_info object + const std::vector &inputs = optim_info->inputs(); + const std::vector &workspaces = optim_info->workspaces(); + const std::vector &outputs = optim_info->outputs(); + +// Create an empty vector of vectors of size_t named "shapes" +std::vector> shapes = {}; + +// Create an empty vector of size_t named "indices_shape" +std::vector indices_shape = {}; + +// Add the size of "optim_info->indice_size()" to the back of "indices_shape" +indices_shape.emplace_back(optim_info->indice_size()); + +// Add "indices_shape" to the back of "shapes" +shapes.push_back(indices_shape); + + // Check if the key exists in the original_optim_inputs_shape_ map if (original_optim_inputs_shape_.count(key) != 0) { + // If the key exists, transform the vector of shared pointers to vectors into a vector of vectors std::transform((*(original_optim_inputs_shape_[key])).begin(), (*(original_optim_inputs_shape_[key])).end(), std::back_inserter(shapes), [](const std::shared_ptr> &input_shapes) -> std::vector { return *input_shapes; }); } + + // Reinitialize the optimizer with the updated shapes optimizer->ReInit(shapes); + + // Compute the mean of the shapes using the optim_info object optim_info->ComputeMean(shapes, worker_num_, pserver_num_, server_node_->rank_id()); + + // Execute the optimizer with the inputs, workspaces, and outputs optimizer->Execute(inputs, workspaces, outputs); + + // Reset the optim_info object optim_info->Reset(); } + + // Check if the key is not an embedding if (!is_embedding_[key]) { + // Set the number of tokens for the key to the worker_num_ tokens_[key] = worker_num_; } } + + // Reset the grad_accum_count ResetGradAccumCount(); } } +// Define the function `AccumGrad` belonging to the class `ParameterServer` void ParameterServer::AccumGrad(const Keys &keys, const Values &values, const Lengths &lengths) { + + // Acquire a unique lock on the mutex to ensure exclusive access to the critical section std::unique_lock lock(mutex_); + + // Extract the first key from the `keys` vector const Key &key = keys[0]; + + // Check if there is no sparse gradient by verifying if the size of `values` is 1 and the value is `kGradValue` bool no_sparse_grad = values.size() == 1 && values[0] == kGradValue; + + // If there is no sparse gradient if (!no_sparse_grad) { + + // Retrieve the optimizer information associated with the key from the `optim_infos_` map std::shared_ptr optim_info = optim_infos_[key]; // Create or update the optimizer info + + // Check if the optimizer info for the given key is already present if (optim_info == nullptr) { + + // Get the optimizer info builder for the optimizer associated with the given key const std::shared_ptr &builder = optim_info_builders_[weight_key_to_optims_[key]]; + + // Get the PServerKernel associated with the optimizer std::shared_ptr pserver_kernel = optimizers_[key]; + + // If no optimizer is found for the given key, throw an exception if (pserver_kernel == nullptr) { MS_LOG(EXCEPTION) << "no optimizer found for key " << key << " optim name " << weight_key_to_optims_[key]; } - MS_EXCEPTION_IF_NULL(pserver_kernel); + + // Build the optimizer info using the optimizer info builder, PServerKernel, weights, keys, values, lengths, + // optimizer inputs shape, worker number, and whether the optimizer is for embedding OptimizerInfo *optim = builder->Build(pserver_kernel, weights_[key], keys, values, lengths, optim_inputs_shape_[key], worker_num_, is_embedding_[key]); + + // Reset the optim_info shared pointer to point to the newly created optimizer info optim_info.reset(optim); + + // Store the optim_info in the optim_infos_ map for future reference optim_infos_[key] = optim_info; } else { + + // If the optimizer info is already present, update the values and lengths optim_info->Update(values, lengths); + + // Accumulate the values and lengths in the optimizer info optim_info->Accumulate(values, lengths); } } + // Increment the counter for the given key in the grads_accum_counter_ map grads_accum_counter_[key] += 1; + + // If the counter for the given key reaches the worker_num_, increment the grad_accum_count_ if (grads_accum_counter_[key] == worker_num_) { grad_accum_count_++; } + + // If all workers have accumulated their gradients, notify the apply_grads_cv_ condition variable if (ReadyForUpdateWeights()) { apply_grads_cv_.notify_one(); } } +// Retrieve the weight associated with the given key from the ParameterServer +// The weight is accessed in a thread-safe manner using a unique lock on the mutex WeightPtr ParameterServer::weight(const Key &key) { + + // Acquire a unique lock on the mutex to ensure exclusive access to the weights std::unique_lock lock(mutex_); + + // Check if the weight with the given key exists in the weights map if (weights_.count(key) == 0) { + + // If the weight does not exist, throw an exception with an error message MS_LOG(EXCEPTION) << "Invalid weight key " << key; } + + // Retrieve the weight pointer associated with the key from the weights map WeightPtr weight_ptr = weights_[key]; + + // Check if the retrieved weight pointer is null MS_EXCEPTION_IF_NULL(weight_ptr); + + // Decrement the token count associated with the key by 1 tokens_[key] -= 1; + + // Return the retrieved weight pointer return weight_ptr; } +// Definition of the function "DoEmbeddingLookup" in the class "ParameterServer" + void ParameterServer::DoEmbeddingLookup(Key key, const LookupIds &lookup_ids, KVMessage *res) { + + // Check if recovery is enabled if (EnableRecovery()) { + + // Loop until the recovery process is finished while (!finish_recovery_) { + + // Yield the current thread to allow other threads to execute std::this_thread::yield(); } } +} - std::unique_lock lock(mutex_); - MS_EXCEPTION_IF_NULL(res); - if (weights_.count(key) == 0) { +// Acquire a unique lock on the mutex object 'mutex_' +std::unique_lock lock(mutex_); + +// Check if the pointer 'res' is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(res); + +// Check if the key is present in the 'weights_' map, and log an error message if it is not +if (weights_.count(key) == 0) { MS_LOG(ERROR) << "Invalid embedding table key " << key; return; - } - if (embedding_lookup_ops_.count(key) == 0) { +} + +// Check if the key is present in the 'embedding_lookup_ops_' map, and log an error message if it is not +if (embedding_lookup_ops_.count(key) == 0) { MS_LOG(ERROR) << "Invalid embedding lookup op key " << key; return; - } - WeightPtr table_ptr = weights_[key]; - MS_EXCEPTION_IF_NULL(table_ptr); - std::shared_ptr table_lookup_op = embedding_lookup_ops_[key]; - MS_EXCEPTION_IF_NULL(table_lookup_op); - - // Update shapes of lookup operator - std::vector> shapes = {}; - std::vector indices_shape = {}; - indices_shape.emplace_back(lookup_ids.size()); - shapes.push_back(indices_shape); - table_lookup_op->ReInit(shapes); - - const std::vector output_shapes = table_lookup_op->output_sizes(); - std::vector inputs; - AddressPtr embedding_table = std::make_shared(); - MS_EXCEPTION_IF_NULL(embedding_table); - AddressPtr indices = std::make_shared(); - MS_EXCEPTION_IF_NULL(indices); - inputs.push_back(embedding_table); - inputs.push_back(indices); - embedding_table->addr = table_ptr->data(); - embedding_table->size = table_ptr->size() * sizeof(float); - - std::unique_ptr tmp_ids = std::make_unique(lookup_ids.size()); - MS_EXCEPTION_IF_NULL(tmp_ids); - for (size_t i = 0; i < lookup_ids.size(); i++) { - tmp_ids[i] = static_cast(lookup_ids[i]); - } - indices->addr = tmp_ids.get(); - indices->size = lookup_ids.size() * sizeof(int); - - std::vector workspaces; - std::vector outputs; - AddressPtr output = std::make_shared(); - MS_EXCEPTION_IF_NULL(output); - std::shared_ptr addr = std::make_shared(output_shapes[0] / sizeof(float), 0); - MS_EXCEPTION_IF_NULL(addr); - - output->addr = addr->data(); - output->size = output_shapes[0]; - outputs.push_back(output); - - table_lookup_op->Execute(inputs, workspaces, outputs); - *res->mutable_values() = {addr->begin(), addr->end()}; - res->add_len(res->values_size()); } +// Get the value associated with the key from the 'weights_' map and assign it to the pointer 'table_ptr' +WeightPtr table_ptr = weights_[key]; + +// Check if the pointer 'table_ptr' is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(table_ptr); + +// Get the value associated with the key from the 'embedding_lookup_ops_' map and assign it to the shared pointer 'table_lookup_op' +std::shared_ptr table_lookup_op = embedding_lookup_ops_[key]; + +// Check if the shared pointer 'table_lookup_op' is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(table_lookup_op); + +// Update shapes of lookup operator + +// Create an empty vector of vectors to store the shapes +std::vector> shapes = {}; + +// Create a vector to store the shape of the indices +std::vector indices_shape = {}; + +// Add the size of lookup_ids to the indices_shape vector +indices_shape.emplace_back(lookup_ids.size()); + +// Push the indices_shape vector into the shapes vector +shapes.push_back(indices_shape); + +// Reinitialize the table_lookup_op with the updated shapes +table_lookup_op->ReInit(shapes); + +// Declare a constant vector of size_t called output_shapes and assign it the value returned by the output_sizes() function of the table_lookup_op object +const std::vector output_shapes = table_lookup_op->output_sizes(); + +// Declare a vector of kernel::AddressPtr called inputs +std::vector inputs; + +// Declare a shared pointer of kernel::Address called embedding_table and initialize it using std::make_shared() +AddressPtr embedding_table = std::make_shared(); + +// Check if embedding_table is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(embedding_table); + +// Declare a shared pointer of kernel::Address called indices and initialize it using std::make_shared() +AddressPtr indices = std::make_shared(); + +// Check if indices is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(indices); + +// Add embedding_table and indices to the inputs vector +inputs.push_back(embedding_table); +inputs.push_back(indices); + +// Set the address of embedding_table to the data pointer of the table_ptr object +embedding_table->addr = table_ptr->data(); + +// Set the size of embedding_table to the size of the table_ptr object multiplied by the size of a float +embedding_table->size = table_ptr->size() * sizeof(float); + +// Create a unique_ptr object named tmp_ids and allocate memory for an array of integers with the size of lookup_ids +std::unique_ptr tmp_ids = std::make_unique(lookup_ids.size()); + +// Check if tmp_ids is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(tmp_ids); + +// Iterate through the lookup_ids vector and assign each element to the corresponding index in tmp_ids +for (size_t i = 0; i < lookup_ids.size(); i++) { + tmp_ids[i] = static_cast(lookup_ids[i]); +} + +// Set the address of the tmp_ids array as the address of the indices pointer +indices->addr = tmp_ids.get(); + +// Set the size of the indices pointer to the size of lookup_ids multiplied by the size of an integer +indices->size = lookup_ids.size() * sizeof(int); + +// Create an empty vector of kernel::AddressPtr objects named "workspaces" +std::vector workspaces; + +// Create an empty vector of kernel::AddressPtr objects named "outputs" +std::vector outputs; + +// Create a shared pointer to a kernel::Address object named "output" using std::make_shared +AddressPtr output = std::make_shared(); + +// Check if the output pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(output); + +// Create a shared pointer to a Values object named "addr" using std::make_shared +std::shared_ptr addr = std::make_shared(output_shapes[0] / sizeof(float), 0); + +// Check if the addr pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(addr); + +// Set the `addr` member of the `output` object to the address of the `data` member of the `addr` object +output->addr = addr->data(); + +// Set the `size` member of the `output` object to the value of the first element in the `output_shapes` vector +output->size = output_shapes[0]; + +// Add the `output` object to the `outputs` vector +outputs.push_back(output); + +// Call the Execute function of the table_lookup_op object, passing in the inputs, workspaces, and outputs as arguments +table_lookup_op->Execute(inputs, workspaces, outputs); + +// Assign the values of the addr vector to the values field of the res object using the mutable_values() function +*res->mutable_values() = {addr->begin(), addr->end()}; + +// Add the length of the values field to the len field of the res object using the add_len() function +res->add_len(res->values_size()); + +// Definition of the function UpdateEmbeddings in the class ParameterServer + void ParameterServer::UpdateEmbeddings(const Key &key, const LookupIds &lookup_ids, const Values &vals) { + + // Check if recovery is enabled if (EnableRecovery()) { + + // If recovery is enabled, wait until the recovery process is finished while (!finish_recovery_) { std::this_thread::yield(); } } +} - std::unique_lock locker(access_weight_mutex_); +// Create a unique_lock object named "locker" and lock the mutex named "access_weight_mutex_" +std::unique_lock locker(access_weight_mutex_); + // Check if the key exists in the weights_ map if (weights_.count(key) == 0) { + // If the key does not exist, log an error message and return MS_LOG(ERROR) << "Invalid embedding table key " << key; return; } + + // Check if the key exists in the embedding_lookup_ops_ map if (embedding_lookup_ops_.count(key) == 0) { + // If the key does not exist, log an error message and return MS_LOG(ERROR) << "Invalid embedding lookup op key " << key; return; } + + // Get the weight pointer associated with the key from the weights_ map WeightPtr table_ptr = weights_[key]; + + // Check if the weight pointer is null MS_EXCEPTION_IF_NULL(table_ptr); + + // Get the embedding lookup op associated with the key from the embedding_lookup_ops_ map std::shared_ptr lookup_op = embedding_lookup_ops_[key]; + + // Check if the embedding lookup op is null MS_EXCEPTION_IF_NULL(lookup_op); + + // Call the UpdateEmbeddings function of the lookup_op with the appropriate arguments lookup_op->UpdateEmbeddings(table_ptr->data(), lookup_ids.data(), vals.data(), lookup_ids.size()); - UpdateDirtyInfo(key, lookup_ids, lookup_op->offset()); -} +// Call the function UpdateDirtyInfo with the provided arguments: +// - key: the key to update the dirty information for +// - lookup_ids: the lookup IDs associated with the key +// - lookup_op->offset(): the offset of the lookup operation +UpdateDirtyInfo(key, lookup_ids, lookup_op->offset()); + +// Update the dirty information for a given key, lookup_ids, and offset in the ParameterServer class void ParameterServer::UpdateDirtyInfo(const Key &key, const LookupIds &lookup_ids, int64_t offset) { + + // Check if recovery is enabled if (EnableRecovery()) { + + // Create a set to store sorted ids std::set sorted_ids; + + // Iterate over each lookup id in the lookup_ids vector (void)std::for_each(lookup_ids.begin(), lookup_ids.end(), [&](uint64_t id) { + + // Calculate the index by subtracting the offset from the lookup id int index = SizeToInt(id) - LongToInt(offset); + + // Insert the index into the sorted_ids set (void)sorted_ids.insert(index); }); + // Find the iterator for the given key in the weights_dirty_info_ map auto iter = weights_dirty_info_.find(key); + + // If the iterator points to the end of the map, it means the key was not found if (iter == weights_dirty_info_.end()) { - MS_LOG(EXCEPTION) << "Cannot find dirty info for embedding table, key: " << key; + // Log an exception with a descriptive error message + MS_LOG(EXCEPTION) << "Cannot find dirty info for embedding table, key: " << key; } + + // Get a reference to the dirty_info object associated with the key distributed::storage::DirtyInfo &dirty_info = iter->second; + + // Use std::for_each algorithm to iterate over the sorted_ids vector + // and push each element to the dirty_info vector (void)std::for_each(sorted_ids.begin(), sorted_ids.end(), [&](int id) { dirty_info.push_back(id); }); } } +// Check if the size of the grads_accum_counter_ vector is greater than 0 +// and if the grad_accum_count_ is equal to the size of the grads_accum_counter_ vector inline bool ParameterServer::ReadyForUpdateWeights() const { return grads_accum_counter_.size() > 0 && grad_accum_count_ == grads_accum_counter_.size(); } +// Check if the ParameterServer is ready to receive a push for a given key inline bool ParameterServer::ReadyForPush(const Key &key) { + + // Acquire a lock on the mutex to ensure thread safety std::unique_lock lock(mutex_); + + // Check if the weights in the server are empty if (weights_.empty()) { + + // If the weights are empty, throw an exception with a descriptive error message MS_LOG(EXCEPTION) << "The weights in server is empty. Many reasons could cause this: 1.The Worker didn't send " "kInitWeightsCmd command. 2.The Server failed to initialize weights."; } + + // Check if the number of gradient accumulations is less than the number of weights, + // and if the token count for the given key is 0 return grad_accum_count_ < weights_.size() && tokens_[key] == 0; } +// Check if the ParameterServer is ready for a pull operation for a given key inline bool ParameterServer::ReadyForPull(const Key &key) { + + // Acquire a unique lock on the mutex to ensure thread safety std::unique_lock lock(mutex_); + + // Check if the key is present in the tokens_ map or if the weight for the key is 0 if (tokens_.count(key) == 0 || weights_[key] == 0) { + + // If the key is not present or the weight is 0, throw an exception with an error message MS_LOG(EXCEPTION) << "Invalid weight key " << key; } + + // Log an informational message indicating if the token count for the key is greater than 0 MS_LOG(INFO) << "ReadyForPull: " << (tokens_[key] > 0); + + // Return true if the token count for the key is greater than 0, otherwise return false return tokens_[key] > 0; } +// Reset the grad_accum_count_ variable to 0 inline void ParameterServer::ResetGradAccumCount() { grad_accum_count_ = 0; + + // Iterate over each element in the grads_accum_counter_ map for (auto iter = grads_accum_counter_.begin(); iter != grads_accum_counter_.end(); iter++) { + // Set the value of the current element to 0 grads_accum_counter_[iter->first] = 0; } } +// Define a constant member function named GetCNode in the class ParameterServer, which returns a constant CNodePtr object const CNodePtr ParameterServer::GetCNode(const std::string &name) const { + + // Get a list of CNodePtr objects by calling the GetOrderedCnodes function of the func_graph_ member variable std::list cnodes = func_graph_->GetOrderedCnodes(); + + // Iterate through each CNodePtr object in the cnodes list for (CNodePtr cnode : cnodes) { + + // Check if the cnode pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(cnode); + + // Get the full name of the cnode object, including its scope, and store it in the fullname variable std::string fullname = cnode->fullname_with_scope(); + + // Check if the fullname string contains the given name and also contains the substring "Push" if (fullname.find(name) != std::string::npos && fullname.find("Push") != std::string::npos) { + + // If the conditions are met, return the cnode object return cnode; } } + + // If no matching cnode is found, return a null pointer return nullptr; } -inline std::mutex &ParameterServer::mutex() { return mutex_; } +// Define an inline function named "mutex" that returns a reference to the mutex member variable of the ParameterServer class +inline std::mutex &ParameterServer::mutex() { + return mutex_; +} + +// Definition of the function GetEmbeddingTableParamPtr in the ParameterServer class void ParameterServer::GetEmbeddingTableParamPtr() { - if (ps::PsDataPrefetch::GetInstance().cache_enable()) { - return; - } - MS_EXCEPTION_IF_NULL(func_graph_); - auto cnodes = func_graph_->GetOrderedCnodes(); - Key count = 0; - for (auto cnode : cnodes) { - MS_EXCEPTION_IF_NULL(cnode); - std::string cnode_name = Util::GetPrimitiveName(cnode); - if (cnode_name == kEmbeddingLookupOpName || cnode_name == kGatherV2OpName || cnode_name == kSparseGatherV2OpName) { - auto embedding_table = common::AnfAlgo::GetInputNode(cnode, 0); - if (IsPrimitiveCNode(embedding_table, prim::kPrimLoad)) { - auto embedding_cnode = embedding_table->cast(); - embedding_table = common::AnfAlgo::GetInputNode(embedding_cnode, 0); - } - MS_EXCEPTION_IF_NULL(embedding_table); - if (embedding_table->isa()) { - MS_LOG(INFO) << "Embedding table name is " << embedding_table->fullname_with_scope() << ", key is " << count; - (void)embedding_tables_.emplace(count, embedding_table->cast()); - count++; - } - } + // Check if the cache is enabled in the PsDataPrefetch singleton instance + if (ps::PsDataPrefetch::GetInstance().cache_enable()) { + + // If the cache is enabled, return from the function + return; } } -void ParameterServer::CacheEmbeddingTableParamPtr() { - if (embedding_param_ptr_cached_) { - return; - } +// Check if the function graph is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(func_graph_); - MS_EXCEPTION_IF_NULL(func_graph_); - auto cnodes = func_graph_->GetOrderedCnodes(); - for (auto cnode : cnodes) { +// Get the ordered list of CNodes in the function graph +auto cnodes = func_graph_->GetOrderedCnodes(); + +// Initialize a variable to keep track of the count of embedding tables +Key count = 0; + +// Iterate over each CNode in the list +for (auto cnode : cnodes) { + // Check if the CNode is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(cnode); + + // Get the name of the primitive operation associated with the CNode std::string cnode_name = Util::GetPrimitiveName(cnode); + + // Check if the primitive operation is one of the supported embedding lookup operations + if (cnode_name == kEmbeddingLookupOpName || cnode_name == kGatherV2OpName || cnode_name == kSparseGatherV2OpName) { + // Get the input node of the CNode at index 0 + auto embedding_table = common::AnfAlgo::GetInputNode(cnode, 0); + + // Check if the input node is a primitive CNode of type prim::kPrimLoad + if (IsPrimitiveCNode(embedding_table, prim::kPrimLoad)) { + // If it is, cast the input node to a CNode pointer and get its input node at index 0 + auto embedding_cnode = embedding_table->cast(); + embedding_table = common::AnfAlgo::GetInputNode(embedding_cnode, 0); + } + + // Check if the embedding table is a Parameter node + if (embedding_table->isa()) { + // Print the name and key of the embedding table + MS_LOG(INFO) << "Embedding table name is " << embedding_table->fullname_with_scope() << ", key is " << count; + + // Add the embedding table to the map with its key as the count + (void)embedding_tables_.emplace(count, embedding_table->cast()); + + // Increment the count + count++; + } + } +} +// Closing brace to end the main function +} + +// Function to cache the embedding table parameter pointer + +void ParameterServer::CacheEmbeddingTableParamPtr() { + + // Check if the embedding parameter pointer is already cached + if (embedding_param_ptr_cached_) { + return; // If it is already cached, return without doing anything + } + + // If the embedding parameter pointer is not cached, continue with the caching process + +} + + // Check if the func_graph_ is null, if it is, throw an exception + MS_EXCEPTION_IF_NULL(func_graph_); + + // Get the ordered list of CNodes from the func_graph_ + auto cnodes = func_graph_->GetOrderedCnodes(); + + // Iterate over each CNode in the cnodes list + for (auto cnode : cnodes) { + // Check if the cnode is null, if it is, throw an exception + MS_EXCEPTION_IF_NULL(cnode); + + // Get the name of the primitive operation associated with the cnode + std::string cnode_name = Util::GetPrimitiveName(cnode); + + // Check if the cnode_name is not equal to "GatherV2" and "SparseGatherV2", if it is not, skip to the next iteration if (cnode_name != kGatherV2OpName && cnode_name != kSparseGatherV2OpName) { continue; } + // Get the input node of the given CNode and assign it to the variable embedding_table auto embedding_table = common::AnfAlgo::GetInputNode(cnode, 0); + + // Check if the embedding_table is a primitive CNode with the primitive type prim::kPrimLoad if (IsPrimitiveCNode(embedding_table, prim::kPrimLoad)) { + + // If embedding_table is a primitive CNode with prim::kPrimLoad, cast it to CNodePtr and assign it to embedding_cnode auto embedding_cnode = embedding_table->cast(); + + // Get the input node of embedding_cnode and assign it back to embedding_table embedding_table = common::AnfAlgo::GetInputNode(embedding_cnode, 0); } + // Check if the embedding_table is not null MS_EXCEPTION_IF_NULL(embedding_table); + + // Check if the embedding_table is of type Parameter if (embedding_table->isa()) { - (void)embedding_parameter_tables_.emplace(embedding_table->fullname_with_scope(), - embedding_table->cast()); + // Add the embedding_table to the embedding_parameter_tables_ map + // using its fullname_with_scope() as the key and casting it to ParameterPtr + (void)embedding_parameter_tables_.emplace(embedding_table->fullname_with_scope(), + embedding_table->cast()); } } - embedding_param_ptr_cached_ = true; -} +// Set the value of the variable "embedding_param_ptr_cached_" to true +// Definition of the function "RecoverKernels" belonging to the class "ParameterServer" void ParameterServer::RecoverKernels(const std::vector &keys, const std::vector>> &shapes_list, const std::vector ¶m_names) { + // Iterate over the keys vector using a for loop for (size_t i = 0; i < keys.size(); i++) { + // Get the current key from the keys vector size_t key = keys.at(i); + + // Check if the key is not present in the weights_ map if (weights_.count(key) == 0) { - // Recover embedding lookup kernels. + // If the key is not present, recover the embedding lookup kernels + + // Create a shared pointer to a vector of shared pointers to vectors of size_t std::shared_ptr>>> shapes_ptr = std::make_shared>>>(); - const auto &shapes = shapes_list[i]; - for (const auto &shape : shapes) { - std::shared_ptr> shape_ptr = - std::make_shared>(shape.begin(), shape.end()); - shapes_ptr->push_back(shape_ptr); - } - std::shared_ptr lookup = - std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); - lookup->InitKernel(shapes_ptr); - embedding_lookup_ops_[key] = lookup; +// Create a constant reference to the element at index 'i' in the 'shapes_list' vector and name it 'shapes' +const auto &shapes = shapes_list[i]; - // Recover embedding table parameter node address in graph. - const auto ¶m_name = param_names.at(i); - auto iter = embedding_parameter_tables_.find(param_name); - if (iter != embedding_parameter_tables_.end()) { - // Cache embedding table parameter by weight key to parameter node pointer. - (void)embedding_tables_.emplace(key, iter->second); - } - } - } +// Iterate over each element in the 'shapes' vector using a range-based for loop +for (const auto &shape : shapes) { + + // Create a shared pointer to a vector of size_t and name it 'shape_ptr' + // Initialize 'shape_ptr' with a new vector constructed from the 'begin' and 'end' iterators of the 'shape' object + std::shared_ptr> shape_ptr = + std::make_shared>(shape.begin(), shape.end()); + + // Add 'shape_ptr' to the 'shapes_ptr' vector + shapes_ptr->push_back(shape_ptr); } +// Create a shared pointer to a PServerKernel object named "lookup" using the make_shared function +// The PServerKernel object is created using the EmbeddingLookUpPSKernelMod constructor +// The constructor takes three arguments: server_node_->rank_id(), pserver_num_, and worker_num_ +std::shared_ptr lookup = std::make_shared(server_node_->rank_id(), pserver_num_, worker_num_); + +// Initialize the lookup kernel by calling the InitKernel function on the lookup object +// The InitKernel function takes a pointer to a shapes object as an argument +lookup->InitKernel(shapes_ptr); + +// Add the lookup kernel to the embedding_lookup_ops_ map using the key "key" +embedding_lookup_ops_[key] = lookup; + +// Recover embedding table parameter node address in graph. +const auto ¶m_name = param_names.at(i); + +// Find the embedding parameter table in the map using the parameter name +auto iter = embedding_parameter_tables_.find(param_name); + +// If the embedding parameter table is found in the map +if (iter != embedding_parameter_tables_.end()) { + // Cache the embedding table parameter by weight key to parameter node pointer + (void)embedding_tables_.emplace(key, iter->second); +} +} +} +} + +// This function is used to recover parameters based on a given vector of keys void ParameterServer::RecoverParameters(const std::vector &keys) { + + // Iterate over each key in the vector for (size_t i = 0; i < keys.size(); i++) { + + // Get the current key size_t key = keys.at(i); + + // Check if the key is not present in the weights map if (weights_.count(key) == 0) { + + // Find the corresponding embedding lookup operation for the key auto iter = embedding_lookup_ops_.find(key); + + // If the embedding lookup operation is not found, throw an exception if (iter == embedding_lookup_ops_.end()) { MS_LOG(EXCEPTION) << "Cannot find embedding lookup kernel for key: " << key; } + + // Get the shared pointer to the embedding lookup operation std::shared_ptr lookup = iter->second; + + // Throw an exception if the lookup pointer is null MS_EXCEPTION_IF_NULL(lookup); + + // Get the input shapes from the lookup operation const std::vector &input_shapes = lookup->input_sizes(); + + // Calculate the total dimensions by accumulating the input shapes size_t total_dims = std::accumulate(input_shapes.begin(), input_shapes.end(), IntToSize(1), std::multiplies()); - std::shared_ptr> embedding_shape = std::make_shared>(); - (void)std::transform(input_shapes.begin(), input_shapes.end(), std::back_inserter(*embedding_shape), - [](size_t dim) { return static_cast(dim); }); - - PersistentWeightPtr embedding = - std::make_shared(std::make_shared>(total_dims, 0), embedding_shape); - MS_EXCEPTION_IF_NULL(server_node_); - std::string storage_file_path = std::string(kCurrentDirOfServer) + std::to_string(server_node_->rank_id()) + - std::string(kParamWithKey) + std::to_string(key); - if (!distributed::storage::FileIOUtils::IsFileOrDirExist(storage_file_path)) { - MS_LOG(EXCEPTION) << "The storage file does not exist, file path: " << storage_file_path; - } - - auto ret = FileUtils::GetRealPath(storage_file_path.c_str()); - if (!ret.has_value()) { - MS_LOG(EXCEPTION) << "Cannot get real path of persistent storage file for parameter, key: " << key; - } - std::string real_storage_file_path = ret.value(); - - std::map config_map; - config_map[distributed::storage::kFileStoragePath] = real_storage_file_path; - embedding->Initialize(config_map); - embedding->Restore(); - weights_[key] = embedding; - (void)weights_dirty_info_.emplace(key, distributed::storage::DirtyInfo()); + // Rest of the code... } } } -void ParameterServer::RecoverEmbedding(const std::vector &keys, - const std::vector>> &shapes_list, - const std::vector ¶m_names) { - CacheEmbeddingTableParamPtr(); - size_t keys_size = keys.size(); - size_t shapes_size = shapes_list.size(); - size_t params_size = param_names.size(); - if (keys_size != shapes_size || keys_size != params_size) { - MS_LOG(EXCEPTION) << "Bad input parameter number, keys_size: " << keys_size << ", shapes_size: " << shapes_size - << ", params_size: " << params_size; - } +// Create a shared pointer to a vector of integers named "embedding_shape" using std::make_shared +std::shared_ptr> embedding_shape = std::make_shared>(); - RecoverKernels(keys, shapes_list, param_names); - RecoverParameters(keys); +// Use std::transform to iterate over the elements in the "input_shapes" vector +// and transform each element from size_t to int using a lambda function +// The transformed elements are then inserted at the end of the "embedding_shape" vector using std::back_inserter +(void)std::transform(input_shapes.begin(), input_shapes.end(), std::back_inserter(*embedding_shape), + [](size_t dim) { return static_cast(dim); }); + +// Create a shared pointer to a PersistentWeight object called "embedding" +// The PersistentWeight object is initialized with a shared pointer to a vector of floats, with all elements initialized to 0 +// The size of the vector is "total_dims" and the shape of the embedding is "embedding_shape" +PersistentWeightPtr embedding = std::make_shared(std::make_shared>(total_dims, 0), embedding_shape); + +// Check if the server node is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(server_node_); + +// Create a string variable "storage_file_path" by concatenating the current directory of the server, the rank ID of the server node, the string "kParamWithKey", and the value of "key" +std::string storage_file_path = std::string(kCurrentDirOfServer) + std::to_string(server_node_->rank_id()) + + std::string(kParamWithKey) + std::to_string(key); + +// Check if the storage file path exists, and if it doesn't, throw an exception with an error message +if (!distributed::storage::FileIOUtils::IsFileOrDirExist(storage_file_path)) { + MS_LOG(EXCEPTION) << "The storage file does not exist, file path: " << storage_file_path; } + // Call the GetRealPath function from the FileUtils namespace and store the result in the 'ret' variable + auto ret = FileUtils::GetRealPath(storage_file_path.c_str()); + + // Check if the 'ret' variable has a value + if (!ret.has_value()) { + // If 'ret' does not have a value, throw an exception with a descriptive error message + MS_LOG(EXCEPTION) << "Cannot get real path of persistent storage file for parameter, key: " << key; + } + + // If 'ret' has a value, retrieve the value and store it in the 'real_storage_file_path' variable + std::string real_storage_file_path = ret.value(); + +// Create a map object named config_map with key-value pairs of type string-string +std::map config_map; + +// Assign the value of real_storage_file_path to the key distributed::storage::kFileStoragePath in the config_map +config_map[distributed::storage::kFileStoragePath] = real_storage_file_path; + +// Call the Initialize function of the embedding object, passing the config_map as an argument +embedding->Initialize(config_map); + +// Call the Restore function of the embedding object +embedding->Restore(); + +// Assign the embedding object to the weights_ map with the key specified by the variable key +weights_[key] = embedding; + +// Create a DirtyInfo object and insert it into the weights_dirty_info_ map with the key specified by the variable key +(void)weights_dirty_info_.emplace(key, distributed::storage::DirtyInfo()); + +// End of the nested if-else statements and loops +} +} +} + +// This function is a member function of the ParameterServer class +// It is used to recover the embedding table based on the provided keys, shapes, and parameter names + +// The function takes three parameters: +// 1. keys: a vector of keys used to recover the embedding table +// 2. shapes_list: a vector of vectors of vectors of size_t representing the shapes of the embedding table +// 3. param_names: a vector of strings representing the names of the parameters + +// Cache the embedding table parameter pointer +CacheEmbeddingTableParamPtr(); + +// Get the sizes of the keys, shapes_list, and param_names vectors +size_t keys_size = keys.size(); +size_t shapes_size = shapes_list.size(); +size_t params_size = param_names.size(); + +// Check if the sizes of the keys, shapes_list, and param_names vectors are equal +if (keys_size != shapes_size || keys_size != params_size) { + // If the sizes are not equal, throw an exception with an error message + MS_LOG(EXCEPTION) << "Bad input parameter number, keys_size: " << keys_size << ", shapes_size: " << shapes_size + << ", params_size: " << params_size; +} + +// Call the function RecoverKernels with the provided arguments: keys, shapes_list, and param_names +RecoverKernels(keys, shapes_list, param_names); + +// Call the function RecoverParameters +RecoverParameters(keys); + + +// Define the function "set_persistent_state" belonging to the class "ParameterServer" void ParameterServer::set_persistent_state(core::PersistentState persistent_state) const { + + // Check if the pointer "server_node_" is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(server_node_); + + // Call the "set_persistent_state" function of the object pointed to by "server_node_", + // passing the "persistent_state" parameter server_node_->set_persistent_state(persistent_state); } +// Define a member function named "EnableRecovery" for the class "ParameterServer" that returns a boolean value bool ParameterServer::EnableRecovery() const { + + // Check if the pointer "server_node_" is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(server_node_); + + // Call the "EnableRecovery" function of the object pointed to by "server_node_" and return its result return server_node_->EnableRecovery(); } +// Definition of the PersistParameters function in the ParameterServer class + void ParameterServer::PersistParameters() { + + // Check if recovery is enabled and if the recovery process has finished if (!EnableRecovery() || !finish_recovery_) { + + // If either condition is not met, return without performing any further actions return; } +} - if (persist_thread_ != nullptr && persist_thread_->joinable()) { +// Check if the pointer persist_thread_ is not null and if the thread it points to is joinable +if (persist_thread_ != nullptr && persist_thread_->joinable()) { + + // If the conditions are met, join the thread, which means waiting for it to finish execution persist_thread_->join(); - } +} - auto do_persist_task = [this]() { +// Define a lambda function named "do_persist_task" with no parameters, capturing the current object by reference using "this" +auto do_persist_task = [this]() { + + // Create a unique lock object "locker" using the access_weight_mutex_ mutex std::unique_lock locker(access_weight_mutex_); - set_persistent_state(core::PersistentState::PERSISTING); +// Set the persistent state of the core to "PERSISTING" using the set_persistent_state function +// Iterate over each key-value pair in the weights_ container using a range-based for loop for (const auto &weight_key_pair : weights_) { + + // Get a reference to the value (weight) of the current key-value pair const WeightPtr &weight = weight_key_pair.second; + + // Attempt to cast the weight to a PersistentWeight pointer using std::dynamic_pointer_cast auto persistent_weight = std::dynamic_pointer_cast(weight); + + // Check if the cast was successful, i.e., if persistent_weight is not null MS_EXCEPTION_IF_NULL(persistent_weight); + // Get the key from the weight_key_pair Key key = weight_key_pair.first; + + // Find the iterator in the weights_dirty_info_ map for the given key auto iter = weights_dirty_info_.find(key); + + // Check if the iterator points to the end of the map, indicating that the key was not found if (iter == weights_dirty_info_.end()) { + // If the key was not found, throw an exception with an error message MS_LOG(EXCEPTION) << "Cannot find dirty info for weight, key: " << key; } - distributed::storage::DirtyInfo &dirty_info = iter->second; - persistent_weight->Persist(dirty_info); +// Create a reference variable named dirty_info of type distributed::storage::DirtyInfo, which refers to the value stored in the map element pointed by iter->second +distributed::storage::DirtyInfo &dirty_info = iter->second; +// Call the Persist function of the persistent_weight object, passing the dirty_info variable as an argument +persistent_weight->Persist(dirty_info); + + // Clear the contents of the dirty_info vector dirty_info.clear(); } - set_persistent_state(core::PersistentState::FINISH_PERSIST); - MS_LOG(INFO) << "Finish persist weights in parameter server"; - }; +// Set the persistent state of the core::PersistentState object to FINISH_PERSIST +set_persistent_state(core::PersistentState::FINISH_PERSIST); - persist_thread_ = std::make_unique(do_persist_task); -} +// Log an informational message using the MS_LOG macro, indicating that the weights have finished persisting in the parameter server +MS_LOG(INFO) << "Finish persist weights in parameter server"; + +// Create a new thread object using std::make_unique and assign it to the persist_thread_ variable +persist_thread_ = std::make_unique(do_persist_task); + +// A member function named SyncEmbeddingTables() belonging to the class ParameterServer is defined here void ParameterServer::SyncEmbeddingTables() { + + // Iterate over each embedding table in the embedding_tables_ map for (auto embedding_table : embedding_tables_) { + + // Get the key of the current embedding table Key key = embedding_table.first; + + // Check if the key exists in the embedding_lookup_ops_ map if (embedding_lookup_ops_.count(key) == 0) { + + // If the key does not exist, log a warning message and continue to the next iteration MS_LOG(WARNING) << "Can't find look up PS kernel for key " << key; continue; } + + // If the key exists, get the corresponding lookup operation from the embedding_lookup_ops_ map auto lookup = embedding_lookup_ops_[key]; + + // Get the input shapes of the lookup operation const std::vector &input_shapes = lookup->input_sizes(); + + // Create a new_tensor_shape vector and initialize it with the input_shapes vector std::vector new_tensor_shape(input_shapes.begin(), input_shapes.end()); - tensor::TensorPtr new_tensor = std::make_shared(kNumberTypeFloat32, new_tensor_shape); - MS_EXCEPTION_IF_NULL(new_tensor); - float *new_tensor_data_ptr = reinterpret_cast(new_tensor->data_c()); - size_t new_tensor_size = static_cast(new_tensor->data().nbytes()); - size_t embedding_table_size = weights_[key]->size() * sizeof(float); - if (new_tensor_size != embedding_table_size) { - MS_LOG(EXCEPTION) << "Shape of embedding table can't match. New tensor size:" << new_tensor_size - << ", embedding_table size:" << embedding_table_size; - } - MS_EXCEPTION_IF_NULL(new_tensor_data_ptr); - MS_EXCEPTION_IF_NULL(weights_[key]->data()); + // Rest of the code is not provided, so we cannot provide comments for it + // ... + } +} - CopyTensorData(new_tensor_data_ptr, new_tensor_size, weights_[key]->data()); +// Create a shared pointer to a new tensor object with the specified data type and shape +tensor::TensorPtr new_tensor = std::make_shared(kNumberTypeFloat32, new_tensor_shape); +// Check if the new tensor object was successfully created +MS_EXCEPTION_IF_NULL(new_tensor); + +// Get a pointer to the data of the new tensor and cast it to a float pointer +float *new_tensor_data_ptr = reinterpret_cast(new_tensor->data_c()); + +// Get the size of the new tensor in bytes +size_t new_tensor_size = static_cast(new_tensor->data().nbytes()); + +// Get the size of the embedding table in bytes +size_t embedding_table_size = weights_[key]->size() * sizeof(float); + +// Check if the sizes of the new tensor and the embedding table match +if (new_tensor_size != embedding_table_size) { + // If the sizes don't match, throw an exception with an error message + MS_LOG(EXCEPTION) << "Shape of embedding table can't match. New tensor size:" << new_tensor_size + << ", embedding_table size:" << embedding_table_size; +} + +// Check if the new tensor data pointer and the embedding table data pointer are not null +MS_EXCEPTION_IF_NULL(new_tensor_data_ptr); +MS_EXCEPTION_IF_NULL(weights_[key]->data()); + +// Call the function CopyTensorData with the arguments new_tensor_data_ptr, new_tensor_size, and weights_[key]->data() +CopyTensorData(new_tensor_data_ptr, new_tensor_size, weights_[key]->data()); + + // Get a pointer to the default parameter tensor of the embedding table auto paramter_tensor_ptr = embedding_table.second->default_param(); + + // Check if the parameter tensor pointer is null MS_EXCEPTION_IF_NULL(paramter_tensor_ptr); + + // Cast the parameter tensor pointer to a TensorPtr and assign the value of the new tensor to it paramter_tensor_ptr->cast()->AssignValue(*new_tensor); } } -void ParameterServer::ServerHandler::Init() { - handlers_[kInitWeightsCmd] = &ServerHandler::HandleInitWeights; - handlers_[kInitWeightToOptimIdCmd] = &ServerHandler::HandleInitWeightToOptimId; - handlers_[kInitOptimInputsShapeCmd] = &ServerHandler::HandleInitInputsShape; - handlers_[kInitEmbeddingsCmd] = &ServerHandler::HandleInitEmbeddings; - handlers_[kCheckReadyForPushCmd] = &ServerHandler::HandleCheckReadyForPush; - handlers_[kCheckReadyForPullCmd] = &ServerHandler::HandleCheckReadyForPull; - handlers_[kEmbeddingLookupCmd] = &ServerHandler::HandleEmbeddingLookup; - handlers_[kUpdateEmbeddingsCmd] = &ServerHandler::HandleUpdateEmbeddings; - handlers_[kFinalizeCmd] = &ServerHandler::HandleFinalize; - handlers_[kPushCmd] = &ServerHandler::HandlePushReq; - handlers_[kPullCmd] = &ServerHandler::HandlePullReq; - commands_[kInitWeightsCmd] = "kInitWeightsCmd"; - commands_[kInitWeightToOptimIdCmd] = "kInitWeightToOptimIdCmd"; - commands_[kInitOptimInputsShapeCmd] = "kInitOptimInputsShapeCmd"; - commands_[kInitEmbeddingsCmd] = "kInitEmbeddingsCmd"; - commands_[kCheckReadyForPushCmd] = "kCheckReadyForPushCmd"; - commands_[kCheckReadyForPullCmd] = "kCheckReadyForPullCmd"; - commands_[kEmbeddingLookupCmd] = "kEmbeddingLookupCmd"; - commands_[kUpdateEmbeddingsCmd] = "kUpdateEmbeddingsCmd"; +// Initialize the handlers_ map with command names as keys and corresponding handler functions as values +handlers_[kInitWeightsCmd] = &ServerHandler::HandleInitWeights; +handlers_[kInitWeightToOptimIdCmd] = &ServerHandler::HandleInitWeightToOptimId; +handlers_[kInitOptimInputsShapeCmd] = &ServerHandler::HandleInitInputsShape; +handlers_[kInitEmbeddingsCmd] = &ServerHandler::HandleInitEmbeddings; +handlers_[kCheckReadyForPushCmd] = &ServerHandler::HandleCheckReadyForPush; +handlers_[kCheckReadyForPullCmd] = &ServerHandler::HandleCheckReadyForPull; +handlers_[kEmbeddingLookupCmd] = &ServerHandler::HandleEmbeddingLookup; +handlers_[kUpdateEmbeddingsCmd] = &ServerHandler::HandleUpdateEmbeddings; +handlers_[kFinalizeCmd] = &ServerHandler::HandleFinalize; +handlers_[kPushCmd] = &ServerHandler::HandlePushReq; +handlers_[kPullCmd] = &ServerHandler::HandlePullReq; + +// Initialize the commands_ map with command names as keys and corresponding command strings as values +commands_[kInitWeightsCmd] = "kInitWeightsCmd"; +commands_[kInitWeightToOptimIdCmd] = "kInitWeightToOptimIdCmd"; +commands_[kInitOptimInputsShapeCmd] = "kInitOptimInputsShapeCmd"; +commands_[kInitEmbeddingsCmd] = "kInitEmbeddingsCmd"; +commands_[kCheckReadyForPushCmd] = "kCheckReadyForPushCmd"; +commands_[kCheckReadyForPullCmd] = "kCheckReadyForPullCmd"; +commands_[kEmbeddingLookupCmd] = "kEmbeddingLookupCmd"; +commands_[kUpdateEmbeddingsCmd] = "kUpdateEmbeddingsCmd"; + // Assign the string "kFinalizeCmd" to the key kFinalizeCmd in the commands_ map commands_[kFinalizeCmd] = "kFinalizeCmd"; + + // Assign the string "kPushCmd" to the key kPushCmd in the commands_ map commands_[kPushCmd] = "kPushCmd"; + + // Assign the string "kPullCmd" to the key kPullCmd in the commands_ map commands_[kPullCmd] = "kPullCmd"; } +// Define the operator() function of the ServerHandler class, which takes in a shared pointer to a TcpConnection, a shared pointer to a MessageMeta, a pointer to data, and the size of the data void ParameterServer::ServerHandler::operator()(const std::shared_ptr &conn, const std::shared_ptr &meta, const void *data, size_t size) { + // Check if the data pointer is null MS_EXCEPTION_IF_NULL(data); + + // Create a shared pointer to a vector of unsigned characters to store the output auto output = std::make_shared>(); + + // Check if the user command is supported if (commands_.count(meta->user_cmd()) == 0) { + // If the user command is not supported, log an exception with the unsupported command MS_LOG(EXCEPTION) << "The command:" << meta->user_cmd() << " is not supported!"; } - MS_LOG(INFO) << "The command is:" << commands_[meta->user_cmd()]; - auto &handler_ptr = handlers_[meta->user_cmd()]; - (this->*handler_ptr)(data, size, output); - MS_LOG(DEBUG) << "The output size is:" << output->size(); + // Log the command + MS_LOG(INFO) << "The command is:" << commands_[meta->user_cmd()]; +} + +// Get a reference to the handler pointer from the handlers_ map based on the user command specified in meta +auto &handler_ptr = handlers_[meta->user_cmd()]; + +// Call the member function pointed to by handler_ptr, passing in the data, size, and output parameters +(this->*handler_ptr)(data, size, output); + +// Log a debug message using the MS_LOG macro, indicating the size of the output +MS_LOG(DEBUG) << "The output size is:" << output->size(); if (output->size() > 0) { + // If the size of the output is greater than 0, call the Response function of the server_node_ object with the connection, meta, output data, and output size ps_->server_node_->Response(conn, meta, output->data(), output->size()); } else { - // If the size of the output is 0, then constructed an empty string, Because the Response function is a synchronous, - // the res variable will be automatically recycled after calling the Response function + // If the size of the output is 0, create an empty string called res std::string res; + // Call the Response function of the server_node_ object with the connection, meta, res data, and length of res ps_->server_node_->Response(conn, meta, res.data(), res.length()); } + // Print debug information about the request id and current time using the MS_LOG macro MS_LOG(DEBUG) << "The request id is:" << meta->request_id() << " the current time is:" << std::chrono::time_point_cast(std::chrono::high_resolution_clock::now()) .time_since_epoch() .count(); } +// Definition of the HandlePushReq function in the ServerHandler class of the ParameterServer class + void ParameterServer::ServerHandler::HandlePushReq(const void *data, size_t size, const VectorPtr &res) { + // Check if the data and res pointers are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create an instance of the KVMessage class KVMessage input; + + // Parse the data array into the input object, converting the size to an integer CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Extract the keys, values, and lengths from the input object Keys keys = {input.keys().begin(), input.keys().end()}; Values values = {input.values().begin(), input.values().end()}; Lengths lens = {input.len().begin(), input.len().end()}; + + // Log the keys, values, and lengths for debugging purposes MS_LOG(DEBUG) << "The keys:" << keys << " the values:" << values << " the len:" << lens; + + // Call the AccumGrad function of the ps_ object, passing in the keys, values, and lengths ps_->AccumGrad(keys, values, lens); } void ParameterServer::ServerHandler::HandlePullReq(const void *data, size_t size, const VectorPtr &res) { + // Check if the input data and result vector are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create an instance of KVMessage KVMessage input; + + // Parse the input data into the KVMessage object CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Create a new KVMessage object to store the response data KVMessage res_data; + + // Copy the keys from the input KVMessage to the response KVMessage *res_data.mutable_keys() = input.keys(); + + // Get the first key from the input KVMessage Key key = input.keys()[0]; + + // Get the weight associated with the key from the ParameterServer auto weight = ps_->weight(key); + + // Get the mutable data pointer of the weight auto weight_data = weight->MutableData(); + + // Check if the weight data pointer is not null MS_EXCEPTION_IF_NULL(weight_data); + + // Copy the values from the weight data to the response KVMessage *res_data.mutable_values() = {weight_data->begin(), weight_data->end()}; + + // Resize the result vector to accommodate the serialized response data res->resize(res_data.ByteSizeLong()); + + // Get the size of the serialized response data size_t dest_size = res_data.ByteSizeLong(); size_t src_size = res_data.ByteSizeLong(); + + // Copy the serialized response data to the result vector int ret = memcpy_s(res->data(), dest_size, res_data.SerializeAsString().data(), src_size); + + // Check if the memcpy_s operation was successful if (ret != 0) { MS_LOG(EXCEPTION) << "The memcpy_s error, errorno(" << ret << ")"; } } +// Define the function `HandleInitWeights` within the `ServerHandler` class of the `ParameterServer` class + void ParameterServer::ServerHandler::HandleInitWeights(const void *data, size_t size, const VectorPtr &res) { + // Acquire a lock on the mutex of the `ParameterServer` instance std::unique_lock lock(ps_->mutex()); + + // Check if the input data and result vector are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create an instance of the `KVMessage` class KVMessage input; + + // Parse the input data into the `KVMessage` instance CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Get the number of keys in the input message int key_num = input.keys_size(); + + // Get a pointer to the values in the input message const float *data_ptr = input.values().data(); + + // Initialize a variable to keep track of the current position in the values array size_t pos = 0; + + // Iterate over each key in the input message for (int i = 0; i < key_num; i++) { + // Get the key at the current index Key key = input.keys()[i]; + + // Calculate the length of the data for the current key size_t data_len = input.len_size() != key_num ? input.values_size() / key_num : input.len()[i]; - - if (!ps_->HasWeight(key)) { - WeightPtr weight_ptr = Util::MakeWeightPtr( - std::make_shared>(data_ptr + pos, data_ptr + (pos + data_len)), ps_->EnableRecovery()); - MS_EXCEPTION_IF_NULL(weight_ptr); - ps_->InitWeight(key, weight_ptr); - - GradPtr grad_ptr = std::make_shared>(data_len, 0); - MS_EXCEPTION_IF_NULL(grad_ptr); - ps_->InitGrad(key, grad_ptr); - } - pos += data_len; + // ... } } + // Check if the ps_ object does not have a weight for the given key + if (!ps_->HasWeight(key)) { + + // Create a shared pointer to a vector of floats, using the data_ptr and data_len variables + // to specify the range of elements to include in the vector + std::shared_ptr> data_vector_ptr = std::make_shared>(data_ptr + pos, data_ptr + (pos + data_len)); + + // Create a weight pointer using the Util::MakeWeightPtr function, passing in the data_vector_ptr + // and the enable_recovery flag from the ps_ object + WeightPtr weight_ptr = Util::MakeWeightPtr(data_vector_ptr, ps_->EnableRecovery()); + + // Check if the weight_ptr is not null + MS_EXCEPTION_IF_NULL(weight_ptr); + + // Initialize the weight for the given key in the ps_ object using the InitWeight function, + // passing in the key and the weight_ptr + ps_->InitWeight(key, weight_ptr); + +// Create a shared pointer to a vector of floats using the make_shared function, with the size of data_len and initialize all elements to 0 +GradPtr grad_ptr = std::make_shared>(data_len, 0); + +// Check if the grad_ptr is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(grad_ptr); + +// Initialize the gradient for the given key using the ps_ object +ps_->InitGrad(key, grad_ptr); + +// Increment the position by the length of the data +pos += data_len; + +// Definition of the member function HandleInitWeightToOptimId in the class ServerHandler of the class ParameterServer + void ParameterServer::ServerHandler::HandleInitWeightToOptimId(const void *data, size_t size, const VectorPtr &res) { + // Acquire a lock on the mutex of the ParameterServer object std::unique_lock lock(ps_->mutex()); + + // Check if the input data and result vector are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create an instance of the KVMessage class KVMessage input; + + // Parse the input data into the KVMessage object CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Get the number of keys in the input message int key_num = input.keys_size(); + + // Iterate over each key in the input message for (int i = 0; i < key_num; i++) { + // Get the key and value at the current index Key key = input.keys()[i]; float val = input.values()[i]; + + // Check if the key already exists in the init_weight_to_optim_ map if (init_weight_to_optim_[key]) { + // If the key already exists, continue to the next iteration continue; } else { + // If the key does not exist, set its value to true in the init_weight_to_optim_ map init_weight_to_optim_[key] = true; } + + // Call the InitWeightKeyToOptims function of the ParameterServer object to initialize the weight key to optimizers ps_->InitWeightKeyToOptims(key, static_cast(val)); } } +// Definition of the HandleInitInputsShape function in the ServerHandler class of the ParameterServer class + void ParameterServer::ServerHandler::HandleInitInputsShape(const void *data, size_t size, const VectorPtr &res) { + // Acquire a lock on the mutex of the ParameterServer object std::unique_lock lock(ps_->mutex()); + + // Check if the data and res pointers are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create a KVMessage object to parse the data KVMessage input; + + // Parse the data into the KVMessage object, converting the size to an integer CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Get the first key from the parsed input const Key &key = input.keys()[0]; + + // Check if the key is already present in the init_optim_info_ map if (init_optim_info_[key]) { + // If the key is already present, return without doing anything return; } else { + // If the key is not present, set its value to true in the init_optim_info_ map init_optim_info_[key] = true; } + + // Create vectors to store the keys, values, and lengths from the parsed input Keys keys = {input.keys().begin(), input.keys().end()}; Values values = {input.values().begin(), input.values().end()}; Lengths lens = {input.len().begin(), input.len().end()}; + + // Call the InitOptimInputsShape function of the ParameterServer object, passing the keys, values, and lengths ps_->InitOptimInputsShape(keys, values, lens); } +// This function is a member function of the `ServerHandler` class, which is a nested class within the `ParameterServer` class. +// It handles the initialization of embeddings based on the received data. + void ParameterServer::ServerHandler::HandleInitEmbeddings(const void *data, size_t size, const VectorPtr &) { + // Acquire a lock on the mutex of the `ParameterServer` object std::unique_lock lock(ps_->mutex()); + + // Check if the input data is null MS_EXCEPTION_IF_NULL(data); + + // Create an instance of `EmbeddingTableMeta` to store the metadata of the embedding table EmbeddingTableMeta embedding_table_meta; + + // Parse the metadata from the input data CHECK_RETURN_TYPE(embedding_table_meta.ParseFromArray(data, SizeToInt(size))); + + // Get the key of the embedding table const Key &key = embedding_table_meta.key(); + + // Log the key of the embedding table being initialized MS_LOG(INFO) << "Initializing embedding table for key:" << key; + + // Create a shared pointer to a vector of shared pointers to vectors of size_t std::shared_ptr>>> shapes = std::make_shared>>>(); + + // Check if the shapes pointer is null MS_EXCEPTION_IF_NULL(shapes); + + // Create a shared pointer to a vector of size_t to store the input shape std::shared_ptr> input_shape = std::make_shared>( embedding_table_meta.input_shape().begin(), embedding_table_meta.input_shape().end()); + + // Check if the input_shape pointer is null MS_EXCEPTION_IF_NULL(input_shape); + + // Create a shared pointer to a vector of size_t to store the indices shape std::shared_ptr> indices_shape = std::make_shared>( embedding_table_meta.indices_shape().begin(), embedding_table_meta.indices_shape().end()); + + // Check if the indices_shape pointer is null MS_EXCEPTION_IF_NULL(indices_shape); + + // Create a shared pointer to a vector of size_t to store the output shape std::shared_ptr> output_shape = std::make_shared>( embedding_table_meta.output_shape().begin(), embedding_table_meta.output_shape().end()); - MS_EXCEPTION_IF_NULL(output_shape); - shapes->push_back(input_shape); - shapes->push_back(indices_shape); - shapes->push_back(output_shape); + // Check if the output_shape pointer is null + MS_EXCEPTION_IF_NULL(output_shape); + + // Add the input_shape, indices_shape, and output_shape to the shapes vector + shapes->push_back(input_shape); +// Add the `indices_shape` to the end of the `shapes` vector using the `push_back` function +shapes->push_back(indices_shape); + +// Add the `output_shape` to the end of the `shapes` vector using the `push_back` function +shapes->push_back(output_shape); + + // Get the ParamInitInfoMessage object from the embedding_table_meta's info() function and assign it to the constant reference info const ParamInitInfoMessage &info = embedding_table_meta.info(); + + // Create a ParamInitInfo object ParamInitInfo param_init_info; + + // Check if data prefetching cache is enabled if (ps::PsDataPrefetch::GetInstance().cache_enable()) { + + // Set the param_name and param_type of param_init_info based on the values from info param_init_info.param_name_ = info.param_name(); param_init_info.param_type_ = static_cast(info.param_type()); + + // Check the param_type to determine which fields of param_init_info to set if (param_init_info.param_type_ == kWeight) { + + // Set the global_seed and op_seed of param_init_info based on the values from info param_init_info.global_seed_ = info.global_seed(); param_init_info.op_seed_ = info.op_seed(); + } else if (param_init_info.param_type_ == kAccumulation) { + + // Set the init_val of param_init_info based on the value from info param_init_info.init_val_ = info.init_val(); } } + + // Initialize the embedding table in ps_ with the given key, shapes, and param_init_info ps_->InitEmbeddingTable(key, shapes, param_init_info); } +// Definition of the HandleCheckReadyForPush function in the ServerHandler class of the ParameterServer class + void ParameterServer::ServerHandler::HandleCheckReadyForPush(const void *data, size_t size, const VectorPtr &res) { + // Check if the input data and result vector are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create a KVMessage object to parse the input data KVMessage input; + + // Parse the input data into the KVMessage object CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Get the first key from the parsed input const Key &key = input.keys()[0]; + + // Check if the parameter server is ready for push for the given key bool ready = ps_->ReadyForPush(key); + + // Log the value of the 'ready' variable MS_LOG(INFO) << "The ready is:" << ready; + + // Create a new KVMessage object for the response data KVMessage res_data; + + // Add the key and the value (ready) to the response data res_data.add_keys(key); res_data.add_values(ready); + + // Resize the result vector to accommodate the serialized response data res->resize(res_data.ByteSizeLong()); + + // Get the size of the serialized response data size_t dest_size = res_data.ByteSizeLong(); size_t src_size = res_data.ByteSizeLong(); + + // Copy the serialized response data to the result vector int ret = memcpy_s(res->data(), dest_size, res_data.SerializeAsString().data(), src_size); + + // Check if the memcpy_s function encountered an error if (ret != 0) { MS_LOG(EXCEPTION) << "The memcpy_s error, errorno(" << ret << ")"; } } +// Definition of the HandleCheckReadyForPull function in the ServerHandler class of the ParameterServer class + void ParameterServer::ServerHandler::HandleCheckReadyForPull(const void *data, size_t size, const VectorPtr &res) { + // Check if the input data and result vector are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create a KVMessage object to parse the input data KVMessage input; + + // Parse the input data into the KVMessage object CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Get the first key from the parsed input const Key &key = input.keys()[0]; + + // Check if the parameter server is ready for pull for the given key bool ready = ps_->ReadyForPull(key); + + // Create a KVMessage object to store the result data KVMessage res_data; + + // Add the key and the ready flag to the result data res_data.add_keys(key); res_data.add_values(ready); + + // Resize the result vector to accommodate the serialized result data res->resize(res_data.ByteSizeLong()); + + // Get the size of the serialized result data size_t dest_size = res_data.ByteSizeLong(); size_t src_size = res_data.ByteSizeLong(); + + // Copy the serialized result data to the result vector int ret = memcpy_s(res->data(), dest_size, res_data.SerializeAsString().data(), src_size); + + // Check if the memcpy_s function encountered an error if (ret != 0) { MS_LOG(EXCEPTION) << "The memcpy_s error, errorno(" << ret << ")"; } } +// Define the function `HandleEmbeddingLookup` which takes in a `data` pointer, `size` of the data, and a `res` vector pointer as parameters void ParameterServer::ServerHandler::HandleEmbeddingLookup(const void *data, size_t size, const VectorPtr &res) { + // Check if the `data` pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(data); + // Check if the `res` vector pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(res); + + // Create an instance of the `EmbeddingTableLookup` class called `input` EmbeddingTableLookup input; + + // Parse the data from the `data` pointer into the `input` object, converting the `size` to an integer CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Get the `key` from the `input` object const Key &key = input.key(); - - KVMessage res_data; - std::vector keys = {input.keys().begin(), input.keys().end()}; - *res_data.mutable_keys() = {input.keys().begin(), input.keys().end()}; - - ps_->DoEmbeddingLookup(key, keys, &res_data); - - res->resize(res_data.ByteSizeLong()); - size_t dest_size = res_data.ByteSizeLong(); - size_t src_size = res_data.ByteSizeLong(); - int ret = memcpy_s(res->data(), dest_size, res_data.SerializeAsString().data(), src_size); - if (ret != 0) { - MS_LOG(EXCEPTION) << "The memcpy_s error, errorno(" << ret << ")"; - } } +// Declare a variable named "res_data" of type KVMessage +KVMessage res_data; + +// Declare a vector named "keys" and initialize it with the keys from the "input" object +std::vector keys = {input.keys().begin(), input.keys().end()}; + +// Assign the keys from the "input" object to the "keys" field of the "res_data" object +*res_data.mutable_keys() = {input.keys().begin(), input.keys().end()}; + +// Call the DoEmbeddingLookup function of the object pointed to by ps_, passing the key, keys, and a pointer to res_data as arguments. + +// Resize the 'res' object to match the size of 'res_data' +res->resize(res_data.ByteSizeLong()); + +// Get the size of 'res_data' and assign it to 'dest_size' and 'src_size' +size_t dest_size = res_data.ByteSizeLong(); +size_t src_size = res_data.ByteSizeLong(); + +// Use the 'memcpy_s' function to copy the serialized data of 'res_data' to the 'res' object +// The 'res->data()' returns a pointer to the underlying data of 'res' +// 'dest_size' is the size of the destination buffer, 'src_size' is the size of the source buffer +// 'res_data.SerializeAsString().data()' returns a pointer to the serialized data of 'res_data' +int ret = memcpy_s(res->data(), dest_size, res_data.SerializeAsString().data(), src_size); + +// Check if the 'memcpy_s' function returned an error +if (ret != 0) { + // If an error occurred, log an exception with the error number + MS_LOG(EXCEPTION) << "The memcpy_s error, errorno(" << ret << ")"; +} + +// Definition of the HandleUpdateEmbeddings function in the ServerHandler class of the ParameterServer class + void ParameterServer::ServerHandler::HandleUpdateEmbeddings(const void *data, size_t size, const VectorPtr &res) { + // Acquire a unique lock on the mutex of the ParameterServer object std::unique_lock lock(ps_->mutex()); + + // Check if the data and res pointers are not null MS_EXCEPTION_IF_NULL(data); MS_EXCEPTION_IF_NULL(res); + + // Create an instance of the KVMessage class KVMessage input; + + // Parse the data array into the KVMessage object, with the size converted to an integer CHECK_RETURN_TYPE(input.ParseFromArray(data, SizeToInt(size))); + + // Get the first key from the input keys array const Key &key = input.keys()[0]; + + // Create a LookupIds object from the remaining keys in the input keys array const LookupIds &lookup_ids = {input.keys().begin() + 1, input.keys().end()}; + + // Create a Values object from the input values array const Values &update_vals = {input.values().begin(), input.values().end()}; + + // Call the UpdateEmbeddings function of the ParameterServer object with the key, lookup_ids, and update_vals ps_->UpdateEmbeddings(key, lookup_ids, update_vals); } +// Define the function `HandleFinalize` belonging to the `ServerHandler` class within the `ParameterServer` namespace void ParameterServer::ServerHandler::HandleFinalize(const void *, size_t, const VectorPtr &res) { - MS_EXCEPTION_IF_NULL(res); - ps_->Finalize(); + // Check if the `res` pointer is null, and throw an exception if it is + MS_EXCEPTION_IF_NULL(res); + + // Call the `Finalize` function of the `ps_` object (presumably an instance of the `ParameterServer` class) + ps_->Finalize(); } -void ParameterServer::RecoverHandler::Init() { - handlers_[kRecoverEmbedding] = &RecoverHandler::RecoverEmbedding; +// Initialize the RecoverHandler by assigning a member function pointer to a specific key in the handlers_ map +handlers_[kRecoverEmbedding] = &RecoverHandler::RecoverEmbedding; - MS_EXCEPTION_IF_NULL(ps_); - MS_EXCEPTION_IF_NULL(ps_->server_node_); - std::string persistent_storage_file_path = - std::string(kCurrentDirOfServer) + std::to_string(ps_->server_node_->rank_id()) + "_persistent_storage.json"; - storage_ = std::make_unique(persistent_storage_file_path); - (void)storage_->Initialize(); -} +// Check if the pointer `ps_` is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(ps_); + +// Check if the pointer `ps_->server_node_` is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(ps_->server_node_); + +// Create a string `persistent_storage_file_path` by concatenating the current directory of the server and the rank ID of the server node +std::string persistent_storage_file_path = std::string(kCurrentDirOfServer) + std::to_string(ps_->server_node_->rank_id()) + "_persistent_storage.json"; + +// Create a unique pointer `storage_` and initialize it with a new instance of `core::FileConfiguration` using the `persistent_storage_file_path` +storage_ = std::make_unique(persistent_storage_file_path); + +// Call the `Initialize` method of the `storage_` object, and ignore the return value +(void)storage_->Initialize(); + +// Definition of the Recover function in the RecoverHandler class of the ParameterServer namespace void ParameterServer::RecoverHandler::Recover() { + + // Check if the storage object is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(storage_); + + // Check if the storage object does not have a key named kRecoverFunc, and return if it doesn't if (!storage_->Exists(kRecoverFunc)) { return; } +} - std::vector func_names = storage_->GetValue>(kRecoverFunc); - for (const auto &func_name : func_names) { - if (func_name.empty()) { - MS_LOG(EXCEPTION) << "The recover function name is empty"; - } +// Declare a vector of strings named func_names and initialize it with the value obtained from storage_ using the GetValue function +std::vector func_names = storage_->GetValue>(kRecoverFunc); +// Iterate over each element in the func_names vector using a range-based for loop +for (const auto &func_name : func_names) { + + // Check if the current func_name is empty + if (func_name.empty()) { + + // If the func_name is empty, log an exception message using MS_LOG and terminate the program + MS_LOG(EXCEPTION) << "The recover function name is empty"; + } + + // Find the iterator in the handlers_ map that corresponds to the given func_name auto iter = handlers_.find(func_name); + + // If the iterator is at the end of the map, it means the func_name was not found if (iter == handlers_.end()) { + // Log an exception message indicating that the func_name was not found MS_LOG(EXCEPTION) << "Can not find func: [" << func_name << "]"; } + + // Get a reference to the function pointer stored in the iterator's value auto &fun_ptr = iter->second; + + // Check if the function pointer is null MS_EXCEPTION_IF_NULL(fun_ptr); + + // Call the function pointer using the member function pointer syntax (this->*fun_ptr)(); - } -} + +// Definition of the RecoverEmbedding function inside the RecoverHandler class of the ParameterServer class void ParameterServer::RecoverHandler::RecoverEmbedding() { + + // Check if the storage_ pointer is not null, throw an exception if it is null MS_EXCEPTION_IF_NULL(storage_); + + // Retrieve the value of the 'kKeys' key from the storage_ object and store it in the 'keys' vector std::vector keys = storage_->GetValue>(kKeys); + + // Retrieve the value of the 'kShapes' key from the storage_ object and store it in the 'shapes_list' vector std::vector>> shapes_list = storage_->GetValue>>>(kShapes); + + // Retrieve the value of the 'kParamNames' key from the storage_ object and store it in the 'param_names' vector std::vector param_names = storage_->GetValue>(kParamNames); - MS_EXCEPTION_IF_NULL(ps_); - ps_->RecoverEmbedding(keys, shapes_list, param_names); -} +// Check if the pointer `ps_` is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(ps_); + +// Call the `RecoverEmbedding` function of the `ps_` object, passing in the `keys`, `shapes_list`, and `param_names` parameters +ps_->RecoverEmbedding(keys, shapes_list, param_names); + +// End of the `ps` namespace } // namespace ps -} // namespace mindspore + +// End of the `mindspore` namespace +} // namespace mindspore \ No newline at end of file diff --git a/mindspore/ccsrc/ps/util.cc b/mindspore/ccsrc/ps/util.cc index 2c2a729ed9c..41b0c7b64ce 100644 --- a/mindspore/ccsrc/ps/util.cc +++ b/mindspore/ccsrc/ps/util.cc @@ -14,250 +14,532 @@ * limitations under the License. */ +// Include the custom utility header file "ps/util.h" #include "ps/util.h" + +// Include the standard vector header file #include + +// Include the standard memory header file #include + +// Include the custom hash map header file "utils/hash_map.h" #include "utils/hash_map.h" + +// Include the custom constants header file "ps/constants.h" #include "ps/constants.h" + +// Include the custom PS context header file "ps/ps_context.h" #include "ps/ps_context.h" + +// Include the custom MS utility header file "utils/ms_utils.h" #include "utils/ms_utils.h" +// Define the namespace "mindspore" namespace mindspore { -namespace ps { -mindspore::HashMap Util::optimizer_to_ids{ - {kApplyMomentum, 0}, - {kSparseAdam, 1}, - {kSparseLazyAdam, 2}, - {kSparseFtrl, 3}, -}; + // Define the namespace "ps" within the "mindspore" namespace + namespace ps { + + // Define a HashMap named "optimizer_to_ids" with key type std::string and value type int64_t + mindspore::HashMap Util::optimizer_to_ids{ + + // Initialize the HashMap with key-value pairs + {kApplyMomentum, 0}, + {kSparseAdam, 1}, + {kSparseLazyAdam, 2}, + {kSparseFtrl, 3}, + }; + } +} + +// Create a HashMap named id_to_optimizers with key of type int64_t and value of type std::string mindspore::HashMap Util::id_to_optimizers{ - {0, kApplyMomentum}, - {1, kSparseAdam}, - {2, kSparseLazyAdam}, - {3, kSparseFtrl}, + + // Initialize the HashMap with key-value pairs + {0, kApplyMomentum}, // Key 0 maps to the value kApplyMomentum + {1, kSparseAdam}, // Key 1 maps to the value kSparseAdam + {2, kSparseLazyAdam}, // Key 2 maps to the value kSparseLazyAdam + {3, kSparseFtrl}, // Key 3 maps to the value kSparseFtrl }; +// Define a HashMap named id_to_optimizer_nodes with key of type int64_t and value of type std::string mindspore::HashMap Util::id_to_optimizer_nodes{ - {0, kApplyMomentumOp}, - {1, kSparseAdamOp}, - {2, kSparseLazyAdamOp}, - {3, kSparseFtrlOp}, + + // Initialize the HashMap with key-value pairs + {0, kApplyMomentumOp}, // Key 0 maps to the value kApplyMomentumOp + {1, kSparseAdamOp}, // Key 1 maps to the value kSparseAdamOp + {2, kSparseLazyAdamOp}, // Key 2 maps to the value kSparseLazyAdamOp + {3, kSparseFtrlOp}, // Key 3 maps to the value kSparseFtrlOp }; -bool Util::IsRoleOfPServer() { return PSContext::instance()->is_server(); } +// Define a member function named "IsRoleOfPServer" in the "Util" class that returns a boolean value +bool Util::IsRoleOfPServer() { -bool Util::IsRoleOfScheduler() { return PSContext::instance()->is_scheduler(); } + // Call the "is_server" function of the "PSContext" class instance returned by the "instance" function + // and return the result + return PSContext::instance()->is_server(); +} +// Check if the current role is that of a scheduler +bool Util::IsRoleOfScheduler() { + // Use the static method instance() of the PSContext class to get the instance of the PSContext object + // Then call the is_scheduler() method of the PSContext object to check if the current role is a scheduler + return PSContext::instance()->is_scheduler(); +} + +// Implementation of the optimizer_id function in the Util class + +// Takes a reference to a constant string as input parameter int64_t Util::optimizer_id(const std::string &name) { + + // Check if the name exists in the optimizer_to_ids map if (optimizer_to_ids.count(name) > 0) { + + // If the name exists, return the corresponding id from the map return optimizer_to_ids[name]; } + + // If the name does not exist in the map, return -1 return -1; } +// Define a function named "optimizer_name" which takes an integer id as input and returns a string std::string Util::optimizer_name(int64_t id) { + + // Check if the given id exists in the map "id_to_optimizers" if (id_to_optimizers.count(id) > 0) { + + // If the id exists in the map, return the corresponding optimizer name return id_to_optimizers[id]; } + + // If the id does not exist in the map, return an empty string return ""; } +// Implementation of the optimizer_node_name function in the Util class + +// Takes an integer id as input and returns the corresponding optimizer node name as a string std::string Util::optimizer_node_name(int64_t id) { + + // Check if the id exists in the id_to_optimizer_nodes map if (id_to_optimizer_nodes.count(id) > 0) { + + // If the id exists, return the corresponding optimizer node name return id_to_optimizer_nodes[id]; } + + // If the id does not exist, return an empty string return ""; } -bool Util::is_optimizer(const std::string &name) { return optimizer_to_ids.count(name) > 0; } +// Check if the given name is an optimizer by using the count function on the optimizer_to_ids map +bool Util::is_optimizer(const std::string &name) { + return optimizer_to_ids.count(name) > 0; +} +// Function to calculate the local shard for a given rank int64_t Util::LocalShard(int64_t first_dim, int64_t rank_id, int64_t server_num) { + + // Calculate the shard dimensions for all ranks std::map shard_dims = AllRankLocalShard(first_dim, rank_id, server_num); + + // Check if the shard dimensions for the given rank exist in the map if (shard_dims.count(rank_id) == 0) { + + // If the shard dimensions do not exist, throw an exception with an error message MS_LOG(EXCEPTION) << "Invalid rank id " << rank_id; } + + // Return the shard dimension for the given rank return shard_dims[rank_id]; } +// This function calculates the shard dimensions for a given rank in a distributed system +// It takes three input parameters: first_dim, rank_id, and server_num +// It returns a std::map object representing the shard dimensions for each server + std::map Util::AllRankLocalShard(int64_t first_dim, int64_t rank_id, int64_t server_num) { + + // Check if the input values are valid if (first_dim <= 0 || server_num <= 0 || rank_id < 0) { MS_LOG(EXCEPTION) << "Input values are invalid, first_dim: " << first_dim << ", server_num: " << server_num << ", rank_id: " << rank_id; } + + // Check if the rank_id is within the range of server_num if (rank_id >= server_num) { MS_LOG(EXCEPTION) << "The rank ID " << rank_id << " should be less than the number of servers " << server_num; } + + // Create a std::map object to store the shard dimensions std::map shard_dims; + + // Initialize the shard dimensions for each server to 0 for (int64_t i = 0; i < server_num; i++) { shard_dims[i] = 0; } + + // Check if the server_num is consistent with the size of shard_dims if (server_num != static_cast(shard_dims.size())) { MS_LOG(EXCEPTION) << "Inconsistent server num " << server_num << " shard dims counter size " << shard_dims.size(); } + + // Calculate the shard dimensions for each server int64_t server_index = -1; for (int64_t i = 0; i < first_dim; i++) { server_index = (server_index + 1) % server_num; shard_dims[server_index] = shard_dims[server_index] + 1; } + + // Check if the rank_id is valid if (shard_dims.count(rank_id) == 0) { MS_LOG(EXCEPTION) << "Invalid rank id " << rank_id << ", total server num " << server_num; } + + // Return the shard dimensions return shard_dims; } +// Define a function named "ReduceSparseGradient" belonging to the "Util" namespace void Util::ReduceSparseGradient(float *gradients, int *indices, const size_t indices_size, size_t segment_size, const size_t first_dim_size, const size_t outer_dim_size, mindspore::kernel::SparseGradient *unique_sparse_grad) { + + // Calculate the size of each slice segment size_t slice_segment_size = indices_size * segment_size; + + // Create a vector named "workspace_grad" with size equal to "slice_segment_size" to store gradients std::vector workspace_grad(slice_segment_size); + + // Create a vector named "workspace_indices" with size equal to "indices_size" to store indices std::vector workspace_indices(indices_size); - - MS_EXCEPTION_IF_NULL(gradients); - MS_EXCEPTION_IF_NULL(indices); - - mindspore::kernel::SparseGradient workspace_sparse_grad( - {workspace_grad.data(), workspace_indices.data(), indices_size}); - mindspore::kernel::SparseGradient input_sparse_grad({gradients, indices, indices_size}); - mindspore::kernel::ReduceSparseGradientParam param; - param.input_grad_ = &input_sparse_grad; - param.workspace_grad_ = &workspace_sparse_grad; - param.output_grad_ = unique_sparse_grad; - param.max_index_ = first_dim_size; - param.value_stride_ = outer_dim_size; - - mindspore::kernel::SparseOptimizerCpuKernelMod::BucketReduceSparseGradient(param); } +// Check if the "gradients" pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(gradients); + +// Check if the "indices" pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(indices); + +// Create a SparseGradient object named workspace_sparse_grad, which takes in three arguments: +// 1. A pointer to the data of the workspace_grad object +// 2. A pointer to the data of the workspace_indices object +// 3. The size of the indices (indices_size) +mindspore::kernel::SparseGradient workspace_sparse_grad( + {workspace_grad.data(), workspace_indices.data(), indices_size}); + +// Create a SparseGradient object named input_sparse_grad, which takes in three arguments: +// 1. A pointer to the data of the gradients object +// 2. A pointer to the data of the indices object +// 3. The size of the indices (indices_size) +mindspore::kernel::SparseGradient input_sparse_grad({gradients, indices, indices_size}); + +// Create a ReduceSparseGradientParam object named param +mindspore::kernel::ReduceSparseGradientParam param; + +// Set the input_grad_ member of the param object to the address of the input_sparse_grad object +param.input_grad_ = &input_sparse_grad; + +// Set the workspace_grad_ member of the param object to the address of the workspace_sparse_grad object +param.workspace_grad_ = &workspace_sparse_grad; + +// Set the output_grad_ member of the param object to the unique_sparse_grad object +param.output_grad_ = unique_sparse_grad; + +// Set the max_index_ member of the param object to the value of first_dim_size +param.max_index_ = first_dim_size; + +// Set the value_stride_ member of the param object to the value of outer_dim_size + +// Call the `BucketReduceSparseGradient` function of the `SparseOptimizerCpuKernelMod` class from the `mindspore::kernel` namespace, passing `param` as an argument. + +// Function to fuse server communication operations bool Util::FuseServerCommOps(const pipeline::ResourcePtr &res) { + + // Get the function graph from the resource FuncGraphPtr func_graph = res->func_graph(); + + // Throw an exception if the function graph is null MS_EXCEPTION_IF_NULL(func_graph); + + // Fuse the pull weight operation in the function graph DoFusion(func_graph, kPullWeightOpName, kFusedPullWeightOpName); + + // Fuse the push weight operation in the function graph DoFusion(func_graph, kPushWeightOpName, kFusedPushWeightOpName); + + // Return true to indicate successful fusion of server communication operations return true; } +// Define a function named MakeWeightPtr in the Util namespace that takes three arguments: +// 1. A const reference to a shared pointer to a vector of floats named data +// 2. A boolean variable named enable_recovery +// 3. A const reference to a shared pointer to a vector of integers named shape + WeightPtr Util::MakeWeightPtr(const std::shared_ptr> &data, bool enable_recovery, const std::shared_ptr> &shape) { + + // Declare a variable named weight_ptr of type WeightPtr (presumably a typedef for shared_ptr) WeightPtr weight_ptr; + + // Check if enable_recovery is false if (!enable_recovery) { + + // If enable_recovery is false, create a shared pointer to a Weight object using the data and shape arguments weight_ptr = std::make_shared(data, shape); } else { + + // If enable_recovery is true, create a shared pointer to a PersistentWeight object using the data and shape arguments weight_ptr = std::make_shared(data, shape); } + + // Return the weight_ptr variable return weight_ptr; } +// Define a function named "GetPrimitiveName" in the "Util" namespace that takes a reference to a CNodePtr object as input std::string Util::GetPrimitiveName(const CNodePtr &cnode) { + // Check if the input CNodePtr object is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(cnode); + + // Get the inputs of the CNode object auto &inputs = cnode->inputs(); + + // Check if the inputs vector is empty if (inputs.empty()) { + // If it is empty, log an exception with the message "Inputs of node is empty." MS_LOG(EXCEPTION) << "Inputs of node " << cnode->fullname_with_scope() << " is empty."; + + // Return an empty string return ""; } + + // Get the first input of the CNode object auto fn = inputs[0]; + + // Check if the first input is a ValueNode of type Primitive if (!IsValueNode(fn)) { + // If it is not, return an empty string return ""; } + // If it is, continue with the rest of the code - auto node_prim = GetValueNode(fn); - MS_EXCEPTION_IF_NULL(node_prim); - return node_prim->name(); -} +// Assign the value of the function argument 'fn' to the variable 'node_prim' using the 'GetValueNode' template function, +// which returns a pointer to a 'PrimitivePtr' object +auto node_prim = GetValueNode(fn); +// Check if 'node_prim' is a null pointer, and if so, throw an exception +MS_EXCEPTION_IF_NULL(node_prim); + +// Return the name of the 'PrimitivePtr' object pointed to by 'node_prim' +return node_prim->name(); + +// Define the function `DoFusion` in the `Util` namespace, which takes in a `FuncGraphPtr` named `func_graph`, a `std::string` named `cnode_name`, and a `std::string` named `fused_cnode_name` void Util::DoFusion(const FuncGraphPtr &func_graph, const std::string &cnode_name, const std::string &fused_cnode_name) { + + // Check if `func_graph` is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(func_graph); + + // Create a vector of `AnfNodePtr` named `node_list` and initialize it with the result of calling the `TopoSort` function on the `return` node of `func_graph` std::vector node_list = TopoSort(func_graph->get_return()); + // ... +} + // Declare a vector to store pointers to AnfNode objects std::vector single_nodes; + + // Declare a vector to store strings representing weight names std::vector weight_names; + + // Declare a vector to store integers representing indices std::vector indices; + + // Iterate over each AnfNodePtr object in the node_list for (const AnfNodePtr &node : node_list) { + + // Check if the current node is not null and is of type CNode if (node != nullptr && node->isa()) { + + // Check if the primitive name of the current CNode matches the specified cnode_name if (GetPrimitiveName(node->cast()) == cnode_name) { + + // If the conditions are met, add the current node to the single_nodes vector single_nodes.push_back(node); - auto weight_name_value_node = - common::AnfAlgo::GetInputNode(node->cast(), kNodeInputWeightNameOffset)->cast(); + // Get the input node of the current node at the specified offset using AnfAlgo::GetInputNode() function + auto weight_name_value_node = common::AnfAlgo::GetInputNode(node->cast(), kNodeInputWeightNameOffset)->cast(); + + // Get the value stored in the ValueNode as a string using GetValue() function const std::string &weight_name = GetValue(weight_name_value_node->value()); + + // Add the weight name to the weight_names vector weight_names.push_back(weight_name); - auto weight_index_value_node = - common::AnfAlgo::GetInputNode(node->cast(), kNodeInputWeightIndexOffset)->cast(); + // Get the input node of the current CNode using the AnfAlgo::GetInputNode function + auto weight_index_value_node = common::AnfAlgo::GetInputNode(node->cast(), kNodeInputWeightIndexOffset)->cast(); + + // Extract the value of the weight index from the value node int64_t weight_index = GetValue(weight_index_value_node->value()); + + // Add the weight index to the indices vector indices.push_back(weight_index); } } } + // Create a shared pointer to a Primitive object using the provided fused_cnode_name auto prim = std::make_shared(fused_cnode_name); + + // Check if the prim pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(prim); + + // Create an empty vector to store the inputs for the fused node std::vector fused_node_inputs = {}; + + // Add the prim object as a NewValueNode to the fused_node_inputs vector fused_node_inputs.push_back(NewValueNode(prim)); + + // Iterate over the single_nodes vector using a lambda function (void)std::for_each(single_nodes.begin(), single_nodes.end(), [&](const AnfNodePtr &node) { + + // Get the input node of the current node and add it to the fused_node_inputs vector fused_node_inputs.push_back(common::AnfAlgo::GetInputNode(node->cast(), 0)); }); - auto fused_cnode = func_graph->NewCNode(fused_node_inputs); - MS_EXCEPTION_IF_NULL(fused_cnode); - common::AnfAlgo::SetNodeAttr(kAttrPsKey, MakeValue(weight_names), fused_cnode); - common::AnfAlgo::SetNodeAttr(kAttrIndex, MakeValue(indices), fused_cnode); - common::AnfAlgo::SetNodeAttr(kAttrPrimitiveTarget, MakeValue(kCPUDevice), fused_cnode); +// Create a new CNode in the given function graph, using the provided fused_node_inputs as inputs +auto fused_cnode = func_graph->NewCNode(fused_node_inputs); - auto kernel_info = std::make_shared(); - MS_EXCEPTION_IF_NULL(kernel_info); - fused_cnode->set_kernel_info(kernel_info); - auto kernel_build_info = GenerateKernelBuildInfo(single_nodes); - AnfAlgo::SetSelectKernelBuildInfo(kernel_build_info, fused_cnode.get()); +// Check if the fused_cnode is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(fused_cnode); - AbstractBasePtrList abstract_list; - for (const auto &node : single_nodes) { +// Set the attribute with key kAttrPsKey in the fused_cnode to the value of weight_names +common::AnfAlgo::SetNodeAttr(kAttrPsKey, MakeValue(weight_names), fused_cnode); + +// Set the attribute with key kAttrIndex in the fused_cnode to the value of indices +common::AnfAlgo::SetNodeAttr(kAttrIndex, MakeValue(indices), fused_cnode); + +// Set the attribute with key kAttrPrimitiveTarget in the fused_cnode to the value of kCPUDevice +common::AnfAlgo::SetNodeAttr(kAttrPrimitiveTarget, MakeValue(kCPUDevice), fused_cnode); + +// Create a shared pointer to a new instance of the KernelInfo class and assign it to the variable "kernel_info" +auto kernel_info = std::make_shared(); + +// Check if the kernel_info pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(kernel_info); + +// Set the kernel_info of the fused_cnode to the newly created kernel_info +fused_cnode->set_kernel_info(kernel_info); + +// Generate the kernel build information using the single_nodes +auto kernel_build_info = GenerateKernelBuildInfo(single_nodes); + +// Set the kernel build information for the fused_cnode using the AnfAlgo::SetSelectKernelBuildInfo function +AnfAlgo::SetSelectKernelBuildInfo(kernel_build_info, fused_cnode.get()); + +// Create an empty list of AbstractBasePtr objects called abstract_list + +AbstractBasePtrList abstract_list; + +// Iterate over each element in the single_nodes container using a range-based for loop +for (const auto &node : single_nodes) { + + // Cast the current node to a CNodePtr object and assign it to the variable cnode auto cnode = node->cast(); - MS_EXCEPTION_IF_NULL(cnode); - abstract_list.push_back(cnode->abstract()); - } - auto abstract_tuple = std::make_shared(abstract_list); - MS_EXCEPTION_IF_NULL(abstract_tuple); - fused_cnode->set_abstract(abstract_tuple); + // Throw an exception if cnode is a null pointer + MS_EXCEPTION_IF_NULL(cnode); + + // Add the abstract value of cnode to the abstract_list + abstract_list.push_back(cnode->abstract()); +} + +// Create a shared pointer to an AbstractTuple object called abstract_tuple, passing in the abstract_list as an argument +auto abstract_tuple = std::make_shared(abstract_list); + +// Throw an exception if abstract_tuple is a null pointer +MS_EXCEPTION_IF_NULL(abstract_tuple); + +// Set the abstract value of fused_cnode to the abstract_tuple +fused_cnode->set_abstract(abstract_tuple); + + // Get the manager of the given func_graph auto manager = func_graph->manager(); + + // Check if the manager is null, throw an exception if it is MS_EXCEPTION_IF_NULL(manager); + + // Iterate over each node in the single_nodes container for (const auto &node : single_nodes) { + + // Replace the current node with the fused_cnode using the manager's Replace function + // If the replacement fails, throw an exception if (!manager->Replace(node, fused_cnode)) { MS_LOG(EXCEPTION) << "manager replace node failed"; } } - return; -} + // Return from the function + return; + +// Function to generate kernel build information based on a list of nodes kernel::KernelBuildInfoPtr Util::GenerateKernelBuildInfo(const std::vector &node_list) { + + // Initialize vectors to store device format, device type, and output shape information std::vector inputs_device_format; std::vector outputs_device_format; std::vector inputs_device_type; std::vector outputs_device_type; std::vector> outputs_shape; + + // Create a builder object to build the kernel build information kernel::KernelBuildInfo::KernelBuildInfoBuilder builder; + + // Iterate over the node list for (size_t idx = 0; idx < node_list.size(); ++idx) { auto cnode = utils::cast(node_list[idx]); MS_EXCEPTION_IF_NULL(cnode); + + // Get the number of input tensors for the current node size_t input_num = common::AnfAlgo::GetInputTensorNum(cnode); + + // Iterate over the input tensors for (size_t input_index = 0; input_index < input_num; ++input_index) { + // Add the default device format for the input tensor (void)inputs_device_format.emplace_back(kOpFormat_DEFAULT); + + // Get the device type for the input tensor and add it to the vector inputs_device_type.push_back(common::AnfAlgo::GetPrevNodeOutputInferDataType(cnode, input_index)); } + + // Get the number of output tensors for the current node size_t output_num = common::AnfAlgo::GetOutputTensorNum(cnode); + + // Iterate over the output tensors for (size_t output_index = 0; output_index < output_num; ++output_index) { + // Add the default device format for the output tensor (void)outputs_device_format.emplace_back(kOpFormat_DEFAULT); + + // Get the device type for the output tensor and add it to the vector outputs_device_type.push_back(common::AnfAlgo::GetOutputInferDataType(cnode, output_index)); + + // Get the output shape for the output tensor and add it to the vector outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(cnode, output_index)); } } + + // Set the inputs format, outputs format, inputs device type, and outputs device type using the builder builder.SetInputsFormat(inputs_device_format); builder.SetOutputsFormat(outputs_device_format); builder.SetInputsDeviceType(inputs_device_type); builder.SetOutputsDeviceType(outputs_device_type); + + // Build and return the kernel build information return builder.Build(); } -} // namespace ps -} // namespace mindspore + +// End of namespace ps +} // namespace mindspore \ No newline at end of file -- 2.34.1