跳转至

精度对齐(Precision Alignment)

DTorch 中调用 LibTorch,而 PyTorch 则通过 python 调用 LibTorch,两者运行的 Kernel 一致,输出应当时一致的。但也有一些情况,使得两者的输出不一致。

View + Linear

在 NIVIDA 3090 上遇到 View + Linear 可能会导致输出不一致的情况。代码如下,当给 PyTorch 和 DTorch 都添加 hidden_states = hidden_states.contiguous() 后,两种调用的 Kernel 一致。

# https://github.com/huggingface/diffusers/blob/v0.34.0/src/diffusers/models/attention_processor.py#L1488
#
# 在 slice 之后,hidden_states 的以下属性均一致:
#   hidden_states.stride()
#   hidden_states.untyped_storage().nbytes()
#   hidden_states.numel() * hidden_states.element_size()
# 仅发现 _is_view() 属性不一致。
#   PyTorch: hidden_states._is_view() = False
#   DTorch: hidden_states._is_view() = True
#
# 之后在 Linear 中,PyTorch 调用了三个 Kernel,DTorch 则调用两个 Kernel。
#   PyTorch:
#       void at::native::elementwise_kernel<(int)128, (int)4, void at::native::gpu_kernel_impl_nocast<at::native::direct_copy_kernel_cuda(at::TensorIteratorBase &)::[lambda() (instance 3)]::operator ()() const::[lambda() (instance 12)]::operator ()() const::[lambda(c10::BFloat16) (instance 1)]>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
#       ampere_bf16_s16816gemm_bf16_128x128_ldg8_f2f_stages_64x3_tn
#       void at::native::elementwise_kernel<(int)128, (int)4, void at::native::gpu_kernel_impl_nocast<at::native::CUDAFunctor_add<c10::BFloat16>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
#
#   DTorch
#       ampere_bf16_s16816gemm_bf16_64x64_ldg8_f2f_stages_64x5_tn
#       void at::native::elementwise_kernel<(int)128, (int)4, void at::native::gpu_kernel_impl_nocast<at::native::CUDAFunctor_add<c10::BFloat16>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)

hidden_states, encoder_hidden_states = (
    hidden_states[:, : residual.shape[1]],
    hidden_states[:, residual.shape[1] :],
)

# attn.to_out[0] = torch.nn.Linear
hidden_states = attn.to_out[0](hidden_states)