变量管理Variable Manager【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geGE 图引擎提供了一套完整的变量生命周期管理机制覆盖变量的注册、内存分配、格式转换优化、逻辑地址映射、离线序列化/反序列化、运行时地址解析等全流程。该机制支撑了训练场景下多个子图共享同一变量、推理场景下的权重加载与复用、以及 OM 模型离线编译与部署等核心能力。整体架构GE 的变量管理体系采用分层设计由编译期和运行时两个层次组成核心概念变量类型GE 管理以下几类需要持久化内存的算子节点统称为变量类型说明场景VARIABLE训练参数变量可被多个子图共享读写训练场景CONSTANTOP常量节点编译期确定值存储在变量内存区域推理/训练FILECONSTANT外置权重常量权重数据存储在外部文件中运行时按需加载推理大模型CONSTPLACEHOLDER常量占位符支持外部内存注入推理外部权重管理变量的唯一标识变量的唯一键由变量名 格式 数据类型组合而成定义于VarResource::VarKey()var_key batch_var_name _ format _ data_type其中batch_var_name是支持多 batch 训练场景下的变量名映射同一变量在不同 batch 分支中可能使用不同名称但底层共享同一份内存通过batch_var_name_map_建立映射关系。变量管理器VarManager类层次结构VarManagerPool全局单例 └─ mapsession_id, shared_ptrVarManager └─ VarManagerSession 级 ├─ VarResource变量资源地址表、描述表、转换路径 ├─ mapMemType, MemResource内存资源HBM / RDMA / Host └─ MemoryManager物理内存分配器源文件位置base/graph/manager/graph_var_manager.h— VarManager、VarResource、MemResource 定义base/graph/manager/graph_var_manager.cc— 核心实现VarManagerPoolVarManagerPool是全局单例维护session_id → VarManager的映射关系。每个训练/推理 Session 拥有独立的VarManager实例保证 Session 间变量隔离。VarResourceVarResource是变量信息的核心存储维护以下关键数据结构成员类型用途var_addr_mgr_map_mapvar_key, VarAddrMgr变量名格式 → 地址信息映射cur_var_tensor_desc_map_mapvar_name, GeTensorDesc变量当前最新的 Tensor 描述var_offset_map_maplogic_addr, MemType逻辑地址 → 内存类型映射var_dev_addr_mgr_map_maplogic_addr, VarDevAddrMgr逻辑地址 → 设备地址映射var_to_trans_road_mapvar_name, VarTransRoad变量格式转换路径var_names_to_changed_graph_id_mapvar_name, graph_id变量所属的变更图 IDvar_names_to_allocated_graph_id_mapvar_name, graph_id变量首次分配内存的图 IDfile_constant_var_map_mapfile_pathoffset, var_keyFileConstant 文件路径 → 变量键映射device_id_to_var_dev_addr_mgr_map_mapdevice_id, VarDevAddrMgr多设备场景下的设备地址映射batch_var_name_map_mapbatch_var_name, key_name多 batch 变量名映射MemResourceMemResource管理变量内存的分配按内存类型分为三种实现类型类分配策略RT_MEMORY_HBM设备内存HbmMemResource偏移量递增分配512 字节对齐额外保留 1024 字节 guard spaceRT_MEMORY_RDMA_HBMRDMA 内存RdmaMemResource从 RDMA 内存池分配RT_MEMORY_HOST主机内存HostMemResource从 Host 内存池分配HBM 内存的分配逻辑每个变量分配时实际占用空间 对齐后大小 512 字节对齐 1024 字节用于 inner_offset 定位确保每个变量的内存区间互不重叠且有安全间距。变量内存分配编译期分配流程变量内存分配在图编译阶段完成由VarMemAssignUtil驱动源文件位置base/graph/build/memory/var_mem_assign_util.cc内存复用机制VarManager::AssignVarMem()内部实现了多级复用策略格式匹配复用如果变量名已存在于cur_var_tensor_desc_map_且当前格式与已有格式匹配则直接复用已有地址。权重去重复用对于CONSTANTOP类型通过GetReuseAddr()比较权重数据的二进制内容memcmp相同权重共享同一内存地址。FileConstant 路径复用对于FILECONSTANT类型通过file_constant_var_map_按文件路径 偏移量去重同一权重文件的不同算子共享内存。大小兼容复用如果已有变量内存足够容纳新的 Tensor 描述tensor_desc_size cur_tensor_desc_size则复用原有内存。变量逻辑地址GE 采用逻辑地址机制将编译期地址与运行时物理地址解耦。核心常量定义如下常量值含义kVarMemoryLogicBase128 GB变量逻辑地址起始基址逻辑地址的作用编译期生成的 OM 模型中变量引用的是逻辑地址即var_mem_logic_base_ offset。运行时加载模型时根据实际分配的物理内存地址通过GetVarMemoryAddr()将逻辑地址转换为设备物理地址。这种设计使得 OM 模型可以在不同设备上加载执行而无需重新编译。离线场景的特殊处理离线编译atc与运行可能在不同 SoC 版本上执行。为此离线场景将变量逻辑基址固定为 128 GBkVarMemoryLogicBase避免因设备内存布局差异导致地址冲突。运行时地址解析运行时地址解析由VarManager::GetVarMemoryAddr()完成支持以下场景RDMA 内存直接返回逻辑地址RDMA 内存已预分配固定地址。外部变量内存如果通过SetExternalVar()注入了外部内存区域则通过external_var_addr_ (logic_addr - var_mem_logic_base_)计算物理地址。自动分配通过GetAutoMallocVarAddr()实现延迟分配。首次访问变量时自动分配物理内存并缓存到VarDevAddrMgr::dev_addr后续访问直接返回缓存地址。大页内存支持 1GB 大页IsVariableUse1gHugePage通过ExpandableMemoryAllocator管理可扩展内存。变量加速 Pass概述变量加速Variable Acceleration是 GE 编译期的一项重要优化通过VariableOpPass实现。其核心思想是当变量的所有下游算子都以相同的格式消费数据时将变量的存储格式直接改为该目标格式消除运行时的格式转换开销。触发条件变量加速通过选项ge.exec.variable_acc控制开关默认开启。当启用多图并行编译ge.AllowMultiGraphParallelCompile1时自动关闭因为多图并行场景下变量格式变更可能引起冲突。源文件位置compiler/graph/manager/graph_manager.cc工作流程融合判定逻辑FusionIfNeed()采用迭代方式逐层融合变量与下游 TransData/Cast/ReFormat 等转换节点一致性检查CheckSameAndTransOp确认变量的所有下游转换算子输出相同的格式, 数据类型, Shape组合。如果存在不一致跳过该变量。VarRef 合法性检查CheckVariableRefLegally如果存在 VarRef 节点变量回写则检查 Variable 侧的转换路径与 VarRef 侧的转换路径是否互逆。只有互逆时才能安全融合。格式连续性补齐如果变量输出格式与第一个转换节点的输入格式不连续则在转换路径头部插入一个 ReFormat 节点补齐。执行融合DealFusion移除被融合的转换节点将变量的输出直接连接到转换节点的下游。更新描述更新 Variable 和 VarRef 节点的输出描述记录转换路径到 VarManager。变更次数限制GraphRebuildStateCtrl限制每个变量的格式变更次数最多为 1 次kMaxVarChangeTimes_ 1防止变量在多个图的编译过程中反复变更格式导致振荡。首次编译时变量格式被优化为目标格式后后续图的编译将保持该格式不变。支持的转换类型变量加速支持以下转换算子类型类型条件TransData源/目标格式均支持通过formats::IsTransFormatSupport校验TransDataD同 TransDataCast数据类型转换支持通过formats::IsTransDataTypeSupport校验ReFormat无条件支持Reshape无条件支持Shape 变更数据不变SqueezeV2/UnSqueezeV2无条件支持相关 Pass变量加速涉及多个协同工作的 PassPass阶段作用VariablePrepareOpPassO3建立 Variable-VarRef 关系为可写变量Assign/AssignAdd/AssignSub创建 VariableRef 节点VariableOpPassO3OptimizeStage1_1变量格式融合优化VariableRefDeleteOpPass后处理清理不再需要的 VariableRef 节点VariableRefUselessControlOutDeletePass后处理删除 VariableRef 上多余的控制边VarIsInitializedOpPassO0将 VarIsInitializedOp 替换为常量根据变量是否已初始化源文件位置compiler/graph/passes/variable_optimize/变量准备 PassVariablePrepareOpPass功能VariablePrepareOpPass在图优化的早期阶段运行负责建立 Variable 与其引用者之间的 VarRef 关系。这是变量管理和加速的基础——只有正确识别了哪些算子会修改变量后续优化才能安全进行。工作原理识别 Ref 算子遍历图中的所有节点通过输入/输出名称匹配和RefPortIndex属性识别具有引用语义的算子如 Assign、AssignAdd、AssignSub 等。追踪写入链从 Variable 节点出发沿数据边追踪经过的 Ref 算子链找到最后一个写入节点。插入 VarRef在最后一个 Ref 算子的输出后插入一个VariableRef节点类型与原 Variable 相同并设置REF_VAR_SRC_VAR_NAME属性指向原始变量名。控制边保证添加控制边确保 VariableRef 在后续算子之前执行保证变量值的一致性。对于没有对应 Ref 输出的算子如 RefSwitch会额外插入RefIdentity节点作为桥梁。运行时变量管理V1 运行时ModelRtVarManagerModelRtVarManager是 V1 运行时runtime/v1/的变量管理入口每个 Session 对应一个实例。源文件位置inc/framework/runtime/model_rt_var_manager.h、runtime/v1/common/runtime/model_rt_var_manager.cc初始化流程调用Init()设置设备 ID、变量逻辑基址、总变量大小等参数。如果 VarManager 尚未初始化则初始化 VarManager 并配置内存参数。创建ExpandableMemoryAllocator支持按需扩展的变量内存分配。变量恢复RestoreDeviceVariables在模型加载时将变量信息从图中恢复到 VarManager 中。对于已存在的变量直接复用对于新变量调用RestoreVarMem()恢复其内存信息。变量查询GetVarShapeAndMemory根据变量名返回变量的 Shape 和设备内存地址供运行时算子使用。V2 运行时SplitVariable KernelV2 运行时runtime/v2/采用更简洁的设计。变量通过SplitVariableKernel 在执行图Execution Graph中被解析。源文件位置runtime/v2/kernel/common_kernel_impl/variable.ccSplitVariable 的工作方式从RtSession获取RtVarManager运行时变量管理器接口。根据变量 ID字符串查询变量的 Shape 和内存地址。将结果写入输出 Tensor供后续算子使用。V2 设计的核心理念是算子只应按 IR 语义工作不应该感知 Session/Device 等上下文信息。变量的上下文解析应在模型加载阶段完成运行时算子只通过变量 ID 获取所需的地址和 Shape。变量转换器Variable Converter在 V2 的 Lowering 阶段Variable节点通过LoweringVariable转换器被转换为SplitVariableKernel。源文件位置runtime/v2/engine/gelocal/variable_converter.h离线变量管理序列化与反序列化VarManager 支持将完整的变量管理信息序列化为 Protobuf 格式用于 OM 模型的离线编译和部署。Protobuf 定义位于graph_metadef/proto/var_manager.proto核心消息包括消息用途VarManagerInfoVarManager 完整信息版本、Session ID、内存配置、变量资源VarResourceInfo变量资源信息地址映射表、描述表、转换路径、广播信息VarDescInfo变量描述信息当前描述、暂存描述、转换路径VarMatchInfo编译前后的变量描述匹配信息用于模型缓存MemResourceInfo内存资源信息总大小、已用大小序列化流程VarManagerToSerial记录 VarManager 的版本、Session ID、设备 ID、内存配置等元信息。序列化 VarResource 中的地址映射表、描述表、转换路径、广播信息。序列化各内存类型的 MemResource 使用量。反序列化流程VarManagerToDeserial获取当前设备 IDaclrtGetDevice。恢复内存配置参数。从 Protobuf 数据中恢复 VarResource 的所有映射表。恢复 MemResource 的已分配大小。模型缓存与变量在模型缓存Model Cache场景下变量的描述信息格式、数据类型、Shape用于判断缓存是否有效。源文件位置compiler/graph/build/model_cache.h、compiler/graph/build/model_cache.cc编译前描述var_desc_before_compile_记录编译前的变量描述。编译后描述通过VarMatchInfo记录编译后的变量描述。缓存校验加载缓存模型时比较变量描述是否发生变化。如果变化且变量格式变更次数已达上限kMaxVarChangeTimes_ 1则缓存失效需要重新编译。FileConstant 外置权重FileConstant是 GE 提供的外置权重机制将大模型权重数据存储在独立文件中运行时按需加载到设备内存显著降低 OM 模型文件体积。源文件位置base/common/file_constant_utils/file_constant_utils.h权重文件路径获取支持三种方式IR 属性file_path直接指定权重文件路径。IR 属性file_id Optionge.exec.value_bins通过文件 ID 映射到文件路径。私有属性location由 Parser 模块或外置权重功能自动设置。运行时加载FileConstantKernel在首次执行时从权重文件中读取数据并拷贝到设备内存。后续执行时直接使用已加载的内存地址跳过文件读取。同一权重文件的不同算子通过file_constant_var_map_共享内存。外置权重导出通过ge.externalWeight选项控制支持两种模式单独导出1每个权重生成独立文件。合并导出2所有权重合并到同一文件通过 offset 区分。变量初始化与就绪状态变量初始化值变量支持通过_init_value属性设置初始值。在运行时首次分配变量物理内存时GetAutoMallocVarAddr如果检测到该属性会自动将初始值从 Host 拷贝到 Device 内存。源文件位置base/graph/manager/graph_var_manager.ccInitVarIfHasInitValue变量就绪状态VarResource维护变量的就绪状态var_is_instance_按设备 ID 和变量键记录。运行时通过SetVarIsReady()/IsVarReady()查询变量是否已在特定设备上完成初始化避免重复加载。VarIsInitializedOp 处理VarIsInitializedOpPass在编译期将VarIsInitializedOp/IsVariableInitialized算子替换为布尔常量如果变量已在 VarManager 中注册IsVarExist替换为Constant(true)。如果变量尚未注册根据图中的 Assign 操作追踪变量是否已被初始化。替换后的常量值可以被后续的常量折叠 Pass 优化消除运行时开销。内存管理策略静态/动态内存策略GE 通过GE_USE_STATIC_MEMORY环境变量和STATIC_MEMORY_POLICY选项控制内存分配策略策略值含义0默认策略1kStaticMemory静态内存策略2kExtendSizeType扩展大小策略静态图和动态图内存复用3kDynamicExpandable动态可扩展策略4kDynamicAndStaticExpandable静态动态可扩展策略源文件位置base/graph/manager/graph_var_manager.hIsGeUseExtendSizeMemory大页支持变量内存支持使用 1GB 大页Huge Page通过ge.variableUse1gHugePage选项控制0不使用大页默认1仅使用 1GB 大页2优先使用 1GB 大页不够时回退到普通页使用大页可以减少 TLB Miss提升大规模变量的访问性能。广播变量在分布式训练场景中变量需要通过HcomBroadcast或HvdCallbackBroadcast算子进行跨卡同步。VarResource维护了广播信息VarBroadCastInfo记录每个变量的广播输入/输出偏移和大小确保广播操作正确完成。源文件位置base/graph/manager/graph_var_manager.ccSaveBroadCastInfo多设备支持VarManager 支持多设备场景。VarResource通过device_id_to_var_dev_addr_mgr_map_维护每个设备上的变量物理地址映射。运行时通过UpdateDevVarMgrInfo()将编译期的变量信息同步到指定设备然后通过GetVarMgrInfo()按设备 ID 查询对应的物理地址。变量去重加载CheckAndSetVarLoaded()方法检查同一设备上是否已有其他变量占用了相同的内存偏移避免重复加载相同权重的常量。关键 API 汇总VarManager 主要接口接口功能Init()初始化 VarManager版本、Session、设备、JobAssignVarMem()分配变量内存含复用检查RestoreVarMem()恢复变量内存离线加载场景SetVarAddr()设置变量地址GetVarAddr()获取变量逻辑地址GetVarMemoryAddr()逻辑地址 → 物理地址转换SetTransRoad()设置变量格式转换路径GetTransRoad()获取变量格式转换路径RenewCurVarDesc()更新变量当前描述格式/数据类型变更后VarManagerToSerial()序列化到 ProtobufVarManagerToDeserial()从 Protobuf 反序列化FreeVarMemory()释放所有变量内存IsVarExist()检查变量是否存在CheckAndSetVarLoaded()检查并标记变量已加载去重VarResource 主要接口接口功能GetVarAddr()根据变量名和描述获取地址GetReuseAddr()查找可复用的变量地址权重去重SetVarAddr()注册变量地址SaveVarAddr()保存变量地址含逻辑地址计算GetCurVarDesc()获取变量当前 Tensor 描述RenewCurVarDesc()更新变量描述CheckLogicAddrValid()校验逻辑地址合法性SetVarMgrDevAddr()设置设备物理地址文件索引模块文件路径变量管理器核心base/graph/manager/graph_var_manager.h、.cc变量内存分配base/graph/build/memory/var_mem_assign_util.h、.cc变量加速 Passcompiler/graph/passes/variable_optimize/variable_op_pass.h、.cc变量准备 Passcompiler/graph/passes/variable_optimize/variable_prepare_op_pass.h、.ccVarRef 删除 Passcompiler/graph/passes/variable_optimize/variable_ref_delete_op_pass.h、.ccVarRef 控制边清理compiler/graph/passes/variable_optimize/variable_ref_useless_control_out_delete_pass.h、.cc变量初始化检查 Passcompiler/graph/passes/variable_optimize/var_is_initialized_op_pass.h、.cc变量加速状态控制compiler/graph/manager/util/graph_rebuild_state_ctrl.h模型缓存compiler/graph/build/model_cache.h、.cc图管理器Pass 注册compiler/graph/manager/graph_manager.ccV1 运行时变量管理inc/framework/runtime/model_rt_var_manager.h、runtime/v1/common/runtime/model_rt_var_manager.ccV2 运行时变量接口inc/framework/runtime/rt_var_manager.hV2 变量 Kernelruntime/v2/kernel/common_kernel_impl/variable.h、.ccV2 变量转换器runtime/v2/engine/gelocal/variable_converter.h、.cc外置权重工具base/common/file_constant_utils/file_constant_utils.h、.ccProtobuf 定义graph_metadef/proto/var_manager.proto【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考