[ROCm] Fix TorchScript JIT BF16 HIPRTC overload conflict (ROCM-23829)

[srinivamd]

Summary

Fixes ROCM-23829: Megatron-DeepSpeed LLaMA2 pretraining crashes during TorchScript JIT BF16 warmup with __float2bfloat16 return-type overload conflict.

Root Cause

On ROCm >= 7.13 (rocm-systems#4727), HIPRTC headers now bundle amd_hip_bf16.h which defines __float2bfloat16(float) returning __hip_bfloat16. PyTorch's TorchScript JIT fuser (resource_strings.h) emits its own inline __float2bfloat16(const float) returning __nv_bfloat16 into every JIT-generated GPU kernel. Since __hip_bfloat16 and __nv_bfloat16 are different types, these two definitions differ only in return type, producing a fatal HIPRTC compile error:

error: functions that differ only in their return type cannot be overloaded
   51 | __attribute__((device)) __nv_bfloat16 __float2bfloat16(const float a) {
      |                         ~~~~~~~~~~~~~ ^
hiprtc_runtime.h:13560: note: previous definition is here
13560 | __attribute__((device)) static inline __hip_bfloat16 __float2bfloat16(float f) {
      |                                       ~~~~~~~~~~~~~~ ^

This blocks all BF16 JIT fusion workloads (Megatron-DeepSpeed bias_gelu warmup) on MI300X/MI350X.

Fix

In the emitted JIT kernel string (bfloat16_support_literal), detect the HIP bf16 header guard (_HIP_INCLUDE_HIP_AMD_DETAIL_HIP_BF16_H_) at HIPRTC compile time:

• When present (ROCm >= 7.13 with HIPRTC bf16 headers): typedef __hip_bfloat16 __nv_bfloat16 — alias to the native HIP type. HIP-provided __float2bfloat16 / __bfloat162float intrinsics work transparently since __nv_bfloat16 IS __hip_bfloat16.
• When absent (older ROCm): preserve existing inline struct + intrinsic definitions unchanged.

This is backward-compatible — on older ROCm where HIPRTC does not include bf16 headers, the guard is false and behavior is identical to before.

Reproducer

import torch
torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_profiling_executor(False)
torch._C._jit_override_can_fuse_on_cpu(True)
torch._C._jit_override_can_fuse_on_gpu(True)

@torch.jit.script
def bias_gelu(bias, y):
    x = bias + y
    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))

bias = torch.rand(128, dtype=torch.bfloat16, device=\"cuda\")
inp  = torch.rand((32, 4, 128), dtype=torch.bfloat16, device=\"cuda\")
bias.requires_grad, inp.requires_grad = True, True
for _ in range(5):
    out = bias_gelu(bias, inp)
print(\"OK\")

Affected Models

• pyt_deepspeed_megatron_llama2_7b
• pyt_deepspeed_megatron_llama2_13b
• pyt_deepspeed_megatron_llama2_70b
• pyt_deepspeed_megatron_gpt3_13b

Test Plan

• Run reproducer script above on ROCm 7.13 with MI350X (gfx950)
• Verify Megatron-DeepSpeed LLaMA2-70B pretraining starts successfully
• Verify no regression on ROCm 7.12 (where HIPRTC bf16 headers are absent)

Cherry-pick targets

• release/2.11
• release/2.10

Co-Authored-By: Claude Opus 4 (1M context) noreply@anthropic.com

[rocm-repo-management-api]

Jenkins build for 2a10123010ff117f03ba3c6b0a9d616633ab9b17 commit finished as FAILURE
Links: Pipeline Overview / Build artifacts / Test Results