[Hybrid Performance] Move the cast op of AMP which cast fp32 param to fp16 param to the optimizer

[wangxicoding]

PR types

Performance optimization

PR changes

Others

Describe

Move the cast op of AMP which cast fp32 param to fp16 param to the optimizer.
Usage, set optimize_cast to True:

import paddle.distributed.fleet as fleet

strategy = fleet.DistributedStrategy()
strategy.sharding = True
strategy.sharding_configs = {
    "sharding_degree": 1,
    "mp_degree": 1,
    "pp_degree": 2,
    "dp_degree": 2,
    "optimize_cast: True,
}
strategy.pipeline = True
strategy.pipeline_configs = {
    "schedule_mode": "1F1B",
    "micro_batch_size": 2,
    "accumulate_steps": 4,
}
strategy.amp = True

Test

Test in 1node*8cards 32G V100, with Ernie3.0 model.

Model config, Ernie 3.0 base:

	value
hidden size	768
num attention heads	12
num hidden layers	15
num sharing layers	12
branch hidden size	256
branch num attention heads	4

batch size configs:

micro bsz	global bsz
2	256

Performance:

hybrid_config	optimize_cast	throughput(tokens/s)	improve
8pp	false	35849
8pp	true	37747	5.29%
2mp+2pp+2dp	false	41450
2mp+2pp+2dp	true	44543	7.46%

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[sandyhouse]

LGTM

[gongweibao]

需要验证精度对齐以及给出NPU上的性能提升数据?

[JZ-LIANG]

LGTM

[gongweibao]

LGTM