| Line |
Branch |
Exec |
Source |
| 1 |
|
|
#include "framework/pytorch/op_sparse_gradient_commit.h" |
| 2 |
|
|
#include "framework/pytorch/op_internal.h" |
| 3 |
|
|
|
| 4 |
|
|
namespace recstore::framework { |
| 5 |
|
|
namespace { |
| 6 |
|
✗ |
int64_t sparse_gradient_commit_async( |
| 7 |
|
|
const std::string& table, |
| 8 |
|
|
const torch::Tensor& ids, |
| 9 |
|
|
const torch::Tensor& grads, |
| 10 |
|
|
int64_t session, |
| 11 |
|
|
int64_t step, |
| 12 |
|
|
int64_t rank, |
| 13 |
|
|
int64_t world, |
| 14 |
|
|
double lr, |
| 15 |
|
|
int64_t divisor) { |
| 16 |
0/30
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
✗ Branch 18 not taken.
✗ Branch 20 not taken.
✗ Branch 21 not taken.
✗ Branch 22 not taken.
✗ Branch 23 not taken.
✗ Branch 24 not taken.
✗ Branch 25 not taken.
✗ Branch 27 not taken.
✗ Branch 28 not taken.
✗ Branch 30 not taken.
✗ Branch 31 not taken.
✗ Branch 32 not taken.
✗ Branch 33 not taken.
✗ Branch 34 not taken.
✗ Branch 35 not taken.
|
✗ |
TORCH_CHECK(ids.dim() == 1 && ids.scalar_type() == torch::kInt64 && |
| 17 |
|
|
grads.dim() == 2 && grads.scalar_type() == torch::kFloat32 && |
| 18 |
|
|
ids.numel() == grads.size(0), |
| 19 |
|
|
"Invalid sparse gradients"); |
| 20 |
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
✗ |
auto h = MakeSparseGradientCommitHeader( |
| 21 |
|
|
session, step, rank, world, grads.size(1), lr, divisor, 0); |
| 22 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
auto keys = ids.is_cuda() ? StageCudaTensorToPinnedCpuAsyncNoCast(ids) |
| 23 |
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
|
✗ |
: ids.contiguous().clone(); |
| 24 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
auto keys_ready = RecordCudaCompletionForTensor(ids); |
| 25 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
auto values = grads.is_cuda() ? StageCudaTensorToPinnedCpuAsyncNoCast(grads) |
| 26 |
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
|
✗ |
: grads.contiguous().clone(); |
| 27 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
auto grads_ready = RecordCudaCompletionForTensor(grads); |
| 28 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
auto client = GetConcreteKVClientOp(); |
| 29 |
0/16
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 17 not taken.
✗ Branch 18 not taken.
✗ Branch 20 not taken.
✗ Branch 21 not taken.
✗ Branch 23 not taken.
✗ Branch 24 not taken.
|
✗ |
return client->SparseGradientCommitAsync( |
| 30 |
|
|
table, |
| 31 |
|
|
h, |
| 32 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
ToRecTensor(keys, base::DataType::UINT64), |
| 33 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
ToRecTensor(values, base::DataType::FLOAT32), |
| 34 |
|
✗ |
[keys, values, ids, grads, keys_ready, grads_ready] { |
| 35 |
|
✗ |
keys_ready(); |
| 36 |
|
✗ |
grads_ready(); |
| 37 |
|
✗ |
}); |
| 38 |
|
✗ |
} |
| 39 |
|
|
|
| 40 |
|
✗ |
void sparse_gradient_commit_wait(int64_t handle) { |
| 41 |
0/2
✗ Branch 3 not taken.
✗ Branch 4 not taken.
|
✗ |
GetConcreteKVClientOp()->WaitSparseGradientCommit(handle); |
| 42 |
|
✗ |
} |
| 43 |
|
|
|
| 44 |
|
✗ |
void sparse_gradient_commit_close( |
| 45 |
|
|
const std::string& table, |
| 46 |
|
|
int64_t session, |
| 47 |
|
|
int64_t step, |
| 48 |
|
|
int64_t rank, |
| 49 |
|
|
int64_t world, |
| 50 |
|
|
int64_t dim, |
| 51 |
|
|
double lr, |
| 52 |
|
|
int64_t divisor, |
| 53 |
|
|
bool abort) { |
| 54 |
0/4
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
|
✗ |
auto h = MakeSparseGradientCommitHeader( |
| 55 |
|
|
session, step, rank, world, dim, lr, divisor, abort ? 3 : 2); |
| 56 |
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
|
✗ |
GetConcreteKVClientOp()->CloseSparseGradientCommit(table, h); |
| 57 |
|
✗ |
} |
| 58 |
|
|
} // namespace |
| 59 |
|
16 |
void RegisterSparseGradientCommitOps(torch::Library& m) { |
| 60 |
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 |
m.def("sparse_gradient_commit_async", sparse_gradient_commit_async); |
| 61 |
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 |
m.def("sparse_gradient_commit_wait", sparse_gradient_commit_wait); |
| 62 |
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 |
m.def("sparse_gradient_commit_close", sparse_gradient_commit_close); |
| 63 |
|
16 |
} |
| 64 |
|
|
} // namespace recstore::framework |
| 65 |
|
|
|