* @brief Write back an embedding matrix and its moment matrices from GPU cache * @param id id of embedding matrix * * The actual operation depends on the protocol. It is safe to call this function multiple times. */
| 1392 | * The actual operation depends on the protocol. It is safe to call this function multiple times. |
| 1393 | */ |
| 1394 | void write_embedding(int id) { |
| 1395 | Protocol protocol = protocols[id]; |
| 1396 | if (id > 0 && embeddings[id] == embeddings[id - 1]) |
| 1397 | return; |
| 1398 | auto &global_embedding = *(solver->embeddings[id]); |
| 1399 | auto &global_moment = *(solver->moments[id]); |
| 1400 | auto &embedding = *embeddings[id]; |
| 1401 | auto &moment = *moments[id]; |
| 1402 | auto &gradient = *gradients[id]; |
| 1403 | std::vector<Index> none, *mapping = &none; |
| 1404 | |
| 1405 | if (protocol & kHeadPartition) |
| 1406 | mapping = &head_global_ids; |
| 1407 | if (protocol & kTailPartition) |
| 1408 | mapping = &tail_global_ids; |
| 1409 | |
| 1410 | if (protocol & kInPlace) { |
| 1411 | embedding.to_host(); |
| 1412 | embedding.scatter(global_embedding, *mapping); |
| 1413 | } |
| 1414 | else { |
| 1415 | gradient.copy(embedding); |
| 1416 | embedding.to_host(); |
| 1417 | for (Index i = 0; i < embedding.count; i++) |
| 1418 | gradient[i] -= embedding[i]; |
| 1419 | gradient.scatter_sub(global_embedding, *mapping); |
| 1420 | } |
| 1421 | |
| 1422 | // only write back partitioned moments |
| 1423 | if (solver->is_train && !(protocol & kGlobal)) |
| 1424 | for (int i = 0; i < num_moment; i++) { |
| 1425 | moment[i].to_host(); |
| 1426 | moment[i].scatter(global_moment[i], *mapping); |
| 1427 | } |
| 1428 | } |
| 1429 | |
| 1430 | /** |
| 1431 | * @brief Load a partition of the sample pool. Update the cache automatically. |
nothing calls this directly
no test coverage detected