[RFC]: 层次化内存——tensor卸载/预取算子方案

### 背景与目标描述.

【需求背景&价值】
在大模型的训练与推理流程中，随着网络节点数量的增加，显存占用成为显著瓶颈。然而，单纯通过扩充设备侧显存来解决问题成本高昂。为此，我们引入层次化内存能力，在后端编译阶段对计算图执行序列进行深度分析，将部分权重与中间节点数据通过智能卸载与预加载机制进行调度。同时为了易用性，需提供相应的算子来进行预取和卸载等功能。
【应用场景】
需求来源：[RFC]: https://gitee.com/mindspore/mindspore/issues/ICXCQN
[RFC]: https://gitee.com/mindspore/mindspore/issues/ID3KWQ
[RFC]: https://gitee.com/mindspore/mindspore/issues/ID4EHA
![输入图片说明](https://foruda.gitee.com/images/1763361588175364996/4ebf3f5c_13795101.png "屏幕截图")
用于device/host/remote之间的数据搬移等操作
【功能描述】
在两种不同编址下提供两套不同的算子，分别表示卸载/加载和释放卡侧内存功能。
### 建议的方案.
1、`Parameter.to("CPU/Ascend") `和`Tensor.to("CPU/Ascend")`:

> 	接口介绍：将Parameter/Tensor复制到目的设备上，并返回新的Parameter/Tensor对象。
>
> 	接口依赖：CopyToHost(x, sync)/CopyToDevice(x, sync)
>
> ```python
> Parameter.to("CPU") <-------------> y = CopyToHost(x, sync)
> Tensor.to("CPU") 	<-------------> y = CopyToHost(x, sync)
> Parameter.to("Ascend") <-------------> y = CopyToDevice(x, sync)
> Tensor.to("Ascend")    <-------------> y = CopyToDevice(x, sync)
> ```

2、`Parameter.data.operator=()` 和`Tensor.data.operator=()`

> 接口介绍：将右侧data对象赋值回相应的Parameter/Tensor
>
> 接口依赖：SetData(input, value)
>
> ```python
> Parameter.data.operator=()  <--------------->  SetData(x, y)
> Tensor.data.operator=()  <--------------->  SetData(x, y)
> ```

3、`Parameter.data.delete_()`和`Tensor.data.delete_()`

> 接口介绍：删除Parameter和Tensor在device侧上的内存。
>
> 接口依赖：Free(x)
>
> ```python
> Parameter.data.delete_()  <--------------->  Free(x)
> Tensor.data.delete_()  <--------------->  Free(x)
> ```

4、`Parameter.to_(Device)`和`Tensor.to_(device)`

> 接口介绍：为易用性接口，不返回新对象
>
> 接口依赖：CopyToHost/CopyToDevice/FreeDevice/SetData

1、统一编址下：
    在对外接口上，新增两个算子CopyToDevice/CopyToHost和FreeDevice。其中算子功能如下所示：
    CopyToDevice：针对在device侧的tensor，调用该算子，返回原始tensor。在host侧的tensor，调用该算子会将算子拷贝至device侧。
    CopyToHost：针对在host侧的tensor，调用该算子，返回原始tensor。在device侧的tensor，调用该算子会将算子拷贝至device侧。
    Free：调用该算子主动释放tensor在device/host侧的内存。
    SetData：将value的data赋给input，后续共享data。
2、分层编址：
    UpdateToDevice：从remote侧加载tensor至device侧。
    UpdateToRemote：将Device侧数据卸载至remote侧。
    Detach：清除tensor在卡侧的物理内存，保留remote侧。

### 涉及到的对外API
| 编址方式 | 算子名      | 算子定义                                             | 算子功能 | 算子实现 | 特殊说明 |
|------|----------|--------------------------------------------------|------|------|-----------|
| 分层编址 | UpdateToDevice | x = ops.auto_generate.UpdateToDevice()(x, sync=False) | 将uvm从remote侧取值device侧。调用算子后，remote/device侧是否存在tensor副本取决于该tensor是否为多副本状态     |  调用hal接口 UpdateRemoteToDevice   |    实现为inplace算子，输入输出为同一个tensor    |
|   分层编址  |     UpdateToRemote     |           x = ops.auto_generate.UpdateToRemote()(x, sync=False)     |   将uvm从device侧取至remote侧。调用算子后，remote/device侧是否存在tensor副本取决于该tensor是否为多副本状态    |   调用hal接口 UpdateDeviceToRemote   |      实现为inplace算子，输入输出为同一个tensor     |
|  分层编址  |     Detach     |  x = ops.auto_generate.Detach()(x, sync=False)    |    释放tensor在device侧的物理内存。调用该接口后，在device侧不存在副本，仅保留remote侧副本  |  调用hal接口 DetachDevice    |  实现为inplace算子，输入输出为同一个tensor   |
|  统一编址 |  CopyToDevice  |     y = ops.auto_generate.CopyToDevice()(x, sync=False)    |  将数据x拷贝至device侧tensor y上。调用后返回一个与x shape/type/value相同的device侧的tensor  |  调用aclrtMemcpyAsync实现，依据输入的device_type选择不同copy方向  |           |
|   统一编址 |  CopyToHost  |     y = ops.auto_generate.CopyToHost()(x, sync=False)     |  将数据x拷贝至host侧tensor y上。调用后返回一个与x shape/type/value相同的host侧的tensor   |  调用aclrtMemcpyAsync/aclrtMemcpy实现  |   h2h拷贝采用同步接口 |
|  统一编址 |   Free  |     y = ops.auto_generate.Free()(x, sync=False)   | 清除tensor x在device侧的内存，后续在device侧使用x会报错      |  调用ClearDeviceMemory清楚物理内存  |       |
### 测试验证

算子功能验证：
1、调用CopyToDevice后，输出tensor的shape/type/value均与输入一致，device变成对应的卡侧
2、调用CopyToHost后，输出tensor的shape/type/value均与输入一致，device信息变成对新的host侧
3、device侧tensor调用Free后，卡侧内存释放，后续使用该tensor拦截报错处理
用例设计：
1、单CopyToHost/CopyToDevice测试。
```
@jit
def foo():
    x = Tensor(1.0)
    x = ops.auto_generate.CopyToHost()(x, True)
    return x

ret = foo()
assert ret.asnumpy() == 1.0
assert ret.device == "CPU"

@jit
def foo(x):
    x = ops.auto_generate.CopyToDevice()(x)
    return x

x = Tensor([1, 2, 3, 4])
ret = foo(x)
assert np.all(ret.asnumpy() == np.array((1, 2, 3, 4)))
assert ret.device == "Ascend:0"
```

2、CopyToHost/CopyToDevice

```
@jit
def foo(x):
    x = ops.auto_generate.CopyToHost()(x)
    x = ops.auto_generate.CopyToDevice()(x)
    return x

x = Tensor([1, 2, 3, 4])
ret = foo(x)
assert np.all(ret.asnumpy() == np.array((1, 2, 3, 4)))
assert ret.device == "Ascend:0"
```

3、SetData用例设计

```
@jit
def foo1(x, a):
    ops.auto_generate.SetData()(x, a)
    return x

x = Tensor([1, 2, 3, 4])
a = Tensor([5, 6, 7, 8])
ret = foo1(x, a)
assert np.all(ret.asnumpy() == np.array((5, 6, 7, 8)))

@jit
def foo2(x, a):
    ops.auto_generate.SetData()(x, a)
    a.add_(1)
    return x

x = Tensor([1, 2, 3, 4])
a = Tensor([5, 6, 7, 8])
ret = foo2(x, a)
assert np.all(ret.asnumpy() == np.array((6, 7, 8, 9)))

@jit
def foo3(x, a):
    ops.auto_generate.SetData()(x, a)
    a.add_(1)
    x.add_(1)
    return a

x = Tensor([1, 2, 3, 4])
a = Tensor([5, 6, 7, 8])
ret = foo3(x, a)
assert np.all(ret.asnumpy() == np.array((7, 8, 9, 10)))
```

### 期望的反馈时间.

### CC List.
 @baochong  @zh_qh  @liangzhibo  @ginfung  @limingqi107 @zyli2020

### 其他补充信息.

### Before submitting a new issue...

- [x] Make sure you already searched for previous [RFCs](https://gitee.com/mindspore/mindspore/issues?q=is%3Aall+label%3ARFC+sort%3Arecently-updated).

GVP MindSpore/mindspore
关闭

内容风险标识

评论 (6)

GVPMindSpore/mindspore关闭

内容风险标识