0. 參數分配
- 稠密參數 → MPI 集合通信（All-Reduce / Broadcast / Reduce-Scatter）。
- 稀疏參數 → brpc Parameter Server 異步推拉。
完全去掉 NCCL/GPU 相關部分。
1. 整體拓撲
┌----------------┐ ┌----------------┐
│ Worker-0 │ │ PS-0 │
│ Worker-1 │?------? │ PS-1 │
│ ... │ brpc │ ... │
│ Worker-N │ │ PS-M │
└----------------┘ └----------------┘▲│MPI（TCP/InfiniBand）▼
MPI_COMM_WORLD（稠密參數）
- 稠密梯度：通過 MPI 標準集合操作（
MPI_Allreduce、MPI_Bcast等）實現同步。 - 稀疏參數：Worker 與 PS 之間用 brpc + protobuf 通信，異步推拉。
2. 關鍵模塊（C++）
cpu_dist/
├── common/
│ ├── tensor.h // 純 CPU 張量（FP32/FP64）
│ └── mpi_context.h // MPI_Init / Finalize 封裝
├── dense/
│ ├── mpi_allreduce.h // MPI All-Reduce 封裝
│ └── optimizer.h // 本地 SGD / AdamW
├── sparse/
│ ├── ps_server.h/cc // brpc Parameter Server
│ ├── ps_client.h/cc // brpc Client
│ └── table.h // 稀疏表（unordered_map + 鎖）
├── proto/
│ └── message.proto // protobuf 消息
└── launcher.cc // 主進程入口
3. MPI 通信層（稠密參數）
3.1 封裝 MPI All-Reduce
// dense/mpi_allreduce.h
class MPIAllReduce {public:explicit MPIAllReduce(MPI_Comm comm) : comm_(comm) {}template <typename T>void AllReduceSum(std::vector<T>& buf) {std::vector<T> recv(buf.size());MPI_Allreduce(buf.data(), recv.data(), buf.size(),GetMPIType<T>(), MPI_SUM, comm_);buf.swap(recv);}private:MPI_Comm comm_;
};
- 支持 float / double / int。
- 支持 In-place All-Reduce（
MPI_IN_PLACE）。
4. brpc Parameter Server（稀疏參數）
與之前設計一致，僅通信后端為 brpc：
- proto 定義不變（
PullRequest,PushRequest）。 - PS 端 實現
brpc::Service，用brpc::Server啟動。 - Worker 端 用
brpc::Channel連接 PS，支持 輪詢/一致性哈希 負載均衡。
5. 主進程結構（launcher.cc）
int main(int argc, char* argv[]) {MPI_Init(&argc, &argv);int rank, size;MPI_Comm_rank(MPI_COMM_WORLD, &rank);MPI_Comm_size(MPI_COMM_WORLD, &size);bool is_ps = (rank >= FLAGS_worker_num);if (!is_ps) {// WorkerMPIAllReduce ar(MPI_COMM_WORLD);PSClient ps(FLAGS_ps_list);WorkerLoop(ar, ps);} else {// Parameter ServerPSServer server;server.Start(FLAGS_ps_port);}MPI_Finalize();
}
6. Worker 主循環
void WorkerLoop(MPIAllReduce& ar, PSClient& ps) {Model model;DataLoader dl(FLAGS_data_path);for (int step = 0; step < FLAGS_max_step; ++step) {auto batch = dl.Next();std::vector<float> dense_grad;std::vector<int64_t> sparse_keys;std::vector<float> sparse_grad;// 前向 & 反向model.Backward(batch, &dense_grad, &sparse_keys, &sparse_grad);// 1. 稠密梯度 MPI All-Reducear.AllReduceSum(dense_grad);// 2. 稀疏梯度異步 Pushps.PushAsync(0, sparse_keys, sparse_grad);// 3. 稀疏參數 Pullstd::vector<float> sparse_emb;ps.Pull(0, sparse_keys, &sparse_emb);// 4. 參數更新model.Update(dense_grad, sparse_emb);}
}
7. 部署與運行
7.1 啟動腳本（OpenMPI）
# 4 worker + 2 ps
mpirun -np 6 \-x LD_LIBRARY_PATH \./launcher \--worker_num 4 \--ps_list "0.0.0.0:8000,0.0.0.0:8001"
- worker_num 前
rank 0~3為 Worker，后rank 4~5為 PS。 - MPI 負責稠密通信，brpc 負責稀疏通信，兩者互不干擾。
8. 性能調優
| 項 | 建議 |
|---|---|
| MPI | 使用 OpenMPI 4.x 或 Intel MPI（CPU 親和、NUMA 優化）。 |
| brpc | 配置 輪詢 + 批處理（64~256 key/RPC），開啟 8bit 量化壓縮。 |
| 線程 | MPI 與 brpc 線程分離，brpc 用 bthread，避免與 MPI 線程沖突。 |
至此，“CPU + MPI（稠密） + brpc Parameter Server（稀疏）” 的完整框架已就緒。
)
)
