片上内存显存侧动态扩容模式

训练流程介绍

本节介绍如何使用动态扩容模式进行训练，整体操作流程请参见图1。

整理训练流程分为以下部分，整体流程示例代码请参见示例代码。

• 适配模型
• 启动训练

适配模型

关键步骤操作参考如下。

1. 初始化框架。

   调用init接口，设置参数use_dynamic_expansion = True表示启用动态扩容功能，该参数默认为“False”。

2. 稀疏优化器导入。

   调用mx_rec.optimizers包中对应优化器的create_hash_optimizer_by_address接口来创建稀疏表sparse_optimizer。具体可用优化器参考如下。

   • SGDByAddr
   • LazyAdamByAddress
3. 获取嵌入表示结果（emb）和映射地址（addr）。

   使用tf.get_collection("ASCEND_SPARSE_LOOKUP_LOCAL_EMB")接口获取训练用的嵌入表示结果，使用tf.get_collection("ASCEND_SPARSE_LOOKUP_ID_OFFSET")接口获取训练用的映射地址。

4. 反向梯度计算。

   使用tf.gradients(loss, emb)接口对3获取的嵌入表示结果求导，得到梯度（grad）。

5. 反向稀疏表更新。

   使用2.sparse优化器导入。创建的sparse_optimizer.apply_gradients([grad, addr])接口对映射地址地址对应位置的稀疏表进行更新。

启动训练

1. 启动模型训练。
2. 在模型训练完成后调用terminate_config_initializer释放资源。

示例代码

1. 初始化框架。

   use_dynamic_expansion = bool(int(os.getenv("USE_DYNAMIC_EXPANSION", 0)))
   init(use_mpi, train_steps=args.train_steps, eval_steps=args.eval_steps， 
   use_dynamic_expansion=use_dynamic_expansion)

2. 稀疏优化器导入。

   def get_dense_and_sparse_optimizer(cfg):
       dense_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=cfg.learning_rate)
       use_dynamic_expansion = get_use_dynamic_expansion()
       if use_dynamic_expansion:
           sparse_optimizer = create_hash_optimizer_by_address(learning_rate=cfg.learning_rate)
           logging.info("optimizer lazy_adam_by_addr")
       else:
           sparse_optimizer = create_hash_optimizer(learning_rate=cfg.learning_rate)
           logging.info("optimizer lazy_adam")
       return dense_optimizer, sparse_optimizer

3. 获取嵌入表示结果和映射地。

   train_emb_list = tf.compat.v1.get_collection(ASCEND_SPARSE_LOOKUP_LOCAL_EMB)
   train_address_list = tf.compat.v1.get_collection(ASCEND_SPARSE_LOOKUP_ID_OFFSET)

4. 反向梯度计算。

   local_grads = tf.gradients(loss, train_emb_list)  # local_embedding

5. 反向稀疏表更新。

   grads_and_vars = [(grad, address) for grad, address in zip(local_grads, train_address_list)]
   train_ops.append(sparse_optimizer.apply_gradients(grads_and_vars))

   

   • 调用sparse_optimizer.apply_gradients(grads_and_vars)更新梯度时，若使用的vars（如address）是tensor而非variable，需要保证vars的维度和grads的第一个维度相等。
   • train_address_list地址列表需要是有效合法的，需通过3. 获取映射地址获取。若使用非法地址，运行时会抛出AICore Error等错误。