DTensor to() / float() 方法统一至核心层设计

1. 背景

1.1 现状

HyperParallel 的 DTensorBase 分别在两个平台实现中定义：

文件	基类	备注
platform/torch/dtensor.py → DTensorBase(Tensor)	继承 torch.Tensor	已有 to() / float()
platform/mindspore/dtensor.py → DTensorBase(Tensor)	继承 mindspore.Tensor	无 to() / float()
core/dtensor/dtensor.py → DTensor(DTensorBase)	核心层，用户实际使用的类	原先无 to() / float()

当前 to() 仅在 Torch 后端的 DTensorBase 中实现，MindSpore 后端缺失。若在 MindSpore 的 DTensorBase 中补齐 to()，会产生两份逻辑完全相同的代码，违反 DRY 原则。

1.2 核心困难

MindSpore C 描述符限制：MindSpore 的 Tensor 是 C 扩展类型，其 to 方法在元类创建阶段被包装为 C-level 描述符，该描述符强制 isinstance(self, Tensor) 类型检查。因此：

• 无法在 MindSpore 后端编写不依赖硬件的 UT 来测试 DTensorBase.to()
• DTensorBase.__dict__["to"] 会抛出 KeyError，因为元类已将其消费
• 任何试图绕过类型检查的 mock 对象（非 Tensor 子类）都会触发 descriptor 'to' for 'TensorPy' objects doesn't apply to ... 错误

而 PyTorch 后端的 torch.Tensor.to 是纯 Python 实现，不存在此问题，因此 Torch 后端可以写 UT。

1.3 关键发现

DTensor（定义在 core/dtensor/dtensor.py）是一个纯 Python 类（继承自 DTensorBase，但自身没有 C 扩展限制）。因此 DTensor.__dict__["to"] 可以正常获取到绑定的 Python 函数——这为跨平台 UT 提供了可行路径。

2. 目标与非目标

2.1 目标

• 将 to() 和 float() 方法统一放在 core/dtensor/dtensor.py 的 DTensor 类中，消除平台间重复代码
• 抽取 _alias_placements() 辅助方法，统一"获取有效 placements"的逻辑
• 提供硬件无关的 UT，在 CI 中无需 GPU/NPU 即可验证 to() / float() 的正确性
• 保持已有 Torch ST 测试不受影响

2.2 非目标

• 不迁移 detach() / detach_() / type() 等方法（它们逻辑相同但涉及平台特定行为，可作为后续工作）
• 不添加 MindSpore ST 测试（需要真实硬件环境，不在本次范围内）
• 不改变 DTensorBase.__new__ / __init_data__ 等构造逻辑

3. 设计方案

3.1 方法位置迁移

迁移前：
  platform/torch/dtensor.py DTensorBase.to()     ← Torch 独有
  platform/mindspore/dtensor.py DTensorBase      ← 无 to()

迁移后：
  core/dtensor/dtensor.py DTensor.to()           ← 统一，两平台共享
  core/dtensor/dtensor.py DTensor.float()        ← 统一，两平台共享
  platform/torch/dtensor.py DTensorBase          ← 删除 to()/float()

3.2 实现细节

_alias_placements() 辅助方法：

def _alias_placements(self) -> Sequence[Placement]:
    """Return alias_placements from layout, falling back to _placements."""
    if hasattr(self, '_layout') and self._layout:
        return self._layout.alias_placements
    return self._placements

此方法将 Torch 后端 DTensorBase 中反复出现的模式 self._layout.alias_placements if hasattr(self, '_layout') and self._layout else self._placements 抽取为可复用的方法。

to() 方法：

def to(self, *args, **kwargs):
    """Move the DTensor to a different device or dtype."""
    new_local = self._local_tensor.to(*args, **kwargs)
    return self.__class__(new_local, device_mesh=self._device_mesh,
                          placements=self._alias_placements())

float() 方法：

def float(self):
    """Convert the DTensor to float dtype."""
    new_local = self._local_tensor.float()
    return self.__class__(new_local, device_mesh=self._device_mesh,
                          placements=self._alias_placements())

3.3 方法解析顺序（MRO）分析

DTensor.__mro__:
  DTensor → DTensorBase(Tensor) → Tensor → object

• 迁移后，dt.to(...) 先在 DTensor.__dict__ 中找到 to，不再查找 DTensorBase 或 Tensor 的 to
• MindSpore 后端：直接返回 DTensor 的 Python 实现，绕过 C 描述符限制
• Torch 后端：行为与原先 DTensorBase.to() 完全一致，因为逻辑相同

3.4 UT 测试策略

核心技巧：通过 DTensor.__dict__["to"] 提取纯 Python 函数，在非 Tensor 对象上调用：

_to_fn = DTensor.__dict__["to"]

# _Recorder 是一个普通 Python 类（不继承 DTensor/Tensor）
# 用 Mock() 作为 _local_tensor，避免硬件访问
fake = _Recorder.__new__(_Recorder)
fake._local_tensor = Mock()  # mock，不触碰真实硬件
fake._layout = SimpleNamespace(alias_placements=[Shard(0)])

_to_fn(fake, "float16_arg")  # 直接调用 Python 函数

测试用例：

测试	验证点
test_to_delegates_dtype_conversion_to_local_tensor	to() 正确转发参数给 _local_tensor.to()
test_to_preserves_device_mesh	device_mesh 不变传递
test_to_uses_alias_placements_when_layout_exists	有 _layout 时使用 alias_placements
test_to_falls_back_to_placements_when_no_layout	_layout 为 None 时回退到 _placements
test_float_delegates_to_local_tensor	float() 正确调用 _local_tensor.float()

4. 影响范围

4.1 受影响文件

文件	变更
core/dtensor/dtensor.py	新增 _alias_placements()、to()、float()
platform/torch/dtensor.py	删除 to()、float()（逻辑已上移至核心层）
tests/ut/core/dtensor/test_dtensor_to.py	新增 5 个 UT 测试

4.2 不受影响

项目	原因
Torch ST 测试	tests/torch/dtensor/dtensor_api.py 不变，DTensor.to() 行为一致
MindSpore DTensorBase	无需修改，自然继承 DTensor.to()
分布式算子注册	to() / float() 不经过 op dispatch
已有 UT / ST	不涉及 _alias_placements / to / float 的测试不受影响

5. 验证方案

层级	内容	环境
UT	pytest tests/ut/core/dtensor/test_dtensor_to.py -v（5 个用例）	无硬件要求
Torch ST	pytest tests/torch/dtensor/test_dtensor_api.py -v（已有用例）	8 卡 GPU/NPU
MindSpore ST	无（不在本次范围）	—