optimizer/optimizer.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <string> | ||
| 4 | #include <vector> | ||
| 5 | #include <unordered_map> | ||
| 6 | #include <cmath> | ||
| 7 | #include <memory> | ||
| 8 | #include <stdexcept> | ||
| 9 | #include "sparse_tensor.h" | ||
| 10 | #include "ps/base/base_client.h" | ||
| 11 | #include "ps/base/parameters.h" | ||
| 12 | |||
| 13 | using ::ParameterCompressReader; | ||
| 14 | using recstore::EmbeddingTableConfig; | ||
| 15 | |||
| 16 | class Optimizer { | ||
| 17 | protected: | ||
| 18 | std::unordered_map<std::string, SparseTensor*> tensor_map_; | ||
| 19 | |||
| 20 | public: | ||
| 21 | 44 | virtual ~Optimizer() { | |
| 22 | 92 | for (auto& pair : tensor_map_) { | |
| 23 | 48 | delete pair.second; | |
| 24 | } | ||
| 25 | 44 | } | |
| 26 | |||
| 27 | virtual int Init(const std::vector<std::string> table_name, | ||
| 28 | const EmbeddingTableConfig& config, | ||
| 29 | BaseKV* base_kv) = 0; | ||
| 30 | virtual int TensorsPerTable() const = 0; | ||
| 31 | |||
| 32 | virtual void Update(std::string table, | ||
| 33 | const ParameterCompressReader* reader, | ||
| 34 | unsigned tid) = 0; | ||
| 35 | virtual void UpdateFlat( | ||
| 36 | std::string table, | ||
| 37 | const base::ConstArray<uint64_t>& keys, | ||
| 38 | const float* grads, | ||
| 39 | int64_t num_rows, | ||
| 40 | int64_t embedding_dim, | ||
| 41 | unsigned tid) = 0; | ||
| 42 | }; | ||
| 43 | |||
| 44 | class SGD : public Optimizer { | ||
| 45 | private: | ||
| 46 | float learning_rate_; | ||
| 47 | |||
| 48 | public: | ||
| 49 | 34 | explicit SGD(float lr = 0.01) : learning_rate_(lr) {} | |
| 50 | 8 | float LearningRate() const { return learning_rate_; } | |
| 51 | |||
| 52 | int Init(const std::vector<std::string> table_name, | ||
| 53 | const EmbeddingTableConfig& config, | ||
| 54 | BaseKV* base_kv) override; | ||
| 55 | 28 | int TensorsPerTable() const override { return 1; } | |
| 56 | void Update(std::string table, | ||
| 57 | const ParameterCompressReader* reader, | ||
| 58 | unsigned tid) override; | ||
| 59 | void UpdateFlat(std::string table, | ||
| 60 | const base::ConstArray<uint64_t>& keys, | ||
| 61 | const float* grads, | ||
| 62 | int64_t num_rows, | ||
| 63 | int64_t embedding_dim, | ||
| 64 | unsigned tid) override; | ||
| 65 | }; | ||
| 66 | |||
| 67 | class AdaGrad : public Optimizer { | ||
| 68 | private: | ||
| 69 | float learning_rate_; | ||
| 70 | float epsilon_; | ||
| 71 | |||
| 72 | public: | ||
| 73 | explicit AdaGrad(float lr = 0.01, float epsilon = 1e-10) | ||
| 74 | : learning_rate_(lr), epsilon_(epsilon) {} | ||
| 75 | |||
| 76 | int Init(const std::vector<std::string> table_name, | ||
| 77 | const EmbeddingTableConfig& config, | ||
| 78 | BaseKV* base_kv) override; | ||
| 79 | ✗ | int TensorsPerTable() const override { return 2; } | |
| 80 | void Update(std::string table, | ||
| 81 | const ParameterCompressReader* reader, | ||
| 82 | unsigned tid) override; | ||
| 83 | void UpdateFlat(std::string table, | ||
| 84 | const base::ConstArray<uint64_t>& keys, | ||
| 85 | const float* grads, | ||
| 86 | int64_t num_rows, | ||
| 87 | int64_t embedding_dim, | ||
| 88 | unsigned tid) override; | ||
| 89 | }; | ||
| 90 | |||
| 91 | class RowWiseAdaGrad : public Optimizer { | ||
| 92 | private: | ||
| 93 | float learning_rate_; | ||
| 94 | float epsilon_; | ||
| 95 | |||
| 96 | public: | ||
| 97 | 4 | explicit RowWiseAdaGrad(float lr = 0.01, float epsilon = 1e-10) | |
| 98 | 4 | : learning_rate_(lr), epsilon_(epsilon) {} | |
| 99 | |||
| 100 | int Init(const std::vector<std::string> table_name, | ||
| 101 | const EmbeddingTableConfig& config, | ||
| 102 | BaseKV* base_kv) override; | ||
| 103 | 2 | int TensorsPerTable() const override { return 2; } | |
| 104 | void Update(std::string table, | ||
| 105 | const ParameterCompressReader* reader, | ||
| 106 | unsigned tid) override; | ||
| 107 | void UpdateFlat(std::string table, | ||
| 108 | const base::ConstArray<uint64_t>& keys, | ||
| 109 | const float* grads, | ||
| 110 | int64_t num_rows, | ||
| 111 | int64_t embedding_dim, | ||
| 112 | unsigned tid) override; | ||
| 113 | }; | ||
| 114 | |||
| 115 | // Sparse AdamW keeps first/second moments and a persisted step counter in | ||
| 116 | // RecStore. Updates are applied to rows present in the submitted sparse | ||
| 117 | // gradient (the same sparse visibility contract as the existing optimizers). | ||
| 118 | class AdamW : public Optimizer { | ||
| 119 | private: | ||
| 120 | float learning_rate_; | ||
| 121 | float beta1_; | ||
| 122 | float beta2_; | ||
| 123 | float epsilon_; | ||
| 124 | float weight_decay_; | ||
| 125 | |||
| 126 | void UpdateRows(const std::string& table, | ||
| 127 | const uint64_t* keys, | ||
| 128 | const float* grads, | ||
| 129 | int64_t num_rows, | ||
| 130 | int64_t embedding_dim, | ||
| 131 | unsigned tid); | ||
| 132 | |||
| 133 | public: | ||
| 134 | 6 | explicit AdamW(float lr = 0.001, | |
| 135 | float beta1 = 0.9, | ||
| 136 | float beta2 = 0.98, | ||
| 137 | float epsilon = 1e-8, | ||
| 138 | float weight_decay = 0.0) | ||
| 139 | 12 | : learning_rate_(lr), | |
| 140 | 6 | beta1_(beta1), | |
| 141 | 6 | beta2_(beta2), | |
| 142 | 6 | epsilon_(epsilon), | |
| 143 | 6 | weight_decay_(weight_decay) {} | |
| 144 | |||
| 145 | int Init(const std::vector<std::string> table_name, | ||
| 146 | const EmbeddingTableConfig& config, | ||
| 147 | BaseKV* base_kv) override; | ||
| 148 | ✗ | int TensorsPerTable() const override { return 4; } | |
| 149 | void Update(std::string table, | ||
| 150 | const ParameterCompressReader* reader, | ||
| 151 | unsigned tid) override; | ||
| 152 | void UpdateFlat(std::string table, | ||
| 153 | const base::ConstArray<uint64_t>& keys, | ||
| 154 | const float* grads, | ||
| 155 | int64_t num_rows, | ||
| 156 | int64_t embedding_dim, | ||
| 157 | unsigned tid) override; | ||
| 158 | }; | ||
| 159 | |||
| 160 | std::unique_ptr<Optimizer> CreateOptimizer(const json& config); | ||
| 161 |