PR Category

User Experience

PR Types

Bug fixes

Description

pcard-85872
when use AMP(BF16 dtype, FP16 is fine) training in dygraph with pir will report an error: ValueError: (InvalidArgument) The type of data we are trying to retrieve (float32) does not match the type of data (bfloat16) currently contained in the container.

The reason is that when with pir(FLAGS_enable_pir_api=True) the method of _is_dtype_fp16_or_bf16 in the optimizer base class Optimizer incorrectly judges data type of BF16 (FP16 is fine). so when data type is actually BF16 , but it judges that not.

this further leads to issues of judgment logic of the method _add_moments_pows、_create_accumulators and _append_optimize_op in the optimizer. _add_moments_pows creates tensors of moment1, moment2, beta1_pow, and beta2_pow with dtype of BF16 (which should be FP32), which leads to an type mismatch error above when executing beta1_pow.data<MPDType() in the kernel of adam/adamw (adamw_kernel.cu).