[AutoParallel] fix GPT embedding placements

[waliwali777]

Before submitting

• Lint code. If there are lint issues, please format the code first.

# Install and register `pre-commit` in the project folder
pip install pre-commit && pre-commit install

# Process previous code files separately
pre-commit run --file XXXX.py

• Add test cases into tests folder. If there are codecov issues, please add tests cases first.

PR types

Others

PR changes

Others

Description

GPT 动半下，修改 placemetns 以提高性能
1、修改 embedding 和 lmhead 层 weight 的 placemetns 为列切
2、修改 GPTEmbeddingsAuto 层输出的 placemetns 为 replicate

• 修改前：
  GPTEmbeddingsAuto 输出为 [Replicate(), Reshard(2)]，导致每次进入 encoder 层的状态都为 [Replicate(), Reshard(2)]，在进行 residual + dropout 计算时，有 [Replicate(), Reshard(2)] -> [Replicate(), Replicate()]，这在前向和反向中会引入 allgather 通信
  GPT 单机8卡 H卡 最优配置 20layer 吞吐(token/card/s)：7902.79

• 修改后：
  GPTEmbeddingsAuto 输出为 [Replicate(), Replicate()]， encoder 层输入的状态也为 replicate，在进行 residual + dropout 计算时，可以直接进行 replicate 计算，不引入额外通信
  GPT 单机8卡 H卡 最优配置 20layer 吞吐(token/card/s)：8305.75(+402.96,+5.1%)

[paddle-bot]

Thanks for your contribution!

[liym27]

LGTM