Description
I investigated three fp16 linear_cross_entropy test failures on gfx90a / MI200 originally reported in ROCm/TheRock#6993 and was able to reproduce them and trace the numerical behavior further.
The failures affect the reduction='none' backward path for the balanced, compact, and auto accumulation policies.
Related PyTorch CI tracking: pytorch/pytorch#191552.
Root cause analysis
The issue does not appear to be caused by a rocBLAS or hipBLASLt kernel defect.
In the reduction='none' backward path, linear_cross_entropy recomputes the forward logits. This re-computation happens while ROCmBackwardPassGuard is active.
As a result, the fp16 GEMM used to recompute the logits is classified as a backward GEMM and gets the ROCm fp16 alternate implementation flag:
rocblas_gemm_flags_fp16_alt_impl
PyTorch documents this behavior in docs/source/notes/numerical_accuracy.md. On AMD MI200 GPUs, the alternate fp16 implementation is intended to preserve denormal values. The fp16 inputs are converted through an intermediate BF16 representation, GEMM accumulation is performed in fp32, and the result is converted back to fp16. PyTorch enables this alternate implementation by default for backward GEMMs.
The important detail here is that this particular GEMM is logically recomputing a forward activation — the logits — even though it executes during the backward pass.
Therefore, the recomputed logits inherit the backward-only fp16 alternate numerical behavior. The resulting difference in the recomputed logits propagates into the weight gradient and causes the balanced, compact, and auto fp16 tests to exceed their 3 * eps tolerance.
Validation
As a validation, disabling the alternate implementation with:
ROCBLAS_INTERNAL_FP16_ALT_IMPL=0
makes all three tests pass.
The resulting numerical errors also match the results from my exact-GEMM experiments.
This indicates that the failures are specifically associated with using the backward fp16 alternate implementation when recomputing the forward logits, rather than with an incorrect rocBLAS/hipBLASLt GEMM kernel.
Question
This seems to raise a PyTorch-side semantics question:
Should a GEMM that recomputes a forward activation during backward inherit the backward-only ROCm fp16 alternate implementation simply because it executes under ROCmBackwardPassGuard?
Or should this forward-logit re-computation explicitly opt out of the alternate implementation so that its numerical behavior matches the original forward computation?
I would appreciate guidance on what the intended behavior should be here.
References
-
ROCm/TheRock#6993 — original report and reproduction on gfx90a
-
pytorch/pytorch#191552 — MI200 CI failure tracker
-
docs/source/notes/numerical_accuracy.md— documentation for the MI200 fp16 alternate implementation