GCC Code Coverage Report

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