optimizer/optimizer.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "optimizer.h" | ||
| 2 | #include "ps/local_shm/local_shm_stage_report.h" | ||
| 3 | #include "storage/index/dram/extendible_hash.h" | ||
| 4 | #include <algorithm> | ||
| 5 | #include <cstring> | ||
| 6 | #include <limits> | ||
| 7 | |||
| 8 | namespace { | ||
| 9 | |||
| 10 | ✗ | std::vector<uint64_t> CollectReaderKeys(const ParameterCompressReader* reader) { | |
| 11 | ✗ | const int size = reader->item_size(); | |
| 12 | ✗ | std::vector<uint64_t> keys; | |
| 13 | ✗ | keys.reserve(size); | |
| 14 | ✗ | for (int i = 0; i < size; ++i) { | |
| 15 | ✗ | keys.push_back(reader->item(i)->key); | |
| 16 | } | ||
| 17 | ✗ | return keys; | |
| 18 | ✗ | } | |
| 19 | |||
| 20 | 16 | void ValidateFlatUpdateArgs(const base::ConstArray<uint64_t>& keys, | |
| 21 | const float* grads, | ||
| 22 | int64_t num_rows, | ||
| 23 | int64_t embedding_dim) { | ||
| 24 | 16 | if (grads == nullptr) { | |
| 25 | ✗ | throw std::runtime_error("UpdateFlat grads pointer is null"); | |
| 26 | } | ||
| 27 | 16 | if (num_rows < 0 || embedding_dim <= 0) { | |
| 28 | ✗ | throw std::runtime_error("UpdateFlat invalid rows/dim"); | |
| 29 | } | ||
| 30 | 16 | if (keys.Size() != static_cast<size_t>(num_rows)) { | |
| 31 | ✗ | throw std::runtime_error("UpdateFlat keys size mismatch"); | |
| 32 | } | ||
| 33 | 16 | } | |
| 34 | |||
| 35 | // Splits one direct sparse update into its stages so the perf report can tell | ||
| 36 | // index probing, value-store reads, the in-place write-back, and the | ||
| 37 | // missing-row inserts apart. The stage buckets are disjoint. | ||
| 38 | ✗ | void ReportSgdUpdateFlatBreakdown(const BaseKV::UpdateFlatStats& stats, | |
| 39 | const ExtendibleHashInsertCounters& hash) { | ||
| 40 | ✗ | constexpr double kNsPerUs = 1000.0; | |
| 41 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 42 | ✗ | "sgd_update_tag_us", stats.tag_ns / kNsPerUs); | |
| 43 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 44 | ✗ | "sgd_update_batch_get_us", stats.batch_get_ns / kNsPerUs); | |
| 45 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 46 | ✗ | "sgd_update_index_us", stats.index_lookup_ns / kNsPerUs); | |
| 47 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 48 | ✗ | "sgd_update_store_read_us", stats.store_read_ns / kNsPerUs); | |
| 49 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 50 | ✗ | "sgd_update_writeback_us", stats.writeback_ns / kNsPerUs); | |
| 51 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 52 | ✗ | "sgd_update_missing_write_us", stats.missing_write_ns / kNsPerUs); | |
| 53 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 54 | ✗ | "sgd_update_miss_alloc_us", stats.missing_alloc_ns / kNsPerUs); | |
| 55 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 56 | ✗ | "sgd_update_miss_index_us", stats.missing_index_ns / kNsPerUs); | |
| 57 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 58 | ✗ | "sgd_update_miss_track_us", stats.missing_track_ns / kNsPerUs); | |
| 59 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 60 | ✗ | "sgd_update_missing_rows", static_cast<double>(stats.missing_rows)); | |
| 61 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 62 | ✗ | "sgd_update_index_inserts", static_cast<double>(hash.insert_calls)); | |
| 63 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 64 | ✗ | "sgd_update_index_splits", static_cast<double>(hash.split_calls)); | |
| 65 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 66 | ✗ | "sgd_update_index_split_us", hash.split_ns / kNsPerUs); | |
| 67 | ✗ | } | |
| 68 | |||
| 69 | 12 | void CheckUpdateKeyTag( | |
| 70 | uint64_t key, TAG_TYPE expected, const std::string& table) { | ||
| 71 | 12 | const TAG_TYPE got = ExtractKeyTag(key); | |
| 72 | 12 | if (got != expected) { | |
| 73 | ✗ | throw std::runtime_error( | |
| 74 | ✗ | "update key tag mismatch for table " + table + ": got " + | |
| 75 | ✗ | std::to_string(static_cast<unsigned>(got)) + " expected " + | |
| 76 | ✗ | std::to_string(static_cast<unsigned>(expected))); | |
| 77 | } | ||
| 78 | 12 | } | |
| 79 | |||
| 80 | 32 | int InitOrReuseTensor( | |
| 81 | std::unordered_map<std::string, SparseTensor*>* tensor_map, | ||
| 82 | const std::string& name, | ||
| 83 | TensorType type, | ||
| 84 | TAG_TYPE tag, | ||
| 85 | std::vector<uint64_t> shape, | ||
| 86 | BaseKV* base_kv) { | ||
| 87 | 32 | auto it = tensor_map->find(name); | |
| 88 | 32 | if (it != tensor_map->end()) { | |
| 89 | ✗ | if (it->second->Tag() != tag || | |
| 90 | ✗ | it->second->EmbeddingDim() != static_cast<int64_t>(shape[1])) { | |
| 91 | ✗ | throw std::runtime_error( | |
| 92 | ✗ | "table already registered with a different tag/dim: " + name); | |
| 93 | } | ||
| 94 | ✗ | return static_cast<int>(tag); | |
| 95 | } | ||
| 96 | 32 | auto* tensor = new SparseTensor(); | |
| 97 | 32 | std::string mutable_name = name; | |
| 98 | 32 | tensor->init(mutable_name, type, tag, shape, base_kv); | |
| 99 | 32 | (*tensor_map)[name] = tensor; | |
| 100 | 32 | return static_cast<int>(tag); | |
| 101 | 32 | } | |
| 102 | |||
| 103 | } // namespace | ||
| 104 | |||
| 105 | 40 | std::unique_ptr<Optimizer> CreateOptimizer(const json& config) { | |
| 106 | 40 | if (!config.is_object()) { | |
| 107 | ✗ | throw std::invalid_argument("cache_ps.optimizer must be an object"); | |
| 108 | } | ||
| 109 | 40 | const std::string type = config.value("type", "SGD"); | |
| 110 | 40 | const float learning_rate = config.value("learning_rate", 0.01f); | |
| 111 | 40 | if (!std::isfinite(learning_rate) || learning_rate < 0.0f) { | |
| 112 | ✗ | throw std::invalid_argument("cache_ps.optimizer.learning_rate is invalid"); | |
| 113 | } | ||
| 114 | 40 | if (type == "SGD") | |
| 115 | 34 | return std::make_unique<SGD>(learning_rate); | |
| 116 | 6 | if (type == "RowWiseAdagrad") { | |
| 117 | 2 | const float epsilon = config.value("epsilon", 1e-10f); | |
| 118 | 2 | if (!std::isfinite(epsilon) || epsilon < 0.0f) { | |
| 119 | ✗ | throw std::invalid_argument("cache_ps.optimizer.epsilon is invalid"); | |
| 120 | } | ||
| 121 | 2 | return std::make_unique<RowWiseAdaGrad>(learning_rate, epsilon); | |
| 122 | } | ||
| 123 | 4 | if (type == "AdamW") { | |
| 124 | 2 | const float beta1 = config.value("beta1", 0.9f); | |
| 125 | 2 | const float beta2 = config.value("beta2", 0.98f); | |
| 126 | 2 | const float epsilon = config.value("epsilon", 1e-8f); | |
| 127 | 2 | const float weight_decay = config.value("weight_decay", 0.0f); | |
| 128 | 4 | if (!std::isfinite(beta1) || beta1 < 0.0f || beta1 >= 1.0f || | |
| 129 | 2 | !std::isfinite(beta2) || beta2 < 0.0f || beta2 >= 1.0f || | |
| 130 | 2 | !std::isfinite(epsilon) || epsilon < 0.0f || | |
| 131 | 4 | !std::isfinite(weight_decay) || weight_decay < 0.0f) { | |
| 132 | ✗ | throw std::invalid_argument( | |
| 133 | ✗ | "cache_ps.optimizer AdamW parameters are invalid"); | |
| 134 | } | ||
| 135 | 4 | return std::make_unique<AdamW>( | |
| 136 | 2 | learning_rate, beta1, beta2, epsilon, weight_decay); | |
| 137 | } | ||
| 138 | 2 | throw std::invalid_argument("Unsupported cache_ps.optimizer.type: " + type); | |
| 139 | 40 | } | |
| 140 | |||
| 141 | 28 | int SGD::Init(const std::vector<std::string> table_name, | |
| 142 | const EmbeddingTableConfig& config, | ||
| 143 | BaseKV* base_kv) { | ||
| 144 | 28 | LOG(INFO) << "SGD::Init called with " << table_name.size() << " table(s)"; | |
| 145 | 28 | int param_tag = -1; | |
| 146 | 28 | const int k = TensorsPerTable(); | |
| 147 | 56 | for (const auto& name : table_name) { | |
| 148 | 56 | LOG(INFO) << " Initializing table: '" << name << "' with shape [" | |
| 149 | 28 | << config.num_embeddings << ", " << config.embedding_dim | |
| 150 | 28 | << "] table_id=" << config.table_id; | |
| 151 | 28 | std::vector<uint64_t> shape = {config.num_embeddings, config.embedding_dim}; | |
| 152 | 56 | param_tag = InitOrReuseTensor( | |
| 153 | &tensor_map_, | ||
| 154 | name, | ||
| 155 | PARAMETER, | ||
| 156 | 28 | MakeTensorTag(config.table_id, 0, k), | |
| 157 | shape, | ||
| 158 | base_kv); | ||
| 159 | 28 | } | |
| 160 | 56 | LOG(INFO) << "SGD::Init completed. tensor_map_ now has " << tensor_map_.size() | |
| 161 | 28 | << " entries"; | |
| 162 | 28 | return param_tag; | |
| 163 | } | ||
| 164 | |||
| 165 | ✗ | void SGD::Update( | |
| 166 | std::string table, const ParameterCompressReader* reader, unsigned tid) { | ||
| 167 | ✗ | auto it = tensor_map_.find(table); | |
| 168 | ✗ | if (it == tensor_map_.end()) { | |
| 169 | ✗ | LOG(ERROR) << "Table not found in SGD optimizer: '" << table << "'"; | |
| 170 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 171 | } | ||
| 172 | |||
| 173 | ✗ | int size = reader->item_size(); | |
| 174 | ✗ | std::vector<uint64_t> keys = CollectReaderKeys(reader); | |
| 175 | ✗ | if (!keys.empty()) { | |
| 176 | ✗ | CheckUpdateKeyTag(keys[0], it->second->Tag(), table); | |
| 177 | } | ||
| 178 | |||
| 179 | ✗ | std::vector<base::ConstArray<float>> current_values; | |
| 180 | ✗ | it->second->BatchGet(keys, ¤t_values, tid); | |
| 181 | |||
| 182 | ✗ | for (int i = 0; i < size; ++i) { | |
| 183 | ✗ | const auto* item = reader->item(i); | |
| 184 | ✗ | if (current_values[i].Size() == 0) { | |
| 185 | // If key not found, we fallback to Put to initialize it | ||
| 186 | ✗ | std::vector<float> zero_init(item->dim, 0.0f); | |
| 187 | ✗ | for (int j = 0; j < item->dim; ++j) { | |
| 188 | ✗ | zero_init[j] = -learning_rate_ * item->data()[j]; | |
| 189 | } | ||
| 190 | std::string val_str( | ||
| 191 | ✗ | (char*)zero_init.data(), zero_init.size() * sizeof(float)); | |
| 192 | ✗ | it->second->Put(item->key, val_str, tid); | |
| 193 | ✗ | continue; | |
| 194 | ✗ | } | |
| 195 | |||
| 196 | ✗ | float* data = const_cast<float*>(current_values[i].Data()); | |
| 197 | ✗ | int dim = std::min(current_values[i].Size(), item->dim); | |
| 198 | |||
| 199 | ✗ | #pragma omp simd | |
| 200 | for (int j = 0; j < dim; ++j) { | ||
| 201 | ✗ | data[j] -= learning_rate_ * item->data()[j]; | |
| 202 | } | ||
| 203 | } | ||
| 204 | ✗ | } | |
| 205 | |||
| 206 | 8 | void SGD::UpdateFlat( | |
| 207 | std::string table, | ||
| 208 | const base::ConstArray<uint64_t>& keys, | ||
| 209 | const float* grads, | ||
| 210 | int64_t num_rows, | ||
| 211 | int64_t embedding_dim, | ||
| 212 | unsigned tid) { | ||
| 213 | 8 | ValidateFlatUpdateArgs(keys, grads, num_rows, embedding_dim); | |
| 214 | |||
| 215 | 8 | auto it = tensor_map_.find(table); | |
| 216 | 8 | if (it == tensor_map_.end()) { | |
| 217 | ✗ | LOG(ERROR) << "Table not found in SGD optimizer: '" << table << "'"; | |
| 218 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 219 | } | ||
| 220 | 8 | if (it->second->EmbeddingDim() != embedding_dim) { | |
| 221 | ✗ | throw std::runtime_error( | |
| 222 | ✗ | "SGD::UpdateFlat embedding_dim mismatch for table " + table); | |
| 223 | } | ||
| 224 | 8 | if (keys.Size() > 0) { | |
| 225 | 8 | CheckUpdateKeyTag(keys[0], it->second->Tag(), table); | |
| 226 | } | ||
| 227 | |||
| 228 | // Only ask for the sub-stage breakdown when the stage report is actually on; | ||
| 229 | // otherwise the per-update and per-insert clock reads are pure overhead. | ||
| 230 | 8 | const bool collect_stats = recstore::IsLocalShmStageReportEnabled(); | |
| 231 | 8 | const auto direct_update_start = std::chrono::steady_clock::now(); | |
| 232 | 8 | BaseKV::UpdateFlatStats update_stats; | |
| 233 | 8 | ExtendibleHashInsertCounters hash_before; | |
| 234 | 8 | if (collect_stats) { | |
| 235 | ✗ | GetExtendibleHashInsertCounters(&hash_before); | |
| 236 | } | ||
| 237 | 8 | if (it->second->ApplySgdUpdateFlat( | |
| 238 | keys, | ||
| 239 | grads, | ||
| 240 | num_rows, | ||
| 241 | embedding_dim, | ||
| 242 | learning_rate_, | ||
| 243 | tid, | ||
| 244 | collect_stats ? &update_stats : nullptr)) { | ||
| 245 | 8 | recstore::ReportLocalShmStageMetric( | |
| 246 | "sgd_update_direct_us", | ||
| 247 | recstore::LocalShmElapsedUs(direct_update_start)); | ||
| 248 | 8 | if (collect_stats) { | |
| 249 | ✗ | ExtendibleHashInsertCounters hash_after; | |
| 250 | ✗ | GetExtendibleHashInsertCounters(&hash_after); | |
| 251 | const ExtendibleHashInsertCounters hash_delta{ | ||
| 252 | ✗ | hash_after.insert_calls - hash_before.insert_calls, | |
| 253 | ✗ | hash_after.split_calls - hash_before.split_calls, | |
| 254 | ✗ | hash_after.split_ns - hash_before.split_ns}; | |
| 255 | ✗ | ReportSgdUpdateFlatBreakdown(update_stats, hash_delta); | |
| 256 | } | ||
| 257 | 8 | return; | |
| 258 | } | ||
| 259 | |||
| 260 | ✗ | std::vector<uint64_t> key_vec(keys.Data(), keys.Data() + keys.Size()); | |
| 261 | ✗ | const auto batch_get_start = std::chrono::steady_clock::now(); | |
| 262 | ✗ | std::vector<base::ConstArray<float>> current_values; | |
| 263 | ✗ | it->second->BatchGet(key_vec, ¤t_values, tid); | |
| 264 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 265 | "sgd_update_batch_get_us", recstore::LocalShmElapsedUs(batch_get_start)); | ||
| 266 | |||
| 267 | ✗ | const auto apply_start = std::chrono::steady_clock::now(); | |
| 268 | ✗ | int64_t missing_rows = 0; | |
| 269 | ✗ | for (int64_t row = 0; row < num_rows; ++row) { | |
| 270 | ✗ | const float* row_grad = grads + row * embedding_dim; | |
| 271 | ✗ | const auto& current = current_values[static_cast<size_t>(row)]; | |
| 272 | ✗ | if (current.Size() == 0) { | |
| 273 | ✗ | ++missing_rows; | |
| 274 | ✗ | std::vector<float> zero_init(static_cast<size_t>(embedding_dim), 0.0f); | |
| 275 | ✗ | for (int64_t col = 0; col < embedding_dim; ++col) { | |
| 276 | ✗ | zero_init[static_cast<size_t>(col)] = -learning_rate_ * row_grad[col]; | |
| 277 | } | ||
| 278 | ✗ | std::string val_str(reinterpret_cast<char*>(zero_init.data()), | |
| 279 | ✗ | zero_init.size() * sizeof(float)); | |
| 280 | ✗ | it->second->Put(keys[static_cast<size_t>(row)], val_str, tid); | |
| 281 | ✗ | continue; | |
| 282 | ✗ | } | |
| 283 | ✗ | if (static_cast<int64_t>(current.Size()) != embedding_dim) { | |
| 284 | ✗ | throw std::runtime_error( | |
| 285 | ✗ | "SGD::UpdateFlat embedding_dim mismatch for table " + table); | |
| 286 | } | ||
| 287 | |||
| 288 | ✗ | float* data = const_cast<float*>(current.Data()); | |
| 289 | ✗ | #pragma omp simd | |
| 290 | for (int64_t col = 0; col < embedding_dim; ++col) { | ||
| 291 | ✗ | data[col] -= learning_rate_ * row_grad[col]; | |
| 292 | } | ||
| 293 | } | ||
| 294 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 295 | "sgd_update_apply_us", recstore::LocalShmElapsedUs(apply_start)); | ||
| 296 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 297 | "sgd_update_missing_rows", static_cast<double>(missing_rows)); | ||
| 298 | ✗ | } | |
| 299 | |||
| 300 | ✗ | int AdaGrad::Init(const std::vector<std::string> table_name, | |
| 301 | const EmbeddingTableConfig& config, | ||
| 302 | BaseKV* base_kv) { | ||
| 303 | ✗ | int param_tag = -1; | |
| 304 | ✗ | const int k = TensorsPerTable(); | |
| 305 | ✗ | for (const auto& name : table_name) { | |
| 306 | ✗ | std::vector<uint64_t> shape = {config.num_embeddings, config.embedding_dim}; | |
| 307 | ✗ | param_tag = InitOrReuseTensor( | |
| 308 | &tensor_map_, | ||
| 309 | name, | ||
| 310 | PARAMETER, | ||
| 311 | ✗ | MakeTensorTag(config.table_id, 0, k), | |
| 312 | shape, | ||
| 313 | base_kv); | ||
| 314 | ✗ | InitOrReuseTensor( | |
| 315 | &tensor_map_, | ||
| 316 | ✗ | name + "_accumulated_grad", | |
| 317 | MOMENT_1, | ||
| 318 | ✗ | MakeTensorTag(config.table_id, 1, k), | |
| 319 | shape, | ||
| 320 | base_kv); | ||
| 321 | ✗ | } | |
| 322 | ✗ | return param_tag; | |
| 323 | } | ||
| 324 | |||
| 325 | ✗ | void AdaGrad::Update( | |
| 326 | std::string table, const ParameterCompressReader* reader, unsigned tid) { | ||
| 327 | ✗ | auto param_it = tensor_map_.find(table); | |
| 328 | ✗ | if (param_it == tensor_map_.end()) { | |
| 329 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 330 | } | ||
| 331 | |||
| 332 | ✗ | std::string acc_table = table + "_accumulated_grad"; | |
| 333 | ✗ | auto acc_it = tensor_map_.find(acc_table); | |
| 334 | ✗ | if (acc_it == tensor_map_.end()) { | |
| 335 | ✗ | throw std::runtime_error( | |
| 336 | ✗ | "Accumulated gradient table not found: " + acc_table); | |
| 337 | } | ||
| 338 | |||
| 339 | ✗ | int size = reader->item_size(); | |
| 340 | ✗ | std::vector<uint64_t> keys = CollectReaderKeys(reader); | |
| 341 | ✗ | if (!keys.empty()) { | |
| 342 | ✗ | CheckUpdateKeyTag(keys[0], param_it->second->Tag(), table); | |
| 343 | } | ||
| 344 | |||
| 345 | ✗ | std::vector<base::ConstArray<float>> current_values; | |
| 346 | ✗ | std::vector<base::ConstArray<float>> acc_values; | |
| 347 | ✗ | param_it->second->BatchGet(keys, ¤t_values, tid); | |
| 348 | ✗ | acc_it->second->BatchGet(keys, &acc_values, tid); | |
| 349 | |||
| 350 | ✗ | for (int i = 0; i < size; ++i) { | |
| 351 | ✗ | const auto* item = reader->item(i); | |
| 352 | ✗ | if (current_values[i].Size() == 0 || acc_values[i].Size() == 0) { | |
| 353 | // Fallback to sequential initialization if not found | ||
| 354 | // (This is rare in training but kept for robustness) | ||
| 355 | ✗ | continue; | |
| 356 | } | ||
| 357 | |||
| 358 | ✗ | float* param_data = const_cast<float*>(current_values[i].Data()); | |
| 359 | ✗ | float* acc_data = const_cast<float*>(acc_values[i].Data()); | |
| 360 | ✗ | int dim = std::min(current_values[i].Size(), item->dim); | |
| 361 | |||
| 362 | ✗ | #pragma omp simd | |
| 363 | for (int j = 0; j < dim; ++j) { | ||
| 364 | ✗ | acc_data[j] += item->data()[j] * item->data()[j]; | |
| 365 | ✗ | float adaptive_lr = learning_rate_ / (std::sqrt(acc_data[j]) + epsilon_); | |
| 366 | ✗ | param_data[j] -= adaptive_lr * item->data()[j]; | |
| 367 | } | ||
| 368 | } | ||
| 369 | ✗ | } | |
| 370 | |||
| 371 | ✗ | void AdaGrad::UpdateFlat( | |
| 372 | std::string table, | ||
| 373 | const base::ConstArray<uint64_t>& keys, | ||
| 374 | const float* grads, | ||
| 375 | int64_t num_rows, | ||
| 376 | int64_t embedding_dim, | ||
| 377 | unsigned tid) { | ||
| 378 | ✗ | ValidateFlatUpdateArgs(keys, grads, num_rows, embedding_dim); | |
| 379 | |||
| 380 | ✗ | auto param_it = tensor_map_.find(table); | |
| 381 | ✗ | if (param_it == tensor_map_.end()) { | |
| 382 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 383 | } | ||
| 384 | |||
| 385 | ✗ | std::string acc_table = table + "_accumulated_grad"; | |
| 386 | ✗ | auto acc_it = tensor_map_.find(acc_table); | |
| 387 | ✗ | if (acc_it == tensor_map_.end()) { | |
| 388 | ✗ | throw std::runtime_error( | |
| 389 | ✗ | "Accumulated gradient table not found: " + acc_table); | |
| 390 | } | ||
| 391 | |||
| 392 | ✗ | if (keys.Size() > 0) { | |
| 393 | ✗ | CheckUpdateKeyTag(keys[0], param_it->second->Tag(), table); | |
| 394 | } | ||
| 395 | ✗ | std::vector<uint64_t> key_vec(keys.Data(), keys.Data() + keys.Size()); | |
| 396 | ✗ | std::vector<base::ConstArray<float>> current_values; | |
| 397 | ✗ | std::vector<base::ConstArray<float>> acc_values; | |
| 398 | ✗ | param_it->second->BatchGet(key_vec, ¤t_values, tid); | |
| 399 | ✗ | acc_it->second->BatchGet(key_vec, &acc_values, tid); | |
| 400 | |||
| 401 | ✗ | for (int64_t row = 0; row < num_rows; ++row) { | |
| 402 | ✗ | const auto& current = current_values[static_cast<size_t>(row)]; | |
| 403 | ✗ | const auto& acc = acc_values[static_cast<size_t>(row)]; | |
| 404 | ✗ | if (current.Size() == 0 || acc.Size() == 0) { | |
| 405 | ✗ | continue; | |
| 406 | } | ||
| 407 | ✗ | if (static_cast<int64_t>(current.Size()) != embedding_dim || | |
| 408 | ✗ | static_cast<int64_t>(acc.Size()) != embedding_dim) { | |
| 409 | ✗ | throw std::runtime_error( | |
| 410 | ✗ | "AdaGrad::UpdateFlat embedding_dim mismatch for table " + table); | |
| 411 | } | ||
| 412 | |||
| 413 | ✗ | const float* row_grad = grads + row * embedding_dim; | |
| 414 | ✗ | float* param_data = const_cast<float*>(current.Data()); | |
| 415 | ✗ | float* acc_data = const_cast<float*>(acc.Data()); | |
| 416 | ✗ | #pragma omp simd | |
| 417 | for (int64_t col = 0; col < embedding_dim; ++col) { | ||
| 418 | ✗ | acc_data[col] += row_grad[col] * row_grad[col]; | |
| 419 | float adaptive_lr = | ||
| 420 | ✗ | learning_rate_ / (std::sqrt(acc_data[col]) + epsilon_); | |
| 421 | ✗ | param_data[col] -= adaptive_lr * row_grad[col]; | |
| 422 | } | ||
| 423 | } | ||
| 424 | ✗ | } | |
| 425 | |||
| 426 | 2 | int RowWiseAdaGrad::Init(const std::vector<std::string> table_name, | |
| 427 | const EmbeddingTableConfig& config, | ||
| 428 | BaseKV* base_kv) { | ||
| 429 | 2 | int param_tag = -1; | |
| 430 | 2 | const int k = TensorsPerTable(); | |
| 431 | 4 | for (const auto& name : table_name) { | |
| 432 | 2 | std::vector<uint64_t> shape = {config.num_embeddings, config.embedding_dim}; | |
| 433 | 2 | std::vector<uint64_t> acc_shape = {config.num_embeddings, 1}; | |
| 434 | 4 | param_tag = InitOrReuseTensor( | |
| 435 | &tensor_map_, | ||
| 436 | name, | ||
| 437 | PARAMETER, | ||
| 438 | 2 | MakeTensorTag(config.table_id, 0, k), | |
| 439 | shape, | ||
| 440 | base_kv); | ||
| 441 | 4 | InitOrReuseTensor( | |
| 442 | &tensor_map_, | ||
| 443 | 2 | name + "_rowwise_accumulated_grad", | |
| 444 | MOMENT_1, | ||
| 445 | 2 | MakeTensorTag(config.table_id, 1, k), | |
| 446 | acc_shape, | ||
| 447 | base_kv); | ||
| 448 | 2 | } | |
| 449 | 2 | return param_tag; | |
| 450 | } | ||
| 451 | |||
| 452 | ✗ | void RowWiseAdaGrad::Update( | |
| 453 | std::string table, const ParameterCompressReader* reader, unsigned tid) { | ||
| 454 | ✗ | auto param_it = tensor_map_.find(table); | |
| 455 | ✗ | if (param_it == tensor_map_.end()) { | |
| 456 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 457 | } | ||
| 458 | |||
| 459 | ✗ | std::string acc_table = table + "_rowwise_accumulated_grad"; | |
| 460 | ✗ | auto acc_it = tensor_map_.find(acc_table); | |
| 461 | ✗ | if (acc_it == tensor_map_.end()) { | |
| 462 | ✗ | throw std::runtime_error( | |
| 463 | ✗ | "Row-wise accumulated gradient table not found: " + acc_table); | |
| 464 | } | ||
| 465 | |||
| 466 | ✗ | int size = reader->item_size(); | |
| 467 | ✗ | std::vector<uint64_t> keys = CollectReaderKeys(reader); | |
| 468 | ✗ | if (!keys.empty()) { | |
| 469 | ✗ | CheckUpdateKeyTag(keys[0], param_it->second->Tag(), table); | |
| 470 | } | ||
| 471 | |||
| 472 | ✗ | std::vector<base::ConstArray<float>> current_values; | |
| 473 | ✗ | std::vector<base::ConstArray<float>> acc_values; | |
| 474 | ✗ | param_it->second->BatchGet(keys, ¤t_values, tid); | |
| 475 | ✗ | acc_it->second->BatchGet(keys, &acc_values, tid); | |
| 476 | |||
| 477 | ✗ | for (int i = 0; i < size; ++i) { | |
| 478 | ✗ | const auto* item = reader->item(i); | |
| 479 | ✗ | const auto& current = current_values[static_cast<size_t>(i)]; | |
| 480 | ✗ | const auto& acc = acc_values[static_cast<size_t>(i)]; | |
| 481 | ✗ | const int64_t expected_dim = param_it->second->EmbeddingDim(); | |
| 482 | ✗ | if (item->dim != expected_dim || | |
| 483 | ✗ | (current.Size() != 0 && current.Size() != expected_dim) || | |
| 484 | ✗ | (acc.Size() != 0 && acc.Size() != 1)) { | |
| 485 | ✗ | throw std::runtime_error( | |
| 486 | ✗ | "RowWiseAdaGrad::Update embedding_dim mismatch for table " + table); | |
| 487 | } | ||
| 488 | ✗ | const int dim = item->dim; | |
| 489 | |||
| 490 | ✗ | float grad_square_mean = 0.0; | |
| 491 | ✗ | #pragma omp simd reduction(+ : grad_square_mean) | |
| 492 | for (int j = 0; j < dim; ++j) { | ||
| 493 | ✗ | grad_square_mean += item->data()[j] * item->data()[j]; | |
| 494 | } | ||
| 495 | ✗ | grad_square_mean /= dim; | |
| 496 | |||
| 497 | ✗ | float accumulated_grad = acc.Size() == 0 ? 0.0f : acc.Data()[0]; | |
| 498 | ✗ | accumulated_grad += grad_square_mean; | |
| 499 | |||
| 500 | const float adaptive_lr = | ||
| 501 | ✗ | learning_rate_ / (std::sqrt(accumulated_grad) + epsilon_); | |
| 502 | ✗ | if (current.Size() == 0) { | |
| 503 | ✗ | std::vector<float> initial_value(static_cast<size_t>(dim), 0.0f); | |
| 504 | ✗ | for (int j = 0; j < dim; ++j) { | |
| 505 | ✗ | initial_value[static_cast<size_t>(j)] = -adaptive_lr * item->data()[j]; | |
| 506 | } | ||
| 507 | const std::string value( | ||
| 508 | ✗ | reinterpret_cast<const char*>(initial_value.data()), | |
| 509 | ✗ | initial_value.size() * sizeof(float)); | |
| 510 | ✗ | param_it->second->Put(item->key, value, tid); | |
| 511 | ✗ | } else { | |
| 512 | ✗ | float* param_data = const_cast<float*>(current.Data()); | |
| 513 | ✗ | #pragma omp simd | |
| 514 | for (int j = 0; j < dim; ++j) { | ||
| 515 | ✗ | param_data[j] -= adaptive_lr * item->data()[j]; | |
| 516 | } | ||
| 517 | } | ||
| 518 | |||
| 519 | ✗ | if (acc.Size() == 0) { | |
| 520 | const std::string value( | ||
| 521 | ✗ | reinterpret_cast<const char*>(&accumulated_grad), sizeof(float)); | |
| 522 | ✗ | acc_it->second->Put(item->key, value, tid); | |
| 523 | ✗ | } else { | |
| 524 | ✗ | const_cast<float*>(acc.Data())[0] = accumulated_grad; | |
| 525 | } | ||
| 526 | } | ||
| 527 | ✗ | } | |
| 528 | |||
| 529 | 4 | void RowWiseAdaGrad::UpdateFlat( | |
| 530 | std::string table, | ||
| 531 | const base::ConstArray<uint64_t>& keys, | ||
| 532 | const float* grads, | ||
| 533 | int64_t num_rows, | ||
| 534 | int64_t embedding_dim, | ||
| 535 | unsigned tid) { | ||
| 536 | 4 | ValidateFlatUpdateArgs(keys, grads, num_rows, embedding_dim); | |
| 537 | |||
| 538 | 4 | auto param_it = tensor_map_.find(table); | |
| 539 | 4 | if (param_it == tensor_map_.end()) { | |
| 540 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 541 | } | ||
| 542 | |||
| 543 | 4 | std::string acc_table = table + "_rowwise_accumulated_grad"; | |
| 544 | 4 | auto acc_it = tensor_map_.find(acc_table); | |
| 545 | 4 | if (acc_it == tensor_map_.end()) { | |
| 546 | ✗ | throw std::runtime_error( | |
| 547 | ✗ | "Row-wise accumulated gradient table not found: " + acc_table); | |
| 548 | } | ||
| 549 | |||
| 550 | 4 | if (keys.Size() > 0) { | |
| 551 | 4 | CheckUpdateKeyTag(keys[0], param_it->second->Tag(), table); | |
| 552 | } | ||
| 553 | 4 | std::vector<uint64_t> key_vec(keys.Data(), keys.Data() + keys.Size()); | |
| 554 | 4 | std::vector<base::ConstArray<float>> current_values; | |
| 555 | 4 | std::vector<base::ConstArray<float>> acc_values; | |
| 556 | 4 | param_it->second->BatchGet(key_vec, ¤t_values, tid); | |
| 557 | 4 | acc_it->second->BatchGet(key_vec, &acc_values, tid); | |
| 558 | |||
| 559 | 8 | for (int64_t row = 0; row < num_rows; ++row) { | |
| 560 | 4 | const auto& current = current_values[static_cast<size_t>(row)]; | |
| 561 | 4 | const auto& acc = acc_values[static_cast<size_t>(row)]; | |
| 562 | 4 | if ((current.Size() != 0 && | |
| 563 | 10 | static_cast<int64_t>(current.Size()) != embedding_dim) || | |
| 564 | 6 | (acc.Size() != 0 && acc.Size() != 1)) { | |
| 565 | ✗ | throw std::runtime_error( | |
| 566 | ✗ | "RowWiseAdaGrad::UpdateFlat embedding_dim mismatch for table " + | |
| 567 | ✗ | table); | |
| 568 | } | ||
| 569 | |||
| 570 | 4 | const float* row_grad = grads + row * embedding_dim; | |
| 571 | 4 | float grad_square_mean = 0.0f; | |
| 572 | 4 | #pragma omp simd reduction(+ : grad_square_mean) | |
| 573 | for (int64_t col = 0; col < embedding_dim; ++col) { | ||
| 574 | 8 | grad_square_mean += row_grad[col] * row_grad[col]; | |
| 575 | } | ||
| 576 | 4 | grad_square_mean /= static_cast<float>(embedding_dim); | |
| 577 | |||
| 578 | 4 | float accumulated_grad = acc.Size() == 0 ? 0.0f : acc.Data()[0]; | |
| 579 | 4 | accumulated_grad += grad_square_mean; | |
| 580 | const float adaptive_lr = | ||
| 581 | 4 | learning_rate_ / (std::sqrt(accumulated_grad) + epsilon_); | |
| 582 | |||
| 583 | 4 | if (current.Size() == 0) { | |
| 584 | std::vector<float> initial_value( | ||
| 585 | 2 | static_cast<size_t>(embedding_dim), 0.0f); | |
| 586 | 6 | for (int64_t col = 0; col < embedding_dim; ++col) { | |
| 587 | 4 | initial_value[static_cast<size_t>(col)] = -adaptive_lr * row_grad[col]; | |
| 588 | } | ||
| 589 | const std::string value( | ||
| 590 | 2 | reinterpret_cast<const char*>(initial_value.data()), | |
| 591 | 2 | initial_value.size() * sizeof(float)); | |
| 592 | 2 | param_it->second->Put(keys[static_cast<size_t>(row)], value, tid); | |
| 593 | 2 | } else { | |
| 594 | 2 | float* param_data = const_cast<float*>(current.Data()); | |
| 595 | 2 | #pragma omp simd | |
| 596 | for (int64_t col = 0; col < embedding_dim; ++col) { | ||
| 597 | 4 | param_data[col] -= adaptive_lr * row_grad[col]; | |
| 598 | } | ||
| 599 | } | ||
| 600 | |||
| 601 | 4 | if (acc.Size() == 0) { | |
| 602 | const std::string value( | ||
| 603 | 2 | reinterpret_cast<const char*>(&accumulated_grad), sizeof(float)); | |
| 604 | 2 | acc_it->second->Put(keys[static_cast<size_t>(row)], value, tid); | |
| 605 | 2 | } else { | |
| 606 | 2 | const_cast<float*>(acc.Data())[0] = accumulated_grad; | |
| 607 | } | ||
| 608 | } | ||
| 609 | 4 | } | |
| 610 | |||
| 611 | namespace { | ||
| 612 | |||
| 613 | // The step is stored as a tagged scalar in its own table. Keep its key away | ||
| 614 | // from normal embedding ids (the top 8 bits are reserved for TensorType). | ||
| 615 | constexpr uint64_t kAdamWStepKey = (std::numeric_limits<uint64_t>::max() >> 8); | ||
| 616 | |||
| 617 | } // namespace | ||
| 618 | |||
| 619 | 4 | int AdamW::Init(const std::vector<std::string> table_name, | |
| 620 | const EmbeddingTableConfig& config, | ||
| 621 | BaseKV* base_kv) { | ||
| 622 | 8 | for (const auto& name : table_name) { | |
| 623 | const std::vector<uint64_t> shape = { | ||
| 624 | 4 | config.num_embeddings, config.embedding_dim}; | |
| 625 | 4 | auto* param_tensor = new SparseTensor(); | |
| 626 | 4 | auto mutable_name = name; | |
| 627 | 4 | auto mutable_shape = shape; | |
| 628 | 4 | param_tensor->init( | |
| 629 | mutable_name, PARAMETER, PARAMETER, mutable_shape, base_kv); | ||
| 630 | 4 | tensor_map_[name] = param_tensor; | |
| 631 | |||
| 632 | 4 | auto* first_moment = new SparseTensor(); | |
| 633 | 4 | const std::string first_name = name + "_adamw_m"; | |
| 634 | 4 | auto mutable_first_name = first_name; | |
| 635 | 4 | first_moment->init( | |
| 636 | mutable_first_name, MOMENT_1, MOMENT_1, mutable_shape, base_kv); | ||
| 637 | 4 | tensor_map_[first_name] = first_moment; | |
| 638 | |||
| 639 | 4 | auto* second_moment = new SparseTensor(); | |
| 640 | 4 | const std::string second_name = name + "_adamw_v"; | |
| 641 | 4 | auto mutable_second_name = second_name; | |
| 642 | 4 | second_moment->init( | |
| 643 | mutable_second_name, MOMENT_2, MOMENT_2, mutable_shape, base_kv); | ||
| 644 | 4 | tensor_map_[second_name] = second_moment; | |
| 645 | |||
| 646 | 4 | auto* step_tensor = new SparseTensor(); | |
| 647 | 4 | const std::string step_name = name + "_adamw_step"; | |
| 648 | 4 | auto mutable_step_name = step_name; | |
| 649 | 4 | std::vector<uint64_t> step_shape = {1, 1}; | |
| 650 | 4 | step_tensor->init( | |
| 651 | mutable_step_name, MOMENT_1, MOMENT_1, step_shape, base_kv); | ||
| 652 | 4 | tensor_map_[step_name] = step_tensor; | |
| 653 | 4 | } | |
| 654 | 4 | return static_cast<int>(PARAMETER); | |
| 655 | } | ||
| 656 | |||
| 657 | ✗ | void AdamW::Update( | |
| 658 | std::string table, const ParameterCompressReader* reader, unsigned tid) { | ||
| 659 | ✗ | auto param_it = tensor_map_.find(table); | |
| 660 | ✗ | if (param_it == tensor_map_.end()) { | |
| 661 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 662 | } | ||
| 663 | ✗ | const int size = reader->item_size(); | |
| 664 | ✗ | const int64_t dim = param_it->second->EmbeddingDim(); | |
| 665 | ✗ | std::vector<uint64_t> keys; | |
| 666 | ✗ | std::vector<float> grads; | |
| 667 | ✗ | keys.reserve(size); | |
| 668 | ✗ | grads.reserve(static_cast<size_t>(size) * static_cast<size_t>(dim)); | |
| 669 | ✗ | for (int i = 0; i < size; ++i) { | |
| 670 | ✗ | const auto* item = reader->item(i); | |
| 671 | ✗ | if (item->dim != dim) { | |
| 672 | ✗ | throw std::runtime_error( | |
| 673 | ✗ | "AdamW::Update embedding_dim mismatch for table " + table); | |
| 674 | } | ||
| 675 | ✗ | keys.push_back(item->key); | |
| 676 | ✗ | grads.insert(grads.end(), item->data(), item->data() + dim); | |
| 677 | } | ||
| 678 | ✗ | UpdateRows(table, keys.data(), grads.data(), size, dim, tid); | |
| 679 | ✗ | } | |
| 680 | |||
| 681 | 4 | void AdamW::UpdateFlat( | |
| 682 | std::string table, | ||
| 683 | const base::ConstArray<uint64_t>& keys, | ||
| 684 | const float* grads, | ||
| 685 | int64_t num_rows, | ||
| 686 | int64_t embedding_dim, | ||
| 687 | unsigned tid) { | ||
| 688 | 4 | ValidateFlatUpdateArgs(keys, grads, num_rows, embedding_dim); | |
| 689 | 4 | auto it = tensor_map_.find(table); | |
| 690 | 4 | if (it == tensor_map_.end()) { | |
| 691 | ✗ | throw std::runtime_error("Table not found: " + table); | |
| 692 | } | ||
| 693 | 4 | if (it->second->EmbeddingDim() != embedding_dim) { | |
| 694 | ✗ | throw std::runtime_error( | |
| 695 | ✗ | "AdamW::UpdateFlat embedding_dim mismatch for table " + table); | |
| 696 | } | ||
| 697 | 4 | UpdateRows(table, keys.Data(), grads, num_rows, embedding_dim, tid); | |
| 698 | 4 | } | |
| 699 | |||
| 700 | 4 | void AdamW::UpdateRows( | |
| 701 | const std::string& table, | ||
| 702 | const uint64_t* keys, | ||
| 703 | const float* grads, | ||
| 704 | int64_t num_rows, | ||
| 705 | int64_t embedding_dim, | ||
| 706 | unsigned tid) { | ||
| 707 | 4 | auto param_it = tensor_map_.find(table); | |
| 708 | 4 | auto first_it = tensor_map_.find(table + "_adamw_m"); | |
| 709 | 4 | auto second_it = tensor_map_.find(table + "_adamw_v"); | |
| 710 | 4 | auto step_it = tensor_map_.find(table + "_adamw_step"); | |
| 711 | 12 | if (param_it == tensor_map_.end() || first_it == tensor_map_.end() || | |
| 712 | 12 | second_it == tensor_map_.end() || step_it == tensor_map_.end()) { | |
| 713 | ✗ | throw std::runtime_error("AdamW state table not found for table " + table); | |
| 714 | } | ||
| 715 | |||
| 716 | 4 | std::string step_value; | |
| 717 | 4 | step_it->second->Get(kAdamWStepKey, step_value, tid); | |
| 718 | float step = | ||
| 719 | 4 | step_value.empty() ? 0.0f : base::ConstArray<float>(step_value).Data()[0]; | |
| 720 | 4 | step += 1.0f; | |
| 721 | 4 | const float bias1 = 1.0f - std::pow(beta1_, step); | |
| 722 | 4 | const float bias2 = 1.0f - std::pow(beta2_, step); | |
| 723 | 4 | if (!(bias1 > 0.0f) || !(bias2 > 0.0f)) { | |
| 724 | ✗ | throw std::runtime_error("AdamW bias correction underflow"); | |
| 725 | } | ||
| 726 | |||
| 727 | 4 | std::vector<uint64_t> key_vec(keys, keys + num_rows); | |
| 728 | 4 | std::vector<base::ConstArray<float>> params; | |
| 729 | 4 | std::vector<base::ConstArray<float>> first; | |
| 730 | 4 | std::vector<base::ConstArray<float>> second; | |
| 731 | 4 | param_it->second->BatchGet(key_vec, ¶ms, tid); | |
| 732 | 4 | first_it->second->BatchGet(key_vec, &first, tid); | |
| 733 | 4 | second_it->second->BatchGet(key_vec, &second, tid); | |
| 734 | |||
| 735 | 4 | const float decay = 1.0f - learning_rate_ * weight_decay_; | |
| 736 | 4 | const float correction1 = 1.0f / bias1; | |
| 737 | 4 | const float correction2 = 1.0f / bias2; | |
| 738 | 8 | for (int64_t row = 0; row < num_rows; ++row) { | |
| 739 | 4 | const auto index = static_cast<size_t>(row); | |
| 740 | 4 | const float* row_grad = grads + row * embedding_dim; | |
| 741 | 4 | std::vector<float> param(static_cast<size_t>(embedding_dim), 0.0f); | |
| 742 | 4 | std::vector<float> moment1(static_cast<size_t>(embedding_dim), 0.0f); | |
| 743 | 4 | std::vector<float> moment2(static_cast<size_t>(embedding_dim), 0.0f); | |
| 744 | 4 | if (params[index].Size() != 0) { | |
| 745 | 4 | if (params[index].Size() != embedding_dim) { | |
| 746 | ✗ | throw std::runtime_error( | |
| 747 | ✗ | "AdamW parameter dimension mismatch for table " + table); | |
| 748 | } | ||
| 749 | 8 | std::copy(params[index].Data(), | |
| 750 | 4 | params[index].Data() + embedding_dim, | |
| 751 | param.begin()); | ||
| 752 | } | ||
| 753 | 4 | if (first[index].Size() != 0) { | |
| 754 | 2 | if (first[index].Size() != embedding_dim) { | |
| 755 | ✗ | throw std::runtime_error( | |
| 756 | ✗ | "AdamW first moment dimension mismatch for table " + table); | |
| 757 | } | ||
| 758 | 4 | std::copy(first[index].Data(), | |
| 759 | 2 | first[index].Data() + embedding_dim, | |
| 760 | moment1.begin()); | ||
| 761 | } | ||
| 762 | 4 | if (second[index].Size() != 0) { | |
| 763 | 2 | if (second[index].Size() != embedding_dim) { | |
| 764 | ✗ | throw std::runtime_error( | |
| 765 | ✗ | "AdamW second moment dimension mismatch for table " + table); | |
| 766 | } | ||
| 767 | 4 | std::copy(second[index].Data(), | |
| 768 | 2 | second[index].Data() + embedding_dim, | |
| 769 | moment2.begin()); | ||
| 770 | } | ||
| 771 | 12 | for (int64_t col = 0; col < embedding_dim; ++col) { | |
| 772 | 8 | const float grad = row_grad[col]; | |
| 773 | 16 | moment1[static_cast<size_t>(col)] = | |
| 774 | 8 | beta1_ * moment1[static_cast<size_t>(col)] + (1.0f - beta1_) * grad; | |
| 775 | 16 | moment2[static_cast<size_t>(col)] = | |
| 776 | 8 | beta2_ * moment2[static_cast<size_t>(col)] + | |
| 777 | 8 | (1.0f - beta2_) * grad * grad; | |
| 778 | 8 | const float m_hat = moment1[static_cast<size_t>(col)] * correction1; | |
| 779 | 8 | const float v_hat = moment2[static_cast<size_t>(col)] * correction2; | |
| 780 | 8 | param[static_cast<size_t>(col)] = | |
| 781 | 8 | decay * param[static_cast<size_t>(col)] - | |
| 782 | 8 | learning_rate_ * m_hat / (std::sqrt(v_hat) + epsilon_); | |
| 783 | } | ||
| 784 | 4 | const std::string param_value(reinterpret_cast<const char*>(param.data()), | |
| 785 | 4 | param.size() * sizeof(float)); | |
| 786 | 4 | const std::string first_value(reinterpret_cast<const char*>(moment1.data()), | |
| 787 | 4 | moment1.size() * sizeof(float)); | |
| 788 | const std::string second_value( | ||
| 789 | 4 | reinterpret_cast<const char*>(moment2.data()), | |
| 790 | 4 | moment2.size() * sizeof(float)); | |
| 791 | 4 | param_it->second->Put(keys[index], param_value, tid); | |
| 792 | 4 | first_it->second->Put(keys[index], first_value, tid); | |
| 793 | 4 | second_it->second->Put(keys[index], second_value, tid); | |
| 794 | 4 | } | |
| 795 | const std::string next_step( | ||
| 796 | 4 | reinterpret_cast<const char*>(&step), sizeof(float)); | |
| 797 | 4 | step_it->second->Put(kAdamWStepKey, next_step, tid); | |
| 798 | 4 | } | |
| 799 |