ps/base/sparse_gradient_commit.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <algorithm> | ||
| 4 | #include <cmath> | ||
| 5 | #include "base/array.h" | ||
| 6 | #include "base/tensor.h" | ||
| 7 | #include <cstdint> | ||
| 8 | #include <cstring> | ||
| 9 | #include <functional> | ||
| 10 | #include <limits> | ||
| 11 | #include <mutex> | ||
| 12 | #include <stdexcept> | ||
| 13 | #include <string> | ||
| 14 | #include <unordered_map> | ||
| 15 | #include <utility> | ||
| 16 | #include <vector> | ||
| 17 | |||
| 18 | namespace recstore { | ||
| 19 | |||
| 20 | // Homogeneous-host wire format, shared by RPC, SHM and RDMA transports. | ||
| 21 | struct SparseGradientCommitHeader { | ||
| 22 | uint64_t session = 0; | ||
| 23 | int64_t step = 0; | ||
| 24 | uint32_t rank = 0; | ||
| 25 | uint32_t world = 1; | ||
| 26 | uint32_t dim = 0; | ||
| 27 | uint32_t rows = 0; | ||
| 28 | uint32_t chunk = 0; | ||
| 29 | uint32_t chunks = 1; | ||
| 30 | uint32_t kind = 0; // 0: submit, 1: query, 2: close, 3: abort | ||
| 31 | float learning_rate = 0.01f; | ||
| 32 | uint32_t divisor = 1; | ||
| 33 | uint32_t version = 1; | ||
| 34 | }; | ||
| 35 | static_assert(sizeof(SparseGradientCommitHeader) == 56, | ||
| 36 | "SparseGradientCommitHeader wire size"); | ||
| 37 | |||
| 38 | inline SparseGradientCommitHeader | ||
| 39 | 270 | DecodeSparseGradientCommit(const std::string& payload) { | |
| 40 | 270 | SparseGradientCommitHeader h; | |
| 41 | 270 | if (payload.size() < sizeof(h)) { | |
| 42 | ✗ | throw std::invalid_argument("Truncated sparse gradient commit header"); | |
| 43 | } | ||
| 44 | 270 | std::memcpy(&h, payload.data(), sizeof(h)); | |
| 45 | 268 | if (h.version != 1 || !h.session || !h.world || h.world > 65536 || | |
| 46 | 266 | h.rank >= h.world || !h.dim || h.dim > 65536 || h.step < 0 || | |
| 47 | 266 | h.kind > 3 || !h.chunks || h.chunks > 1048576 || h.chunk >= h.chunks || | |
| 48 | 266 | (h.divisor != 1 && h.divisor != h.world) || | |
| 49 | 538 | !std::isfinite(h.learning_rate) || h.learning_rate < 0) { | |
| 50 | 4 | throw std::invalid_argument("Invalid sparse gradient commit metadata"); | |
| 51 | } | ||
| 52 | 266 | const uint64_t bytes = | |
| 53 | 266 | sizeof(h) + uint64_t(h.rows) * (8 + uint64_t(h.dim) * 4); | |
| 54 | 266 | if (bytes != payload.size() || (h.kind != 0 && h.rows != 0)) { | |
| 55 | ✗ | throw std::invalid_argument("Invalid sparse gradient commit payload size"); | |
| 56 | } | ||
| 57 | 266 | return h; | |
| 58 | } | ||
| 59 | |||
| 60 | 92 | inline std::string PackSparseGradientCommit(SparseGradientCommitHeader h, | |
| 61 | const uint64_t* keys = nullptr, | ||
| 62 | const float* grads = nullptr) { | ||
| 63 | std::string result( | ||
| 64 | 92 | sizeof(h) + size_t(h.rows) * (8 + size_t(h.dim) * 4), '\0'); | |
| 65 | 92 | std::memcpy(result.data(), &h, sizeof(h)); | |
| 66 | 92 | if (h.rows) { | |
| 67 | 40 | if (!keys || !grads) | |
| 68 | ✗ | throw std::invalid_argument("Missing sparse gradient commit rows"); | |
| 69 | 40 | std::memcpy(result.data() + sizeof(h), keys, size_t(h.rows) * 8); | |
| 70 | 40 | std::memcpy(result.data() + sizeof(h) + size_t(h.rows) * 8, | |
| 71 | grads, | ||
| 72 | 40 | size_t(h.rows) * h.dim * 4); | |
| 73 | } | ||
| 74 | 92 | DecodeSparseGradientCommit(result); | |
| 75 | 90 | return result; | |
| 76 | 2 | } | |
| 77 | |||
| 78 | class SparseGradientCommitClient { | ||
| 79 | public: | ||
| 80 | 154 | virtual ~SparseGradientCommitClient() = default; | |
| 81 | ✗ | virtual int64_t ExchangeSparseGradientCommit(const std::string& table, | |
| 82 | const std::string& payload) { | ||
| 83 | ✗ | throw std::runtime_error( | |
| 84 | ✗ | "Sparse gradient commit protocol unsupported by this PS backend"); | |
| 85 | } | ||
| 86 | |||
| 87 | ✗ | virtual size_t SparseGradientCommitPayloadLimit() const { return 256 * 1024; } | |
| 88 | |||
| 89 | 8 | virtual int64_t SparseGradientCommit( | |
| 90 | const std::string& table, | ||
| 91 | SparseGradientCommitHeader h, | ||
| 92 | const base::ConstArray<uint64_t>& keys, | ||
| 93 | const base::RecTensor& grads) { | ||
| 94 | 8 | if (h.kind != 0) | |
| 95 | 2 | return ExchangeSparseGradientCommit(table, PackSparseGradientCommit(h)); | |
| 96 | 6 | if (grads.dim() != 2 || | |
| 97 | 6 | grads.shape(0) != static_cast<int64_t>(keys.Size()) || | |
| 98 | 12 | grads.shape(1) != h.dim || grads.dtype() != base::DataType::FLOAT32) { | |
| 99 | ✗ | throw std::invalid_argument("Invalid sparse gradient shape"); | |
| 100 | } | ||
| 101 | 6 | const size_t limit = SparseGradientCommitPayloadLimit(); | |
| 102 | 6 | const size_t row_bytes = 8 + size_t(h.dim) * 4; | |
| 103 | 6 | if (limit < sizeof(h) + row_bytes) | |
| 104 | ✗ | throw std::runtime_error("Gradient row exceeds transport slot"); | |
| 105 | 6 | const size_t rows = (limit - sizeof(h)) / row_bytes; | |
| 106 | 6 | h.chunks = std::max<size_t>(1, (keys.Size() + rows - 1) / rows); | |
| 107 | 6 | int64_t committed = -1; | |
| 108 | 12 | for (h.chunk = 0; h.chunk < h.chunks; ++h.chunk) { | |
| 109 | 6 | const size_t begin = size_t(h.chunk) * rows; | |
| 110 | 6 | h.rows = std::min(rows, keys.Size() - begin); | |
| 111 | 6 | committed = ExchangeSparseGradientCommit( | |
| 112 | table, | ||
| 113 | 12 | PackSparseGradientCommit( | |
| 114 | h, | ||
| 115 | 6 | h.rows ? keys.Data() + begin : nullptr, | |
| 116 | 6 | h.rows ? grads.data_as<float>() + begin * h.dim : nullptr)); | |
| 117 | } | ||
| 118 | 6 | return committed; | |
| 119 | } | ||
| 120 | }; | ||
| 121 | |||
| 122 | // One in-flight step per table. RPC handlers acknowledge chunks immediately; | ||
| 123 | // they never block a server worker waiting for another trainer. | ||
| 124 | class SparseGradientCommitCoordinator { | ||
| 125 | public: | ||
| 126 | using Apply = std::function<void( | ||
| 127 | const std::vector<uint64_t>&, const std::vector<float>&)>; | ||
| 128 | |||
| 129 | int64_t | ||
| 130 | 108 | Exchange(const std::string& table, std::string payload, const Apply& apply) { | |
| 131 | 108 | const auto h = DecodeSparseGradientCommit(payload); | |
| 132 | // ponytail: one lock per PS shard; split by table if independent jobs need | ||
| 133 | // concurrent commits. The fused BagPipe table already requires ordering. | ||
| 134 | 108 | std::lock_guard<std::mutex> guard(mu_); | |
| 135 | 108 | auto it = states_.find(table); | |
| 136 | 108 | if (it == states_.end()) { | |
| 137 | 14 | if (h.kind != 0 || h.step != 0) { | |
| 138 | ✗ | throw std::runtime_error("Unknown sparse gradient session"); | |
| 139 | } | ||
| 140 | 14 | State state; | |
| 141 | 14 | state.header = h; | |
| 142 | 14 | state.ranks.resize(h.world); | |
| 143 | 14 | state.closed.resize(h.world, false); | |
| 144 | 14 | it = states_.emplace(table, std::move(state)).first; | |
| 145 | 14 | } | |
| 146 | 108 | auto& state = it->second; | |
| 147 | 108 | const auto& initial = state.header; | |
| 148 | 108 | if (initial.session != h.session || initial.world != h.world || | |
| 149 | 108 | initial.dim != h.dim || initial.learning_rate != h.learning_rate || | |
| 150 | 108 | initial.divisor != h.divisor) { | |
| 151 | ✗ | throw std::runtime_error( | |
| 152 | ✗ | "Sparse gradient session configuration mismatch"); | |
| 153 | } | ||
| 154 | 108 | if (h.kind == 3) | |
| 155 | 2 | state.failed = true; | |
| 156 | 108 | if (state.failed) | |
| 157 | 14 | throw std::runtime_error( | |
| 158 | 28 | "Sparse gradient session aborted or commit failed"); | |
| 159 | 94 | if (h.kind == 1) | |
| 160 | 4 | return state.committed; | |
| 161 | 90 | if (h.kind == 2) { | |
| 162 | 6 | if (state.committed < h.step) | |
| 163 | ✗ | throw std::runtime_error("Closing an uncommitted step"); | |
| 164 | 6 | state.closed[h.rank] = true; | |
| 165 | 6 | const auto committed = state.committed; | |
| 166 | 6 | bool all_closed = true; | |
| 167 | 24 | for (bool closed : state.closed) | |
| 168 | 18 | all_closed &= closed; | |
| 169 | 6 | if (all_closed) | |
| 170 | 2 | states_.erase(it); | |
| 171 | 6 | return committed; | |
| 172 | } | ||
| 173 | // A transport retry must never apply an already committed step twice. | ||
| 174 | 84 | if (h.step <= state.committed) | |
| 175 | 8 | return state.committed; | |
| 176 | 76 | if (h.step != state.committed + 1 || state.closed[h.rank]) { | |
| 177 | 2 | throw std::runtime_error( | |
| 178 | 4 | "Sparse gradient commits must be submitted in order"); | |
| 179 | } | ||
| 180 | 74 | auto& chunks = state.ranks[h.rank]; | |
| 181 | 74 | if (chunks.empty()) | |
| 182 | 42 | chunks.resize(h.chunks); | |
| 183 | 74 | if (chunks.size() != h.chunks) | |
| 184 | ✗ | throw std::runtime_error("Gradient chunk count changed"); | |
| 185 | 74 | if (!chunks[h.chunk].empty()) { | |
| 186 | 16 | if (chunks[h.chunk] != payload) | |
| 187 | 2 | throw std::runtime_error("Conflicting gradient chunk retry"); | |
| 188 | 14 | return state.committed; | |
| 189 | } | ||
| 190 | 58 | chunks[h.chunk] = std::move(payload); | |
| 191 | 126 | for (const auto& rank : state.ranks) { | |
| 192 | 108 | if (rank.empty()) | |
| 193 | 40 | return state.committed; | |
| 194 | 188 | for (const auto& chunk : rank) | |
| 195 | 120 | if (chunk.empty()) | |
| 196 | 16 | return state.committed; | |
| 197 | } | ||
| 198 | 18 | std::vector<uint64_t> keys; | |
| 199 | 18 | std::vector<float> grads; | |
| 200 | 18 | std::unordered_map<uint64_t, size_t> positions; | |
| 201 | // Fixed rank/chunk order makes network arrival order irrelevant. | ||
| 202 | 56 | for (const auto& rank : state.ranks) { | |
| 203 | 92 | for (const auto& chunk : rank) { | |
| 204 | 54 | const auto ch = DecodeSparseGradientCommit(chunk); | |
| 205 | 54 | const char* key_data = chunk.data() + sizeof(ch); | |
| 206 | 54 | const char* grad_data = key_data + size_t(ch.rows) * 8; | |
| 207 | 88 | for (size_t row = 0; row < ch.rows; ++row) { | |
| 208 | uint64_t key; | ||
| 209 | 34 | std::memcpy(&key, key_data + row * 8, 8); | |
| 210 | 34 | auto entry = positions.emplace(key, keys.size()); | |
| 211 | 34 | if (entry.second) { | |
| 212 | 18 | keys.push_back(key); | |
| 213 | 18 | grads.resize(keys.size() * h.dim, 0.0f); | |
| 214 | } | ||
| 215 | 34 | float* dst = grads.data() + entry.first->second * h.dim; | |
| 216 | 84 | for (size_t col = 0; col < h.dim; ++col) { | |
| 217 | float value; | ||
| 218 | 50 | std::memcpy(&value, grad_data + (row * h.dim + col) * 4, 4); | |
| 219 | 50 | dst[col] += value; | |
| 220 | } | ||
| 221 | } | ||
| 222 | } | ||
| 223 | } | ||
| 224 | 18 | if (h.divisor != 1) | |
| 225 | 28 | for (auto& value : grads) | |
| 226 | 16 | value /= h.divisor; | |
| 227 | try { | ||
| 228 | 18 | if (!keys.empty()) | |
| 229 | 14 | apply(keys, grads); | |
| 230 | 2 | } catch (...) { | |
| 231 | 2 | state.failed = true; | |
| 232 | 2 | throw; | |
| 233 | 2 | } | |
| 234 | 16 | state.committed = h.step; | |
| 235 | 52 | for (auto& rank : state.ranks) | |
| 236 | 36 | rank.clear(); | |
| 237 | 16 | return state.committed; | |
| 238 | 114 | } | |
| 239 | |||
| 240 | private: | ||
| 241 | struct State { | ||
| 242 | SparseGradientCommitHeader header; | ||
| 243 | int64_t committed = -1; | ||
| 244 | bool failed = false; | ||
| 245 | std::vector<std::vector<std::string>> ranks; | ||
| 246 | std::vector<bool> closed; | ||
| 247 | }; | ||
| 248 | std::mutex mu_; | ||
| 249 | std::unordered_map<std::string, State> states_; | ||
| 250 | }; | ||
| 251 | } // namespace recstore | ||
| 252 |