| 20 | namespace embedding { |
| 21 | |
| 22 | std::string CreateKeyValueStore(const std::string& key_value_store_options, int64_t local_rank_id, |
| 23 | int64_t rank_id, int64_t world_size) { |
| 24 | oneflow::embedding::KeyValueStoreOptions options(key_value_store_options); |
| 25 | #ifdef WITH_CUDA |
| 26 | oneflow::Singleton<oneflow::embedding::EmbeddingManager>::Get()->CreateKeyValueStore( |
| 27 | options, local_rank_id, rank_id, world_size); |
| 28 | return options.Name(); |
| 29 | #else |
| 30 | UNIMPLEMENTED() << "OneEmbedding Only Support with CUDA"; |
| 31 | #endif |
| 32 | return ""; |
| 33 | } |
| 34 | |
| 35 | void LoadSnapshot(const std::string& snapshot_name, const std::string& embedding_name, |
| 36 | int64_t local_rank_id, int64_t rank_id) { |
no test coverage detected