关联 Issue (Related Issue)

https://gitcode.com/mindspore/mindformers/issues/2213

修改描述 (Description)

修复内容

修复在 hyper-parallel 模式下 vocab embedding 的 Tensor Parallel (TP) 报错问题。

变更说明

文件: mindformers/pynative/base_models/common/embeddings/vocab_embedding.py

主要修改:

1. 将 mint.nn.functional.embedding 替换为 mint.gather
   原使用 mint.nn.functional.embedding 会导致 hyper-parallel 模式下报错，现改用 mint.gather 实现相同功能。

2. 新增 mint.tile 操作
   用于扩展输入维度以适配 gather 操作的需求。

3. 重构 forward 逻辑

   • 使用 tile 将 input_ 从 (batch_size, seq_len) 扩展为 (batch_size, seq_len, embedding_dim)
   • 使用 gather 从 weight 中按指定维度收集 embedding
   • 最终 reshape 为 (-1, seq_len, embedding_dim) 保持输出形状一致

影响范围

• 仅影响 VocabEmbedding 类的 forward 实现
• 修复了 hyper-parallel TP 模式下的兼容性问题
• 功能行为保持不变，仅更换了底层实现方式

修改类型 (Type of Change)

• Bug 修复 (Bug fix)
• 新特性 (New feature)
• 文档更新 (Documentation update)
• 代码重构 (Refactoring)
• 性能优化 (Performance improvement)
• 其他 (Other)

测试结果 (Test Results)

修改之前，embedding结果错误：

修改之后，开启tp后结果正确：

检查清单 (Checklist)

• 我已自验过该功能/修复 (Self-checked)
• 我已通过了本地的单元测试 (Passed local UTs)
• (如适用) 我已更新了对应文档 (Updated documentation)
• (如适用) 我已添加了新的测试用例 (Added new tests)