拓冰网站定制 
MXNet mxnet.kvstoreKVStore 参数同步接口 API 与分布式训练实现解析【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet本文基于 API 文档页 mxnet.kvstore该页通过 automodule 指令渲染 python/mxnet/kvstore.py 的完整模块文档展开系统讲解mxnet.kvstore中 KVStore 参数同步接口的创建方式、init/push/pull 数据流、优化器注册与梯度压缩等全部公开 API并深入 src/kvstore 下的 C 实现帮助读者掌握在单机多卡与多机分布式训练中完成参数聚合、更新与同步的完整方案。一、模块定位mxnet.kvstore 是做什么的mxnet.kvstore 模块的官方定义只有一句话Key value store interface of MXNet for parameter synchronization.——它是 MXNet 面向参数同步的键值存储接口。在训练流程中各设备/各机器上的 worker 通过统一的 KVStore 完成三件事初始化把权重等参数的初始值登记进存储init聚合更新各 worker 把本地梯度push上去由存储内的 updater优化器归约并更新权重拉取参数各 worker 把最新权重pull回各自的设备。对外只有一个工厂函数入口 mx.kv.create其内部调用 C APIMXKVStoreCreate声明见 include/mxnet/c_api.h创建 store并顺带把 handle 注册给 profilerset_kvstore_handle以便后续性能剖析关联到 kvstore 节点。二、创建 KVStore 与六种 store 类型create(namelocal) 的文档字符串给出了六类 store 的语义这是选型的核心依据类型适用场景行为local单机把所有梯度拷贝到 CPU 内存并在那里更新权重device单机多卡在 GPU 上聚合梯度并更新权重并尝试使用 GPU peer-to-peer 通信加速nccl单机/多机 NCCL基于 NCCL 的集合通信实现KVStoreNCCLdist_sync多机同步类似local但 batch-size 指每台机器的 batch sizen 台机器、单机 batch b 时等价于local的 n*b 同步批量训练dist_device_sync多机同步与dist_sync的区别类似device与local的差别聚合在 GPU 上做dist_async多机异步收到任何一台机器的梯度就立即更新权重同一权重上不会有两个更新并发但到达顺序不保证从源码结构看类型分派发生在 C 侧的 KVStore::Create名称中含dist创建kvstore::KVStoreDistps-lite 后端。若类型不含_async且当前是 rank 0 的 worker会向 server 发送kSyncMode命令把 server 配置为同步模式编译时未定义MXNET_USE_DIST_KVSTORE则会LOG(FATAL)提示compile with USE_DIST_KVSTORE1即分布式类型需要以USE_DIST_KVSTORE1重新编译名称中含nccl创建kvstore::KVStoreNCCL同样要求编译期USE_NCCL1其余local、device创建kvstore::KVStoreLocal其中含device的走设备通信路径use_device_comm true。import mxnet as mx kv mx.kv.create(local) # 默认 kv mx.kv.create(device) # 多卡聚合放 GPU kv mx.kv.create(dist_sync) # 多机同步需编译时开启 USE_DIST_KVSTOREKVStore对象本身只是 C handle 的 Python 包装KVStore.init断言 handle 为KVStoreHandle析构时调用MXKVStoreFree释放。三、核心数据流init、push、pull 与 pushpull3.1 init先登记后读写init(key, value) 的语义要点key 支持str、int或二者的序列value 支持NDArray、RowSparseNDArray及其列表每个 key 必须先init再push/pull否则 C 侧会触发key k has not been inited检查失败见 KVStoreLocal::PushImpl 中的CHECK(!local.is_none())多 worker 对同一 key 并发init时只有 rank 0 提供的值生效。shape (2, 3) kv mx.kv.create(local) kv.init(3, mx.nd.ones(shape) * 2) # 初始化单个 key a mx.nd.zeros(shape) kv.pull(3, outa) # 拉出初始值全为 2 keys [5, 7, 9] kv.init(keys, [mx.nd.ones(shape)] * len(keys)) # 批量初始化 # 稀疏row_sparse值同样可以初始化 kv.init(4, mx.nd.ones(shape).tostype(row_sparse)) b mx.nd.sparse.zeros(row_sparse, shape) kv.row_sparse_pull(4, row_idsmx.nd.array([0, 1]), outb)实现细节上KVStoreLocal在 InitImpl 中把初始值拷贝到 pinned 内存上下文pinned_ctx_作为local_权威副本并调用comm_-Init登记形状/类型字符串 key 会经 Init 映射成递增整数str_key_dict_并禁止重复 init 与混用 key 类型Mixed key types are not allowed。3.2 push异步提交与归约push(key, value, priority0) 的文档明确了几个重要语义函数把操作加入执行引擎engine后立即返回实际操作异步执行对同一 key 连续 push不保证序列化push 的完成也不意味着此前所有 push 都已完成worker 之间没有内建同步——需要全局同步时应调用_barrier()priority越大越可能被优先调度。单设备 push 直接覆盖值多设备 pushvalue 为各设备的 NDArray 列表会先归约聚合# 4 张 GPU 上的梯度各为 1聚合后拉出即 4 gpus [mx.gpu(i) for i in range(4)] b [mx.nd.ones(shape, gpu) for gpu in gpus] kv.push(3, b) kv.pull(3, outa) print(a.asnumpy()) # [[4 4 4] [4 4 4]]批量多 key 的写法keys 与 values 等长列表多设备时为二维列表keys [4, 5, 6] kv.push(keys, [mx.nd.ones(shape)] * len(keys)) b [mx.nd.zeros(shape)] * len(keys) kv.pull(keys, outb) # 多设备keys 个 key每个 key 有 len(gpus) 份设备值 keys [7, 8, 9] b [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys) kv.push(keys, b)Python 层通过 _ctype_key_value 把 key/value 转成 ctypes 数组字符串 key 走MXKVStorePushEx/MXKVStoreInitEx整数 key 走MXKVStorePush/MXKVStoreInitC 原型见 include/mxnet/c_api.h。C 侧 PushImpl 按 key 分组后调用comm_-Reduce完成跨设备归约如果注册了 updater归约结果会经过 updater 更新local_副本没有 updater 时直接以归约值替换副本——这解释了为什么push 全 1、无优化器时 pull 出 4。3.3 pull阻塞语义与最新值保证pull(key, out, priority0, ignore_sparseTrue) 的语义提交后立即返回但后续读取out会阻塞直到 pull 完成pull 在所有先前 pull、以及同一 key 的最后一次 push 完成后才执行返回值保证是 store 中最新值对分布式 kvstorepull不支持RowSparseNDArray应改用row_sparse_pull。# 拉到多个设备 b [mx.nd.ones(shape, gpu) for gpu in gpus] kv.pull(3, outb)3.4 pushpull一次调用完成推与拉pushpull(key, value, outNone, priority0) 是 push 与 pull 的合并形式先把value推上去再把 store 更新后的值拉入outout缺省时写回value返回值保证是最新值。分布式 kvstore 下同样不支持RowSparseNDArray。从 KVStoreLocal::PushPullImpl 看其实现就是PushImpl后接PullImpl合并的收益在于减少一次引擎调度和一轮通信往返kv.pushpull(3, mx.nd.ones(shape) * 8, outa) # 多设备聚合 kv.pushpull(3, b, outa)3.5 row_sparse_pull按行拉取稀疏参数row_sparse_pull(key, out, priority0, row_ids) 专为 embedding 类稀疏参数设计只按row_ids1-D NDArray不要求唯一、不要求有序拉取对应行。文档示例覆盖了三种情况shape (3, 3) kv.init(3, mx.nd.ones(shape).tostype(row_sparse)) a mx.nd.sparse.zeros(row_sparse, shape) kv.row_sparse_pull(3, outa, row_idsmx.nd.array([0, 2], dtypeint64)) print(a.asnumpy()) # 第 0、2 行为 1第 1 行为 0 kv.row_sparse_pull(3, outa, row_idsmx.nd.array([2, 2], dtypeint64)) # 重复 row_id kv.row_sparse_pull(3, outa, row_idsmx.nd.array([1, 0], dtypeint64)) # 无序 row_id一个 Python 层的优化值得注意当所有 out 共享同一份row_ids时single_rowid分支第 453-L470 行C API 只调用一次其余设备副本直接copyto避免重复的稀疏拉取。C 侧 PullRowSparseImpl 会先对 row_id 做Unique去重再广播。四、优化器注册与更新规则4.1 set_optimizer单机本地更新多机发送到 serverset_optimizer(optimizer) 把优化器注册进 store其行为随 store 类型分叉单机调用opt.get_updater(optimizer)得到 updater 函数后走本地_set_updaterpush 到的梯度直接在本地执行更新多机dist 且当前是 worker 节点用pickleASCII 协议 0序列化整个优化器经_send_command_to_servers下发到所有 server 节点函数在所有 server 更新完成后才返回若优化器是multi_precision还会追加发送kSetMultiPrecision命令。kv mx.kv.create() shape (2, 2) weight mx.nd.zeros(shape) kv.init(3, weight) kv.set_optimizer(mx.optimizer.SGD()) # 默认学习率 0.01 kv.push(3, mx.nd.ones(shape)) kv.pull(3, outweight) weight.asnumpy() # 全 -0.01SGD 已按梯度下降更新server 端如何接住这条命令见 python/mxnet/kvstore_server.pyKVStoreServer._controller中cmd_id 0时反序列化出 optimizer 并kvstore.set_optimizer(optimizer)。而 worker 侧可用命令编号表定义在 _get_kvstore_server_command_type命令编号kController下发优化器等控制命令0kSetMultiPrecision1kStopServer2kSyncMode3kSetGradientCompression4kSetProfilerParams54.2 _set_updater注册原始更新函数_set_updater(updater) 允许直接注册形如updater(key, input, stored)的更新函数Python 层经 _updater_wrapper 包成 ctypes 回调并同时注册 int key 与 str key 两个 C 回调MXKVStoreSetUpdaterEx。它只改本地 store多机场景必须用set_optimizer。文档示例def update(key, input, stored): print(update on key: %d % key) stored input * 2 kv._set_updater(update) kv.push(3, mx.nd.ones(shape))注意 key 类型一致性从 SetKeyType 可见一个 store 内 int key 与 str key 不可混用否则CHECK_EQ直接失败。4.3 优化器状态持久化save_optimizer_states(fname, dump_optimizerFalse) 与 load_optimizer_states(fname) 用于训练中 checkpoint 时保存/恢复优化器updater内部状态如动量项dump_optimizerTrue时连优化器本身学习率、weight decay 调度等一并保存。两者都断言self._updater is not None即分布式训练下不适用状态在 server 侧。五、梯度压缩set_gradient_compression 与 2bit 量化set_gradient_compression(compression_params) 为通信开启低比特梯度量化。目前仅支持type2bit可选参数threshold正浮点例如kv.set_gradient_compression({type: 2bit, threshold: 0.5})文档对 2bit 方案的描述是按threshold把梯度三值化大于阈值钳到阈值、小于负阈值钳到负阈值、其余置 0每个值用 2 bit 表示原梯度每 16 个 float 只需 1 个 float 即可表示被截断的残差保留在发送端累加进下一轮梯度发送从而在压缩通信的同时不丢信息。文档同时给出两类生效范围kvstore 为local此处实指设备间路径时压缩 GPU 到合并 GPU 的通信各 GPU 需额外存 residual 数组显存占用增加kvstore 为dist时只压缩 worker 到 server 的方向多 GPU 机器内 GPU-GPU / GPU-CPU 以及 server 到 workerpull方向不压缩且各 worker 的 CPU 内存因 residual 增加。实现上Python 层会先检查self.type含device或dist否则抛出 Gradient compression is not supported for this type of kvstore即纯localCPU 路径不支持C 侧 KVStoreDist::SetGradientCompression 在本地记录参数的同时由 rank 0 向所有 server 发送kSetGradientCompression命令并附带EncodeParams()编码串保证 server 同步启用。压缩算子本体位于 src/kvstore/gradient_compression.cc 与 gradient_compression.cu。六、分布式角色rank、num_workers、barrier 与 KVStoreServerKVStore提供三个与分布式拓扑相关的只读属性typestore 类型字符串如dist_syncrank本 worker 节点在[0, num_workers)内的编号底层调MXKVStoreGetRanknum_workersworker 节点数底层调MXKVStoreGetGroupSize。_barrier() 在所有 worker 间建立全局屏障典型用途是机器 0 先init完成再让所有 workerpull的时序保证_send_command_to_servers(head, body) 则向所有 server 节点下发命令server 会调用其KVStoreServer.controller执行命令全部执行完才返回。分布式 kvstore 的后端是 ps-liteKVStoreDist 继承KVStoreLocalworker 节点创建ps::KVWorker并ps::StartAsync启动相关环境变量从源码结构看MXNET_KVSTORE_BIGARRAY_BOUND大数组边界默认1000*1000kvstore_dist.h 第 58 行MXNET_KVSTORE_USETREE单机 device 通信是否使用二叉树拓扑见 KVStoreLocal 构造选CommDeviceTree或CommDevicelocal类型为CommCPU。server/scheduler 进程的入口是 python/mxnet/kvstore_server.py模块导入时执行 _init_kvstore_server_module若当前节点不是worker 节点MXKVStoreIsWorkerNode返回 0则以create(dist)建 store、KVStoreServer.run()进入服务循环对应 C APIMXKVStoreRunServer声明见 include/mxnet/c_api.h收到控制命令时由 controller 处理cmd 0 安装 worker 下发的 pickled 优化器。这与 C 侧 KVStoreDist::RunServer 的 ps-lite 启动流程一一对应。分布式路径的验证用例可参考 tests/nightly/dist_sync_kvstore.py、tests/nightly/dist_async_kvstore.py、tests/nightly/dist_device_sync_kvstore.py 与单机单测 tests/python/unittest/test_kvstore.py多卡 device 路径另有 tests/python/gpu/test_kvstore_gpu.py。七、C API 一览与工程要点小结MXKVStore*系列 C 接口集中声明在 include/mxnet/c_api.h与 Python 方法一一对应MXKVStoreCreate/MXKVStoreFree、MXKVStoreInit(Ex)、MXKVStorePush(Ex)、MXKVStorePullWithSparse(Ex)/MXKVStorePull(Ex)/MXKVStorePullRowSparse(Ex)、MXKVStorePushPull(Ex)、MXKVStoreSetUpdater(Ex)、MXKVStoreSetGradientCompression、MXKVStoreGetType/GetRank/GetGroupSize、MXKVStoreBarrier、MXKVStoreRunServer、MXKVStoreSendCommmandToServers等。工程使用要点汇总选型单机默认local多卡聚合放 GPU 用device多机训练用dist_sync/dist_device_sync等价批量为 n*b或dist_async真异步、到达顺序不保证顺序先init后push/pull跨 worker 需要确定时序时用_barrier()语义push 无序列化保证pull/pushpull 保证读到最新值且读取out会阻塞优化器单机用set_optimizer/_set_updater本地更新多机自动 pickled 下发 servercheckpoint 用save_optimizer_states/load_optimizer_states分布式不适用压缩{type: 2bit, threshold: ...}仅对 device/dist 生效代价是残差缓存带来的额外内存编译前提dist_*与nccl类型分别需要USE_DIST_KVSTORE1、USE_NCCL1编译选项未开启时mx.kv.create会直接致命报错。【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 MXNet mxnet.kvstore：KVStore 参数同步接口 API 与分布式训练实现解析
MXNet mxnet.kvstoreKVStore 参数同步接口 API 与分布式训练实现解析【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet本文基于 API 文档页 mxnet.kvstore该页通过 automodule 指令渲染 python/mxnet/kvstore.py 的完整模块文档展开系统讲解mxnet.kvstore中 KVStore 参数同步接口的创建方式、init/push/pull 数据流、优化器注册与梯度压缩等全部公开 API并深入 src/kvstore 下的 C 实现帮助读者掌握在单机多卡与多机分布式训练中完成参数聚合、更新与同步的完整方案。一、模块定位mxnet.kvstore 是做什么的mxnet.kvstore 模块的官方定义只有一句话Key value store interface of MXNet for parameter synchronization.——它是 MXNet 面向参数同步的键值存储接口。在训练流程中各设备/各机器上的 worker 通过统一的 KVStore 完成三件事初始化把权重等参数的初始值登记进存储init聚合更新各 worker 把本地梯度push上去由存储内的 updater优化器归约并更新权重拉取参数各 worker 把最新权重pull回各自的设备。对外只有一个工厂函数入口 mx.kv.create其内部调用 C APIMXKVStoreCreate声明见 include/mxnet/c_api.h创建 store并顺带把 handle 注册给 profilerset_kvstore_handle以便后续性能剖析关联到 kvstore 节点。二、创建 KVStore 与六种 store 类型create(namelocal) 的文档字符串给出了六类 store 的语义这是选型的核心依据类型适用场景行为local单机把所有梯度拷贝到 CPU 内存并在那里更新权重device单机多卡在 GPU 上聚合梯度并更新权重并尝试使用 GPU peer-to-peer 通信加速nccl单机/多机 NCCL基于 NCCL 的集合通信实现KVStoreNCCLdist_sync多机同步类似local但 batch-size 指每台机器的 batch sizen 台机器、单机 batch b 时等价于local的 n*b 同步批量训练dist_device_sync多机同步与dist_sync的区别类似device与local的差别聚合在 GPU 上做dist_async多机异步收到任何一台机器的梯度就立即更新权重同一权重上不会有两个更新并发但到达顺序不保证从源码结构看类型分派发生在 C 侧的 KVStore::Create名称中含dist创建kvstore::KVStoreDistps-lite 后端。若类型不含_async且当前是 rank 0 的 worker会向 server 发送kSyncMode命令把 server 配置为同步模式编译时未定义MXNET_USE_DIST_KVSTORE则会LOG(FATAL)提示compile with USE_DIST_KVSTORE1即分布式类型需要以USE_DIST_KVSTORE1重新编译名称中含nccl创建kvstore::KVStoreNCCL同样要求编译期USE_NCCL1其余local、device创建kvstore::KVStoreLocal其中含device的走设备通信路径use_device_comm true。import mxnet as mx kv mx.kv.create(local) # 默认 kv mx.kv.create(device) # 多卡聚合放 GPU kv mx.kv.create(dist_sync) # 多机同步需编译时开启 USE_DIST_KVSTOREKVStore对象本身只是 C handle 的 Python 包装KVStore.init断言 handle 为KVStoreHandle析构时调用MXKVStoreFree释放。三、核心数据流init、push、pull 与 pushpull3.1 init先登记后读写init(key, value) 的语义要点key 支持str、int或二者的序列value 支持NDArray、RowSparseNDArray及其列表每个 key 必须先init再push/pull否则 C 侧会触发key k has not been inited检查失败见 KVStoreLocal::PushImpl 中的CHECK(!local.is_none())多 worker 对同一 key 并发init时只有 rank 0 提供的值生效。shape (2, 3) kv mx.kv.create(local) kv.init(3, mx.nd.ones(shape) * 2) # 初始化单个 key a mx.nd.zeros(shape) kv.pull(3, outa) # 拉出初始值全为 2 keys [5, 7, 9] kv.init(keys, [mx.nd.ones(shape)] * len(keys)) # 批量初始化 # 稀疏row_sparse值同样可以初始化 kv.init(4, mx.nd.ones(shape).tostype(row_sparse)) b mx.nd.sparse.zeros(row_sparse, shape) kv.row_sparse_pull(4, row_idsmx.nd.array([0, 1]), outb)实现细节上KVStoreLocal在 InitImpl 中把初始值拷贝到 pinned 内存上下文pinned_ctx_作为local_权威副本并调用comm_-Init登记形状/类型字符串 key 会经 Init 映射成递增整数str_key_dict_并禁止重复 init 与混用 key 类型Mixed key types are not allowed。3.2 push异步提交与归约push(key, value, priority0) 的文档明确了几个重要语义函数把操作加入执行引擎engine后立即返回实际操作异步执行对同一 key 连续 push不保证序列化push 的完成也不意味着此前所有 push 都已完成worker 之间没有内建同步——需要全局同步时应调用_barrier()priority越大越可能被优先调度。单设备 push 直接覆盖值多设备 pushvalue 为各设备的 NDArray 列表会先归约聚合# 4 张 GPU 上的梯度各为 1聚合后拉出即 4 gpus [mx.gpu(i) for i in range(4)] b [mx.nd.ones(shape, gpu) for gpu in gpus] kv.push(3, b) kv.pull(3, outa) print(a.asnumpy()) # [[4 4 4] [4 4 4]]批量多 key 的写法keys 与 values 等长列表多设备时为二维列表keys [4, 5, 6] kv.push(keys, [mx.nd.ones(shape)] * len(keys)) b [mx.nd.zeros(shape)] * len(keys) kv.pull(keys, outb) # 多设备keys 个 key每个 key 有 len(gpus) 份设备值 keys [7, 8, 9] b [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys) kv.push(keys, b)Python 层通过 _ctype_key_value 把 key/value 转成 ctypes 数组字符串 key 走MXKVStorePushEx/MXKVStoreInitEx整数 key 走MXKVStorePush/MXKVStoreInitC 原型见 include/mxnet/c_api.h。C 侧 PushImpl 按 key 分组后调用comm_-Reduce完成跨设备归约如果注册了 updater归约结果会经过 updater 更新local_副本没有 updater 时直接以归约值替换副本——这解释了为什么push 全 1、无优化器时 pull 出 4。3.3 pull阻塞语义与最新值保证pull(key, out, priority0, ignore_sparseTrue) 的语义提交后立即返回但后续读取out会阻塞直到 pull 完成pull 在所有先前 pull、以及同一 key 的最后一次 push 完成后才执行返回值保证是 store 中最新值对分布式 kvstorepull不支持RowSparseNDArray应改用row_sparse_pull。# 拉到多个设备 b [mx.nd.ones(shape, gpu) for gpu in gpus] kv.pull(3, outb)3.4 pushpull一次调用完成推与拉pushpull(key, value, outNone, priority0) 是 push 与 pull 的合并形式先把value推上去再把 store 更新后的值拉入outout缺省时写回value返回值保证是最新值。分布式 kvstore 下同样不支持RowSparseNDArray。从 KVStoreLocal::PushPullImpl 看其实现就是PushImpl后接PullImpl合并的收益在于减少一次引擎调度和一轮通信往返kv.pushpull(3, mx.nd.ones(shape) * 8, outa) # 多设备聚合 kv.pushpull(3, b, outa)3.5 row_sparse_pull按行拉取稀疏参数row_sparse_pull(key, out, priority0, row_ids) 专为 embedding 类稀疏参数设计只按row_ids1-D NDArray不要求唯一、不要求有序拉取对应行。文档示例覆盖了三种情况shape (3, 3) kv.init(3, mx.nd.ones(shape).tostype(row_sparse)) a mx.nd.sparse.zeros(row_sparse, shape) kv.row_sparse_pull(3, outa, row_idsmx.nd.array([0, 2], dtypeint64)) print(a.asnumpy()) # 第 0、2 行为 1第 1 行为 0 kv.row_sparse_pull(3, outa, row_idsmx.nd.array([2, 2], dtypeint64)) # 重复 row_id kv.row_sparse_pull(3, outa, row_idsmx.nd.array([1, 0], dtypeint64)) # 无序 row_id一个 Python 层的优化值得注意当所有 out 共享同一份row_ids时single_rowid分支第 453-L470 行C API 只调用一次其余设备副本直接copyto避免重复的稀疏拉取。C 侧 PullRowSparseImpl 会先对 row_id 做Unique去重再广播。四、优化器注册与更新规则4.1 set_optimizer单机本地更新多机发送到 serverset_optimizer(optimizer) 把优化器注册进 store其行为随 store 类型分叉单机调用opt.get_updater(optimizer)得到 updater 函数后走本地_set_updaterpush 到的梯度直接在本地执行更新多机dist 且当前是 worker 节点用pickleASCII 协议 0序列化整个优化器经_send_command_to_servers下发到所有 server 节点函数在所有 server 更新完成后才返回若优化器是multi_precision还会追加发送kSetMultiPrecision命令。kv mx.kv.create() shape (2, 2) weight mx.nd.zeros(shape) kv.init(3, weight) kv.set_optimizer(mx.optimizer.SGD()) # 默认学习率 0.01 kv.push(3, mx.nd.ones(shape)) kv.pull(3, outweight) weight.asnumpy() # 全 -0.01SGD 已按梯度下降更新server 端如何接住这条命令见 python/mxnet/kvstore_server.pyKVStoreServer._controller中cmd_id 0时反序列化出 optimizer 并kvstore.set_optimizer(optimizer)。而 worker 侧可用命令编号表定义在 _get_kvstore_server_command_type命令编号kController下发优化器等控制命令0kSetMultiPrecision1kStopServer2kSyncMode3kSetGradientCompression4kSetProfilerParams54.2 _set_updater注册原始更新函数_set_updater(updater) 允许直接注册形如updater(key, input, stored)的更新函数Python 层经 _updater_wrapper 包成 ctypes 回调并同时注册 int key 与 str key 两个 C 回调MXKVStoreSetUpdaterEx。它只改本地 store多机场景必须用set_optimizer。文档示例def update(key, input, stored): print(update on key: %d % key) stored input * 2 kv._set_updater(update) kv.push(3, mx.nd.ones(shape))注意 key 类型一致性从 SetKeyType 可见一个 store 内 int key 与 str key 不可混用否则CHECK_EQ直接失败。4.3 优化器状态持久化save_optimizer_states(fname, dump_optimizerFalse) 与 load_optimizer_states(fname) 用于训练中 checkpoint 时保存/恢复优化器updater内部状态如动量项dump_optimizerTrue时连优化器本身学习率、weight decay 调度等一并保存。两者都断言self._updater is not None即分布式训练下不适用状态在 server 侧。五、梯度压缩set_gradient_compression 与 2bit 量化set_gradient_compression(compression_params) 为通信开启低比特梯度量化。目前仅支持type2bit可选参数threshold正浮点例如kv.set_gradient_compression({type: 2bit, threshold: 0.5})文档对 2bit 方案的描述是按threshold把梯度三值化大于阈值钳到阈值、小于负阈值钳到负阈值、其余置 0每个值用 2 bit 表示原梯度每 16 个 float 只需 1 个 float 即可表示被截断的残差保留在发送端累加进下一轮梯度发送从而在压缩通信的同时不丢信息。文档同时给出两类生效范围kvstore 为local此处实指设备间路径时压缩 GPU 到合并 GPU 的通信各 GPU 需额外存 residual 数组显存占用增加kvstore 为dist时只压缩 worker 到 server 的方向多 GPU 机器内 GPU-GPU / GPU-CPU 以及 server 到 workerpull方向不压缩且各 worker 的 CPU 内存因 residual 增加。实现上Python 层会先检查self.type含device或dist否则抛出 Gradient compression is not supported for this type of kvstore即纯localCPU 路径不支持C 侧 KVStoreDist::SetGradientCompression 在本地记录参数的同时由 rank 0 向所有 server 发送kSetGradientCompression命令并附带EncodeParams()编码串保证 server 同步启用。压缩算子本体位于 src/kvstore/gradient_compression.cc 与 gradient_compression.cu。六、分布式角色rank、num_workers、barrier 与 KVStoreServerKVStore提供三个与分布式拓扑相关的只读属性typestore 类型字符串如dist_syncrank本 worker 节点在[0, num_workers)内的编号底层调MXKVStoreGetRanknum_workersworker 节点数底层调MXKVStoreGetGroupSize。_barrier() 在所有 worker 间建立全局屏障典型用途是机器 0 先init完成再让所有 workerpull的时序保证_send_command_to_servers(head, body) 则向所有 server 节点下发命令server 会调用其KVStoreServer.controller执行命令全部执行完才返回。分布式 kvstore 的后端是 ps-liteKVStoreDist 继承KVStoreLocalworker 节点创建ps::KVWorker并ps::StartAsync启动相关环境变量从源码结构看MXNET_KVSTORE_BIGARRAY_BOUND大数组边界默认1000*1000kvstore_dist.h 第 58 行MXNET_KVSTORE_USETREE单机 device 通信是否使用二叉树拓扑见 KVStoreLocal 构造选CommDeviceTree或CommDevicelocal类型为CommCPU。server/scheduler 进程的入口是 python/mxnet/kvstore_server.py模块导入时执行 _init_kvstore_server_module若当前节点不是worker 节点MXKVStoreIsWorkerNode返回 0则以create(dist)建 store、KVStoreServer.run()进入服务循环对应 C APIMXKVStoreRunServer声明见 include/mxnet/c_api.h收到控制命令时由 controller 处理cmd 0 安装 worker 下发的 pickled 优化器。这与 C 侧 KVStoreDist::RunServer 的 ps-lite 启动流程一一对应。分布式路径的验证用例可参考 tests/nightly/dist_sync_kvstore.py、tests/nightly/dist_async_kvstore.py、tests/nightly/dist_device_sync_kvstore.py 与单机单测 tests/python/unittest/test_kvstore.py多卡 device 路径另有 tests/python/gpu/test_kvstore_gpu.py。七、C API 一览与工程要点小结MXKVStore*系列 C 接口集中声明在 include/mxnet/c_api.h与 Python 方法一一对应MXKVStoreCreate/MXKVStoreFree、MXKVStoreInit(Ex)、MXKVStorePush(Ex)、MXKVStorePullWithSparse(Ex)/MXKVStorePull(Ex)/MXKVStorePullRowSparse(Ex)、MXKVStorePushPull(Ex)、MXKVStoreSetUpdater(Ex)、MXKVStoreSetGradientCompression、MXKVStoreGetType/GetRank/GetGroupSize、MXKVStoreBarrier、MXKVStoreRunServer、MXKVStoreSendCommmandToServers等。工程使用要点汇总选型单机默认local多卡聚合放 GPU 用device多机训练用dist_sync/dist_device_sync等价批量为 n*b或dist_async真异步、到达顺序不保证顺序先init后push/pull跨 worker 需要确定时序时用_barrier()语义push 无序列化保证pull/pushpull 保证读到最新值且读取out会阻塞优化器单机用set_optimizer/_set_updater本地更新多机自动 pickled 下发 servercheckpoint 用save_optimizer_states/load_optimizer_states分布式不适用压缩{type: 2bit, threshold: ...}仅对 device/dist 生效代价是残差缓存带来的额外内存编译前提dist_*与nccl类型分别需要USE_DIST_KVSTORE1、USE_NCCL1编译选项未开启时mx.kv.create会直接致命报错。【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考