diff --git a/.gitignore b/.gitignore index 4f6c4a0c7..aa893c251 100644 --- a/.gitignore +++ b/.gitignore @@ -34,3 +34,13 @@ dev/cuda/global_norm # log files *.log + +# ROCm/HIP build artifacts +build/ +test_gpt2cu* +train_gpt2cu* +*.bin +*.dll +amd_comgr* +rocblas/ +nul diff --git a/Makefile b/Makefile index 73b83720c..b2334070e 100644 --- a/Makefile +++ b/Makefile @@ -25,6 +25,39 @@ NVCC_CUDNN = # By default we don't build with cudnn because it blows up compile time from a few seconds to ~minute USE_CUDNN ?= 0 +# ROCm / HIP build for AMD GPUs. Set USE_HIP=1 to compile the .cu sources with +# hipcc instead of nvcc; the same target names (train_gpt2cu, test_gpt2cu, ...) +# then build for AMD. The target arch is auto-detected with amdgpu-arch and can +# be overridden with AMDGPU_TARGETS=; gfx90a is the fallback default. +USE_HIP ?= 0 +ifeq ($(USE_HIP), 1) + # Mirror the nvidia-smi compute_cap query below: when the arch is not given, + # detect the installed GPUs with amdgpu-arch (ships with ROCm/LLVM). The tool + # is often not on PATH (it lives in /llvm/bin), so fall back to locating + # it via hipconfig --rocmpath. An absent tool yields empty output, so the + # strip-check below falls back to gfx90a. + ifndef AMDGPU_TARGETS + ifneq ($(CI),true) + AMDGPU_ARCH_TOOL := $(shell which amdgpu-arch 2>/dev/null) + ifeq ($(AMDGPU_ARCH_TOOL),) + AMDGPU_ARCH_TOOL := $(shell hipconfig --rocmpath 2>/dev/null)/llvm/bin/amdgpu-arch + endif + AMDGPU_TARGETS := $(shell $(AMDGPU_ARCH_TOOL) 2>/dev/null | sort -u | paste -sd ';') + endif + endif + ifeq ($(strip $(AMDGPU_TARGETS)),) + AMDGPU_TARGETS := gfx90a + endif + # Wavefront width is 64 on CDNA (gfx9xx) and 32 on RDNA (gfx10xx/gfx11xx). + # Derive it from the target arch and pass it to both HIP compile passes; the + # sources use it for host launch geometry and device reductions alike. + ifneq ($(filter gfx9%,$(AMDGPU_TARGETS)),) + LLMC_WARP_SIZE ?= 64 + else + LLMC_WARP_SIZE ?= 32 + endif +endif + # We will place .o files in the `build` directory (create it if it doesn't exist) BUILD_DIR = build ifeq ($(OS), Windows_NT) @@ -47,7 +80,8 @@ endef endif ifneq ($(CI),true) # if not in CI, then use the GPU query - ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY= + ifeq ($(USE_HIP), 1) # HIP build: arch comes from AMDGPU_TARGETS, skip nvidia-smi + else ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY= ifneq ($(call file_exists_in_path, nvidia-smi),) # Get the compute capabilities of all GPUs # Remove decimal points, sort numerically in ascending order, and select the first (lowest) value @@ -66,8 +100,47 @@ endif $(info ---------------------------------------------) ifneq ($(OS), Windows_NT) - NVCC := $(shell which nvcc 2>/dev/null) - NVCC_LDFLAGS += -lnvidia-ml + ifeq ($(USE_HIP), 1) + # HIP toolchain: hipcc compiles the .cu sources directly (no hipify step). + # The compat header (llmc/cuda_to_hip.h) is force-included on every HIP TU so + # the CUDA-spelled symbols resolve, and llmc/hip_shims is on the include path + # so the CUDA-named toolkit headers (, , ...) + # forward to it. NVCC is repointed at hipcc so the existing build rules apply. + HIPCC ?= $(shell which hipcc 2>/dev/null) + NVCC := $(HIPCC) + # NOTE: do NOT use clang's -ffast-math here. It is far more aggressive than + # nvcc's --use_fast_math (it enables -fassociative-math / -funsafe-math- + # optimizations / -fno-signed-zeros), which reassociates the online-softmax + # and layernorm-backward reductions on gfx90a into NaN gradients (the forward + # loss stays correct, only the backward NaNs). -ffp-contract=fast gives the + # FMA contraction that matters for perf while keeping IEEE semantics. + NVCC_FLAGS := -O$(FORCE_NVCC_O) -std=c++17 -ffp-contract=fast -fno-math-errno \ + $(addprefix --offload-arch=,$(AMDGPU_TARGETS)) \ + -DUSE_HIP=1 -DLLMC_WARP_SIZE=$(LLMC_WARP_SIZE) \ + -include llmc/cuda_to_hip.h -I llmc/hip_shims + # ROCm's clang selects the highest /usr/lib/gcc// dir even when + # that GCC's libstdc++ headers are absent (e.g. Ubuntu installs libgcc-14-dev + # without libstdc++-14-dev), failing with "Could not find standard C++ header". + # Probe for that and pin --gcc-install-dir to the newest GCC version that has + # matching headers under /usr/include/c++/. + HIP_STDLIB_OK := $(shell $(HIPCC) -x c++ -fsyntax-only -include cmath /dev/null >/dev/null 2>&1 && echo 1) + ifneq ($(HIP_STDLIB_OK),1) + HIP_GCC_DIR := $(shell for d in /usr/lib/gcc/*/*; do v=$$(basename "$$d"); [ -d "/usr/include/c++/$$v" ] && echo "$$d"; done | sort -V | tail -n1) + ifneq ($(strip $(HIP_GCC_DIR)),) + $(info → hipcc cannot find libstdc++ headers; pinning --gcc-install-dir=$(HIP_GCC_DIR)) + NVCC_FLAGS += --gcc-install-dir=$(HIP_GCC_DIR) + endif + endif + NVCC_LDFLAGS := -lhipblas -lhipblaslt + NVCC_INCLUDES := + NVCC_LDLIBS := + # -lineinfo is nvcc-only; hipcc (clang) rejects it. + LINEINFO := + else + NVCC := $(shell which nvcc 2>/dev/null) + NVCC_LDFLAGS += -lnvidia-ml + LINEINFO := -lineinfo + endif # Function to test if the compiler accepts a given flag. define check_and_add_flag @@ -83,7 +156,16 @@ else CFLAGS := REMOVE_FILES = del *.exe,*.obj,*.lib,*.exp,*.pdb && del SHELL_UNAME := Windows - ifneq ($(shell where nvcc 2> nul),"") + ifeq ($(USE_HIP), 1) + # HIP toolchain on Windows. Mirrors the non-Windows USE_HIP branch above but + # uses cmd `where` (not the unix `which`) for detection. hipcc compiles the + # .cu sources directly; the compat header is force-included and hip_shims is + # on the include path so the CUDA-named toolkit headers forward to HIP. + # `where` can return several matches (hipcc.bat, hipcc.exe, ...); take the first. + HIPCC ?= $(firstword $(shell where hipcc 2> nul)) + NVCC := $(HIPCC) + LINEINFO := + else ifneq ($(shell where nvcc 2> nul),"") NVCC := nvcc else NVCC := @@ -94,7 +176,24 @@ else LDFLAGS := LDLIBS := INCLUDES := - NVCC_FLAGS += -I"dev" + ifeq ($(USE_HIP), 1) + # -DNOMINMAX/-DWIN32_LEAN_AND_MEAN: the HIP runtime headers pull in , + # whose min/max macros otherwise break std::min/std::max in the sources. + NVCC_FLAGS := -O3 -std=c++17 -ffp-contract=fast -fno-math-errno \ + $(addprefix --offload-arch=,$(AMDGPU_TARGETS)) \ + -DUSE_HIP=1 -DLLMC_WARP_SIZE=$(LLMC_WARP_SIZE) -DNOMINMAX -DWIN32_LEAN_AND_MEAN \ + -include llmc/cuda_to_hip.h -I llmc/hip_shims -I"dev" + # Windows ROCm lib linkage: hipblas ships an MSVC import lib (hipblas.lib), + # but hipblaslt ships only a GNU import lib (libhipblaslt.dll.a) with no + # hipblaslt.lib, so lld-link must consume it by full path, not via -l. Pass + # HIP_LIB_DIR=/lib on the make command line (the ROCm lib directory). + # -Xlinker for the .dll.a so hipcc's -x hip does not treat it as a source file. + NVCC_LDFLAGS := -L"$(HIP_LIB_DIR)" -lhipblas -Xlinker "$(HIP_LIB_DIR)/libhipblaslt.dll.a" + NVCC_INCLUDES := + NVCC_LDLIBS := + else + NVCC_FLAGS += -I"dev" + endif ifeq ($(WIN_CI_BUILD),1) $(info Windows CI build) OUTPUT_FILE = /link /OUT:$@ @@ -102,7 +201,13 @@ else else $(info Windows local build) OUTPUT_FILE = /link /OUT:$@ && copy /Y $@ $@.exe - CUDA_OUTPUT_FILE = -o $@ && copy /Y $@.exe $@ + # hipcc/clang emits $@.exe from -o $@.exe; copy to the extensionless name the + # README/test commands invoke. (nvcc's -o $@ already produces $@.exe.) + ifeq ($(USE_HIP), 1) + CUDA_OUTPUT_FILE = -o $@.exe && copy /Y $@.exe $@ + else + CUDA_OUTPUT_FILE = -o $@ && copy /Y $@.exe $@ + endif endif endif @@ -283,7 +388,7 @@ test_gpt2fp32cu: test_gpt2_fp32.cu $(NVCC) $(NVCC_FLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) profile_gpt2cu: profile_gpt2.cu $(NVCC_CUDNN) - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $(LINEINFO) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) clean: $(REMOVE_FILES) $(TARGETS) diff --git a/README.md b/README.md index d2d107821..979f58461 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,17 @@ python dev/data/tinyshakespeare.py python train_gpt2.py ``` +## quick start (AMD GPU, ROCm/HIP) + +llm.c also builds and trains on AMD GPUs through ROCm/HIP. With a [ROCm](https://rocm.docs.amd.com/) installation (7.2 or newer), build any of the GPU targets by adding `USE_HIP=1` and your GPU architecture to the make command: + +```bash +make train_gpt2cu USE_HIP=1 AMDGPU_TARGETS=gfx90a +./train_gpt2cu +``` + +Set `AMDGPU_TARGETS` to your GPU (for example `gfx90a` for CDNA2 / MI200, or `gfx1100` for RDNA3). The default NVIDIA/CUDA build is unchanged; `USE_HIP=1` repoints the build at `hipcc`. + ## quick start (CPU) The "I am so GPU poor that I don't even have one GPU" section. You can still enjoy seeing llm.c train! But you won't go too far. Just like the fp32 version above, the CPU version is an even earlier checkpoint in the history of llm.c, back when it was just a simple reference implementation in C. For example, instead of training from scratch, you can finetune a GPT-2 small (124M) to output Shakespeare-like text, as an example: diff --git a/llmc/cuda_common.h b/llmc/cuda_common.h index 6f5bf6564..bcc7e9ce2 100644 --- a/llmc/cuda_common.h +++ b/llmc/cuda_common.h @@ -28,11 +28,21 @@ extern cudaDeviceProp deviceProp; // WarpSize is not a compile time constant // Defining here like this possibly allows the compiler to optimize better +// On ROCm the wavefront is 64 on CDNA (gfx90a/gfx94x) and 32 on RDNA; the build +// derives LLMC_WARP_SIZE from the single target arch and defines it for both the +// host and device compile passes (see llmc/cuda_to_hip.h), so the host launch +// geometry and the device reductions agree. +#if defined(USE_HIP) || defined(__HIP_PLATFORM_AMD__) +#define WARP_SIZE ((unsigned)LLMC_WARP_SIZE) +#else #define WARP_SIZE 32U +#endif // try to make sure that 2 blocks fit on A100/H100 to maximise latency tolerance // this needs to be defines rather than queried to be used for __launch_bounds__ -#if __CUDA_ARCH__ == 800 || __CUDA_ARCH__ >= 900 +#if defined(USE_HIP) || defined(__HIP_PLATFORM_AMD__) +#define MAX_1024_THREADS_BLOCKS 1 +#elif __CUDA_ARCH__ == 800 || __CUDA_ARCH__ >= 900 #define MAX_1024_THREADS_BLOCKS 2 #else #define MAX_1024_THREADS_BLOCKS 1 @@ -98,7 +108,9 @@ typedef __nv_bfloat16 floatX; // our own versions if none already exist, otherwise the compiler will complain. // If not, you easily get "no viable overload" (for sm52) and "function already exists" (sm_80) -#if defined(ENABLE_BF16) && (__CUDACC_VER_MAJOR__ < 12) && !((__CUDA_ARCH__ >= 800) || !defined(__CUDA_ARCH__)) +// On HIP the compat header (llmc/cuda_to_hip.h) supplies generic __ldcs/__stcs, +// so this NVIDIA-only bf16 fallback is excluded. +#if !defined(USE_HIP) && !defined(__HIP_PLATFORM_AMD__) && defined(ENABLE_BF16) && (__CUDACC_VER_MAJOR__ < 12) && !((__CUDA_ARCH__ >= 800) || !defined(__CUDA_ARCH__)) __device__ floatX __ldcs(const floatX* address) { unsigned short bf = __ldcs(reinterpret_cast(address)); return __nv_bfloat16_raw{bf}; diff --git a/llmc/cuda_to_hip.h b/llmc/cuda_to_hip.h new file mode 100644 index 000000000..c40a9e82d --- /dev/null +++ b/llmc/cuda_to_hip.h @@ -0,0 +1,243 @@ +/* +CUDA-to-HIP compatibility shim for the AMD/ROCm build of llm.c. + +This is the only file that knows about HIP. It is force-included on every +HIP translation unit by the Makefile (-include llmc/cuda_to_hip.h), so the +CUDA-spelled sources (cudaMalloc, cublasLt*, __nv_bfloat16, ...) compile +unchanged on ROCm. On NVIDIA this header is never included and the CUDA path +is byte-for-byte the original. + +Layout: + 1. libc headers BEFORE so host memcpy/memset win over + HIP's __device__ overloads inside a .cu compiled as HIP. + 2. runtime / bf16 / fp16 / library headers. + 3. symbol aliases (runtime, bf16/fp16 types, cuBLAS(Lt), profiler, NVTX). + 4. warp-size and full-warp-mask abstractions for the fault classes. + 5. streaming load/store and cooperative-groups reduce shims HIP lacks. +*/ +#ifndef LLMC_CUDA_TO_HIP_H +#define LLMC_CUDA_TO_HIP_H + +#if defined(USE_HIP) || defined(__HIP_PLATFORM_AMD__) + +#include +#include + +#include +#include +#include +#include +#include + +// ---------------------------------------------------------------------------- +// Runtime API +#define cudaError_t hipError_t +#define cudaSuccess hipSuccess +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#define cudaErrorMemoryAllocation hipErrorOutOfMemory + +#define cudaMalloc hipMalloc +#define cudaFree hipFree +#define cudaMallocManaged hipMallocManaged +#define cudaMemAdvise hipMemAdvise +#define cudaMemAdviseSetPreferredLocation hipMemAdviseSetPreferredLocation +#define cudaCpuDeviceId hipCpuDeviceId +#define cudaMallocHost hipHostMalloc +#define cudaFreeHost hipHostFree +#define cudaHostAllocWriteCombined hipHostMallocWriteCombined + +#define cudaMemcpy hipMemcpy +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemsetAsync hipMemsetAsync +#define cudaMemset hipMemset +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost +#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice +#define cudaMemcpyHostToHost hipMemcpyHostToHost +#define cudaMemGetInfo hipMemGetInfo + +#define cudaStream_t hipStream_t +#define cudaStreamCreate hipStreamCreate +#define cudaStreamDestroy hipStreamDestroy +#define cudaStreamSynchronize hipStreamSynchronize +#define cudaStreamWaitEvent hipStreamWaitEvent +#define cudaStreamNonBlocking hipStreamNonBlocking +#define cudaStreamCreateWithPriority hipStreamCreateWithPriority + +#define cudaEvent_t hipEvent_t +#define cudaEventCreate hipEventCreate +#define cudaEventDestroy hipEventDestroy +#define cudaEventRecord hipEventRecord +#define cudaEventSynchronize hipEventSynchronize +#define cudaEventElapsedTime hipEventElapsedTime + +#define cudaDeviceProp hipDeviceProp_t +#define cudaGetDeviceProperties hipGetDeviceProperties +#define cudaSetDevice hipSetDevice +#define cudaGetDevice hipGetDevice +#define cudaGetDeviceCount hipGetDeviceCount +#define cudaDeviceSynchronize hipDeviceSynchronize +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#define cudaFuncSetAttribute hipFuncSetAttribute +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize + +// hipProfilerStart/Stop are deprecated and return hipErrorNotSupported unless a +// profiler is attached, which would trip the cudaCheck wrapper. These are +// profiling-only hooks with no effect on compute, so make them succeed no-ops. +#define cudaProfilerStart() hipSuccess +#define cudaProfilerStop() hipSuccess + +// ---------------------------------------------------------------------------- +// bf16 / fp16 types and intrinsics +#define __nv_bfloat16 __hip_bfloat16 +#define nv_bfloat16 __hip_bfloat16 +#define __nv_bfloat162 __hip_bfloat162 +#define nv_bfloat162 __hip_bfloat162 +#define __nv_bfloat16_raw __hip_bfloat16_raw +#define __nv_bfloat162_raw __hip_bfloat162_raw +// HIP's __float2bfloat16 already rounds to nearest; CUDA spells it _rn. +#define __float2bfloat16_rn __float2bfloat16 + +// ---------------------------------------------------------------------------- +// cuBLAS / cuBLASLt -> hipBLAS / hipBLASLt +#include +#include + +// scalar data types (cudaDataType / cudaDataType_t -> hipDataType from ) +#define cudaDataType hipDataType +#define cudaDataType_t hipDataType +#define cublasDataType_t hipDataType +#define CUDA_R_32F HIP_R_32F +#define CUDA_R_16F HIP_R_16F +#define CUDA_R_16BF HIP_R_16BF + +// handles / status / ops +#define cublasStatus_t hipblasStatus_t +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define cublasOperation_t hipblasOperation_t +#define CUBLAS_OP_N HIPBLAS_OP_N +#define CUBLAS_OP_T HIPBLAS_OP_T +#define cublasComputeType_t hipblasComputeType_t +#define CUBLAS_COMPUTE_32F HIPBLAS_COMPUTE_32F +#define CUBLAS_COMPUTE_16F HIPBLAS_COMPUTE_16F +#define CUBLAS_COMPUTE_32F_FAST_TF32 HIPBLAS_COMPUTE_32F_FAST_TF32 + +// plain cuBLAS (the FP32 driver uses cublasSgemm / cublasSgemmStridedBatched) +#define cublasHandle_t hipblasHandle_t +#define cublasCreate hipblasCreate +#define cublasDestroy hipblasDestroy +#define cublasSetStream hipblasSetStream +#define cublasSgemm hipblasSgemm +#define cublasSgemmStridedBatched hipblasSgemmStridedBatched +#define cublasMath_t hipblasMath_t +#define cublasSetMathMode hipblasSetMathMode +#define CUBLAS_DEFAULT_MATH HIPBLAS_DEFAULT_MATH +#define CUBLAS_TF32_TENSOR_OP_MATH HIPBLAS_TF32_TENSOR_OP_MATH + +// cuBLASLt object types +#define cublasLtHandle_t hipblasLtHandle_t +#define cublasLtCreate hipblasLtCreate +#define cublasLtDestroy hipblasLtDestroy +#define cublasLtMatmul hipblasLtMatmul +#define cublasLtMatmulDesc_t hipblasLtMatmulDesc_t +#define cublasLtMatmulDescCreate hipblasLtMatmulDescCreate +#define cublasLtMatmulDescDestroy hipblasLtMatmulDescDestroy +#define cublasLtMatmulDescSetAttribute hipblasLtMatmulDescSetAttribute +#define cublasLtMatrixLayout_t hipblasLtMatrixLayout_t +#define cublasLtMatrixLayoutCreate hipblasLtMatrixLayoutCreate +#define cublasLtMatrixLayoutDestroy hipblasLtMatrixLayoutDestroy +#define cublasLtMatrixLayoutSetAttribute hipblasLtMatrixLayoutSetAttribute +#define cublasLtMatmulPreference_t hipblasLtMatmulPreference_t +#define cublasLtMatmulPreferenceCreate hipblasLtMatmulPreferenceCreate +#define cublasLtMatmulPreferenceDestroy hipblasLtMatmulPreferenceDestroy +#define cublasLtMatmulPreferenceSetAttribute hipblasLtMatmulPreferenceSetAttribute +#define cublasLtMatmulHeuristicResult_t hipblasLtMatmulHeuristicResult_t +#define cublasLtMatmulAlgoGetHeuristic hipblasLtMatmulAlgoGetHeuristic +#define cublasLtEpilogue_t hipblasLtEpilogue_t + +// cuBLASLt attribute / epilogue enums +#define CUBLASLT_MATMUL_DESC_TRANSA HIPBLASLT_MATMUL_DESC_TRANSA +#define CUBLASLT_MATMUL_DESC_TRANSB HIPBLASLT_MATMUL_DESC_TRANSB +#define CUBLASLT_MATMUL_DESC_EPILOGUE HIPBLASLT_MATMUL_DESC_EPILOGUE +#define CUBLASLT_MATMUL_DESC_BIAS_POINTER HIPBLASLT_MATMUL_DESC_BIAS_POINTER +#define CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE HIPBLASLT_MATMUL_DESC_BIAS_DATA_TYPE +#define CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER +#define CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD HIPBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD +#define CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES +#define CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT +#define CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET +#define CUBLASLT_EPILOGUE_DEFAULT HIPBLASLT_EPILOGUE_DEFAULT +#define CUBLASLT_EPILOGUE_GELU_AUX HIPBLASLT_EPILOGUE_GELU_AUX +#define CUBLASLT_EPILOGUE_GELU_AUX_BIAS HIPBLASLT_EPILOGUE_GELU_AUX_BIAS +#define CUBLASLT_EPILOGUE_DGELU HIPBLASLT_EPILOGUE_DGELU +#define CUBLASLT_EPILOGUE_BIAS HIPBLASLT_EPILOGUE_BIAS +#define CUBLASLT_EPILOGUE_BGRADB HIPBLASLT_EPILOGUE_BGRADB + +// ---------------------------------------------------------------------------- +// NVTX profiling markers -> no-ops (roctx is optional and profiling-only) +#define nvtxRangePush(x) ((void)0) +#define nvtxRangePop() ((void)0) +#define nvtxNameCudaStreamA(stream, name) ((void)0) + +// ---------------------------------------------------------------------------- +// Warp size and full-warp mask (fault classes: warp size 32-vs-64, lane masks) +// +// HIP has two compile passes per .cu: a host pass and a device pass. Arch +// macros like __GFX9__ are defined ONLY in the device pass, so they cannot +// drive a constant that the host launch code (dim3(WARP_SIZE, ...)) also reads +// -- host and device would then disagree and the launch geometry would be +// wrong. llm.c uses WARP_SIZE both as a host launch dimension AND as a device +// constant, so it must be one value per build. The Makefile derives +// LLMC_WARP_SIZE from the single target arch (gfx90a/gfx94x => 64, RDNA => 32) +// and defines it for both passes, keeping host and device identical. +#ifndef LLMC_WARP_SIZE +#error "LLMC_WARP_SIZE must be set by the build (derived from the target arch)" +#endif + +// HIP's __shfl*_sync require a 64-bit mask regardless of the active wave width +// (ROCm static_asserts sizeof(MaskT)==8); the CUDA 0xFFFFFFFF literal will not +// compile. This is NOT keyed on wave width, only on the toolchain. +#define LLMC_FULL_WARP_MASK 0xffffffffffffffffULL + +// ---------------------------------------------------------------------------- +// Streaming cache-hint load/store: HIP lacks __stcs/__stcg and only provides +// __ldcs for __half/__half2 (). The non-template half/half2 +// __ldcs overloads win over this unconstrained template, which covers the +// remaining int4 / float / floatX / unsigned short call sites. The hints are +// advisory; a plain load/store is semantically identical. +template __device__ __forceinline__ T __ldcs(const T* ptr) { return *ptr; } +template __device__ __forceinline__ void __stcs(T* ptr, T value) { *ptr = value; } +template __device__ __forceinline__ void __stcg(T* ptr, T value) { *ptr = value; } + +// ---------------------------------------------------------------------------- +// Cooperative-groups reduce: HIP's CG has no cg::reduce / cg::plus / cg::greater +// (used by train_gpt2_fp32.cu over tiled_partition<32>). Provide a butterfly +// reduction over the tile; width = tile.size() stays within the 32-lane tile on +// a 64-lane wavefront, so it is correct on both wave32 and wave64. +#include +#include +namespace cooperative_groups { +// ROCm 7.2.x's cooperative_groups had neither cg::plus/cg::greater nor cg::reduce. +// Newer ROCm (>= 7.13) ships cg::plus/cg::greater natively (redefining them is an +// error there) but STILL has no cg::reduce -- so version-guard only the functors, +// and always supply reduce. Butterfly reduction over the tile; width = tile.size() +// stays within the 32-lane tile on a 64-lane wavefront (correct on wave32/wave64). +#if HIP_VERSION < 71300000 +template struct plus { __device__ T operator()(T a, T b) const { return a + b; } }; +template struct greater { __device__ T operator()(T a, T b) const { return a > b ? a : b; } }; +#endif +template +__device__ __forceinline__ T reduce(const Group& g, T val, Op op) { + for (int offset = g.size() / 2; offset > 0; offset >>= 1) { + val = op(val, g.shfl_xor(val, offset)); + } + return val; +} +} // namespace cooperative_groups + +#else // NVIDIA: no-op include of the CUDA runtime; the CUDA path is unchanged. +#include +#endif // USE_HIP + +#endif // LLMC_CUDA_TO_HIP_H diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 030ec073e..e34539a17 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -143,17 +143,29 @@ __global__ void copy_and_cast_kernel(Td* dst, const Ts* src, size_t n, ptrdiff_t // ---------------------------------------------------------------------------- // Warp/Block communication primitives +// Full-warp shuffle mask and starting butterfly offset. On HIP the mask must be +// 64-bit (it will not compile otherwise) and the offset must span the whole +// wavefront (WARP_SIZE/2 == 32 on wave64), or a wave64 reduction would silently +// fold only 32 of 64 lanes. On NVIDIA both reduce to the original 0xFFFFFFFF/16. +#if defined(USE_HIP) || defined(__HIP_PLATFORM_AMD__) +#define WARP_REDUCE_MASK LLMC_FULL_WARP_MASK +#define WARP_REDUCE_OFFSET (WARP_SIZE / 2) +#else +#define WARP_REDUCE_MASK 0xFFFFFFFFU +#define WARP_REDUCE_OFFSET 16 +#endif + // warp-level reduction for summing values __device__ inline float warpReduceSum(float val) { - for (int offset = 16; offset > 0; offset /= 2) { - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + for (int offset = WARP_REDUCE_OFFSET; offset > 0; offset /= 2) { + val += __shfl_xor_sync(WARP_REDUCE_MASK, val, offset); } return val; } // warp-level reduction for finding the maximum value __device__ inline float warpReduceMax(float val) { - for (int offset = 16; offset > 0; offset /= 2) { - val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + for (int offset = WARP_REDUCE_OFFSET; offset > 0; offset /= 2) { + val = fmaxf(val, __shfl_xor_sync(WARP_REDUCE_MASK, val, offset)); } return val; } @@ -166,7 +178,12 @@ template __device__ inline float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) { // two reductions of up to 1024 threads: // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) - __shared__ float shared_val[WARP_SIZE]; + // The cross-warp scratch is indexed by warp_id, so it must hold one entry per + // warp. The runtime warp count is not a constant expression, so size it to the + // compile-time upper bound: a 1024-thread block has at most 1024/32 = 32 warps + // (the narrowest wave is 32 lanes), which also covers wave64 (<=16 warps). + constexpr int kMaxWarpsPerBlock = 1024 / 32; + __shared__ float shared_val[kMaxWarpsPerBlock]; const int lane_id = threadIdx.x % WARP_SIZE; const int warp_id = threadIdx.x / WARP_SIZE; const int num_warps = blockDim.x / WARP_SIZE; diff --git a/llmc/hip_shims/cooperative_groups.h b/llmc/hip_shims/cooperative_groups.h new file mode 100644 index 000000000..9ed79b4b1 --- /dev/null +++ b/llmc/hip_shims/cooperative_groups.h @@ -0,0 +1,5 @@ +// HIP-only forwarding shim: -> HIP's CG header plus the +// llm.c compat header (which adds the cg::reduce / cg::plus / cg::greater that +// HIP's cooperative groups lacks). +#include +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/cooperative_groups/reduce.h b/llmc/hip_shims/cooperative_groups/reduce.h new file mode 100644 index 000000000..be36ea367 --- /dev/null +++ b/llmc/hip_shims/cooperative_groups/reduce.h @@ -0,0 +1,3 @@ +// HIP-only forwarding shim: -> the llm.c compat +// header, which provides the cg::reduce shim (HIP's CG has no reduce.h). +#include "../../cuda_to_hip.h" diff --git a/llmc/hip_shims/cublasLt.h b/llmc/hip_shims/cublasLt.h new file mode 100644 index 000000000..93df5ea78 --- /dev/null +++ b/llmc/hip_shims/cublasLt.h @@ -0,0 +1,3 @@ +// HIP-only forwarding shim: -> the llm.c compat header +// (which includes and aliases the cublasLt* symbols). +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/cublas_v2.h b/llmc/hip_shims/cublas_v2.h new file mode 100644 index 000000000..1a8c582c3 --- /dev/null +++ b/llmc/hip_shims/cublas_v2.h @@ -0,0 +1,3 @@ +// HIP-only forwarding shim: -> the llm.c compat header +// (which includes and aliases the cublas* symbols used). +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/cuda_bf16.h b/llmc/hip_shims/cuda_bf16.h new file mode 100644 index 000000000..ab577391e --- /dev/null +++ b/llmc/hip_shims/cuda_bf16.h @@ -0,0 +1,3 @@ +// HIP-only forwarding shim: -> the llm.c compat header +// (which includes and aliases __nv_bfloat16 -> __hip_bfloat16). +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/cuda_fp16.h b/llmc/hip_shims/cuda_fp16.h new file mode 100644 index 000000000..9afd9fbf9 --- /dev/null +++ b/llmc/hip_shims/cuda_fp16.h @@ -0,0 +1,3 @@ +// HIP-only forwarding shim: -> the llm.c compat header +// (which includes ; HIP defines `half`/`__half`). +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/cuda_profiler_api.h b/llmc/hip_shims/cuda_profiler_api.h new file mode 100644 index 000000000..babb3fad0 --- /dev/null +++ b/llmc/hip_shims/cuda_profiler_api.h @@ -0,0 +1,4 @@ +// HIP-only forwarding shim: -> the llm.c compat header +// (which maps cudaProfilerStart/Stop to hipSuccess no-ops -- hipProfilerStart +// returns hipErrorNotSupported and would trip cudaCheck). +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/cuda_runtime.h b/llmc/hip_shims/cuda_runtime.h new file mode 100644 index 000000000..7b2377fbe --- /dev/null +++ b/llmc/hip_shims/cuda_runtime.h @@ -0,0 +1,4 @@ +// HIP-only forwarding shim: -> the llm.c compat header. +// On the HIP build this dir is on the include path; on NVIDIA it is absent so +// the real CUDA toolkit header wins. See llmc/cuda_to_hip.h. +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/cuda_runtime_api.h b/llmc/hip_shims/cuda_runtime_api.h new file mode 100644 index 000000000..d1b0a77b8 --- /dev/null +++ b/llmc/hip_shims/cuda_runtime_api.h @@ -0,0 +1,2 @@ +// HIP-only forwarding shim: -> the llm.c compat header. +#include "../cuda_to_hip.h" diff --git a/llmc/hip_shims/nccl.h b/llmc/hip_shims/nccl.h new file mode 100644 index 000000000..66c7115fb --- /dev/null +++ b/llmc/hip_shims/nccl.h @@ -0,0 +1,3 @@ +// HIP-only forwarding shim: -> RCCL (the ROCm NCCL drop-in). Only used +// by the optional MULTI_GPU build path; the RCCL API mirrors NCCL 1:1. +#include diff --git a/llmc/hip_shims/nvtx3/nvToolsExt.h b/llmc/hip_shims/nvtx3/nvToolsExt.h new file mode 100644 index 000000000..0168275a7 --- /dev/null +++ b/llmc/hip_shims/nvtx3/nvToolsExt.h @@ -0,0 +1,3 @@ +// HIP-only forwarding shim: -> the llm.c compat header +// (which stubs nvtxRangePush/Pop to no-ops). +#include "../../cuda_to_hip.h" diff --git a/llmc/hip_shims/nvtx3/nvToolsExtCudaRt.h b/llmc/hip_shims/nvtx3/nvToolsExtCudaRt.h new file mode 100644 index 000000000..14f4d5f2d --- /dev/null +++ b/llmc/hip_shims/nvtx3/nvToolsExtCudaRt.h @@ -0,0 +1,2 @@ +// HIP-only forwarding shim: -> the llm.c compat header. +#include "../../cuda_to_hip.h" diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index 9777d0658..886cc2db7 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -443,7 +443,9 @@ void layernorm_forward(floatX* out, float* mean, float* rstd, // in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute // this may fail, in which case we fall back to the smem free implementation. cudaCheck(cudaGetLastError()); - auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + // hipFuncSetAttribute takes the kernel as const void* (CUDA also accepts this + // and additionally has a templated T* overload); cast for portability. + auto status = cudaFuncSetAttribute((const void*)layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); cudaCheck(cudaGetLastError()); if (status == cudaSuccess) { layernorm_forward_kernel6<<>>(out, mean, rstd, inp, weight, bias, N, C); @@ -476,7 +478,7 @@ void fused_residual_forward5(floatX* residual, floatX* normed, float* mean, floa // in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute // this may fail, in which case we fall back to the smem free implementation. cudaCheck(cudaGetLastError()); - auto status = cudaFuncSetAttribute(fused_residual_forward_kernel5, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + auto status = cudaFuncSetAttribute((const void*)fused_residual_forward_kernel5, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); cudaCheck(cudaGetLastError()); if(status == cudaSuccess) { fused_residual_forward_kernel5<<>>(residual, normed, @@ -497,7 +499,13 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr const int blocks_per_sm = 2; // supported on every architecture and less cache thrashing than 3 const int grid_size = blocks_per_sm * deviceProp.multiProcessorCount; size_t rounded_C = CEIL_DIV(C, (32 * x128::size)) * (32 * x128::size); - size_t shared_mem_size = (2 * rounded_C + 2 * (block_size - 32) * f128::size) * sizeof(float); + // The kernel bases the dbias/dweight temp-shared regions at offsets of + // WARP_SIZE*f128::size (see *_tmp_shared in layernorm_backward_kernel10), so + // the host reservation uses WARP_SIZE to match the kernel's exact footprint. + // A literal 32 on wave64 reserves 2*(64-32)*f128::size floats MORE than the + // device uses -- a harmless over-allocation, not corruption -- but keeping it + // WARP_SIZE keeps the host reservation and device layout consistent. + size_t shared_mem_size = (2 * rounded_C + 2 * (block_size - WARP_SIZE) * f128::size) * sizeof(float); cudaCheck(cudaMemsetAsync(scratch, 0, 1 * sizeof(float), stream)); // only need to reset the flag to 0 layernorm_backward_kernel10<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); diff --git a/llmc/matmul.cuh b/llmc/matmul.cuh index becc372c6..080c50783 100644 --- a/llmc/matmul.cuh +++ b/llmc/matmul.cuh @@ -50,11 +50,13 @@ __global__ void matmul_backward_bias_kernel9(OutFloat* dbias, const floatX* dout __shared__ float sub_results[x128::size][WARP_SIZE][bdy]; - // reduce within-warp results + // reduce within-warp results. WARP_REDUCE_MASK is the full-warp mask (64-bit + // on HIP, which will not compile with the 0xffffffff literal); the width-4 + // sub-group is wave-size-agnostic (4 <= warpSize on both wave32 and wave64). for (int k = 0; k < x128::size; k++) { float v = accumulators[k]; - v += __shfl_down_sync(0xffffffff, v, 1, 4); - v += __shfl_down_sync(0xffffffff, v, 2, 4); + v += __shfl_down_sync(WARP_REDUCE_MASK, v, 1, 4); + v += __shfl_down_sync(WARP_REDUCE_MASK, v, 2, 4); if(warp_d == 0) { sub_results[k][block_d][warp_c] = v; } @@ -66,8 +68,8 @@ __global__ void matmul_backward_bias_kernel9(OutFloat* dbias, const floatX* dout float a = 0.f; for (int r = warp_d; r < blockDim.z; r += bdx) { float v = sub_results[k][r][warp_c]; - v += __shfl_down_sync(0xffffffff, v, 1, 4); - v += __shfl_down_sync(0xffffffff, v, 2, 4); + v += __shfl_down_sync(WARP_REDUCE_MASK, v, 1, 4); + v += __shfl_down_sync(WARP_REDUCE_MASK, v, 2, 4); a += v; } if(warp_d == 0 && global_oc < OC) { @@ -198,8 +200,12 @@ void matmul_cublaslt(floatX* d, const floatX* a, const floatX* b, const floatX* } // set scale type to FP32 (needs to be FP16 if and only if using CUBLAS_COMPUTE_16F, so it's FP32 even for FP8!) + // hipBLASLt has no CUBLASLT_MATMUL_DESC_SCALE_TYPE attribute; the scale type + // follows the compute type (FP32 here), so this is a no-op on HIP. +#if !defined(USE_HIP) && !defined(__HIP_PLATFORM_AMD__) cublasDataType_t scale_type = CUDA_R_32F; cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type))); +#endif // find a suitable algorithm (cached internally so shouldn't take much CPU time in practice) cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc, ALayout, BLayout, CLayout, DLayout, @@ -255,8 +261,11 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, const int block_size = deviceProp.maxThreadsPerMultiProcessor == 1536 ? 768 : 1024; - dim3 block_dim = {4, 8, (unsigned)block_size/WARP_SIZE}; - const int OC_per_warp = block_dim.y * x128::size; // 64 at BF16 + // block_dim.y must equal the kernel's bdy = WARP_SIZE/bdx (bdx==4), so the + // warp is fully tiled into width-4 sub-groups; hardcoding 8 (the wave32 + // value) trips the kernel's assert(blockDim.y == bdy) on wave64. + dim3 block_dim = {4, WARP_SIZE / 4, (unsigned)block_size/WARP_SIZE}; + const int OC_per_warp = block_dim.y * x128::size; // 64 at BF16 on wave32 const int grid_size_x = CEIL_DIV(OC, OC_per_warp); // e.g. 12 horizontal blocks for 768 OCs at BF16 const int grid_size_y = max(1, deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / (block_size * grid_size_x)); // full GPU! diff --git a/train_gpt2_fp32.cu b/train_gpt2_fp32.cu index df412ea5e..fcdfeaf2f 100644 --- a/train_gpt2_fp32.cu +++ b/train_gpt2_fp32.cu @@ -815,7 +815,16 @@ void matmul_backward(float* dinp, float* dweight, float* dbias, // backward to bias, if given, does a += if (dbias != NULL) { const int block_size = 1024; - const int grid_size = OC / 32; // for now, OC must be divisible by 32 for this kernel to work + // each block reduces warpSize columns (tl = blockIdx.x * warpSize in the + // kernel), so the grid must stride by the wavefront width: 32 on NVIDIA / + // RDNA, 64 on CDNA. Spacing blocks by a literal 32 on wave64 would make + // adjacent blocks overlap and the last block read past dout (-> NaN dbias). +#if defined(USE_HIP) || defined(__HIP_PLATFORM_AMD__) + const int warp_width = LLMC_WARP_SIZE; +#else + const int warp_width = 32; +#endif + const int grid_size = OC / warp_width; // OC must be divisible by the wavefront width matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); cudaCheck(cudaGetLastError()); }