[ROCm] Add AMD GPU (ROCm/HIP) support to EGGROLL#8
Open
jeffdaily wants to merge 11 commits into
Open
Conversation
This adds AMD GPU support to full_cuda_train_egg.cu, the integer-only EGGROLL
language-model trainer, building the existing CUDA source with hipcc. The
NVIDIA build is unchanged: all HIP-specific code is guarded on __HIP__ (which
hipcc -x hip defines), and the two new compat headers expand to nothing under
nvcc.
Two small headers carry the port surface. egg_hip_compat.cuh maps the fixed
set of CUDA runtime symbols this file uses (cudaMalloc/Free/Memset/Memcpy, the
memcpy-kind enums, cudaDeviceProp, cudaGetDeviceProperties,
cudaDeviceSynchronize) to their hip* equivalents and aliases the cub namespace
to hipcub; rocThrust already provides thrust:: at the same paths, so Thrust
needs no remap. egg_warp_compat.cuh holds the one load-bearing semantic fix.
The fix: EGGROLL maps one perturbation to one logical 32-lane warp and
partitions HIDDEN_DIM across exactly those 32 lanes (loops stride by 32,
per-lane arrays sized HIDDEN_DIM_max/32). That 32 is a data-layout stride, not
the hardware wavefront width. On a wave64 device (CDNA, e.g. MI200) the
physical wavefront is 64, so two logical warps share one wavefront. The CUB
warp reduction is pinned to width 32 (cub::WarpReduce<long long, 32> -- hipCUB
otherwise defaults the logical width to the 64-lane physical warp, which would
sum two perturbations together) and the 64-bit broadcast routes through a
width-32 __shfl_sync with a 64-bit all-ones mask. On wave32 devices (RDNA, the
logical warp equals the wavefront) and under nvcc (default warp width 32) these
are no-ops, so the same source is correct everywhere.
Building on AMD selects the architecture from the hipcc command line, so no GPU
architecture is hardcoded:
hipcc -O3 --offload-arch=gfx90a --offload-arch=gfx1100 -x hip \
full_cuda_train_egg.cu -o egg_hip
A few host-only adjustments let the trainer build under MSVC on Windows:
unistd.h is guarded behind a non-Windows check (with a fputs fallback for the
SIGINT handler), and clock_gettime(CLOCK_MONOTONIC) is provided via
QueryPerformanceCounter. None of this touches the GPU code path. An optional
EGG_FIXED_SEED environment variable makes the per-step seed a pure function of
(seed, step) so two runs produce an identical loss sequence; default behavior
(wall-clock seed) is unchanged.
The README gains an "AMD GPU (ROCm/HIP)" build block alongside the existing
CUDA one.
Validation: trained on real AMD GPUs across architectures -- gfx90a (MI250X,
wave64), gfx1100 (Radeon Pro W7800, RDNA3), gfx1201 (RX 9070 XT, RDNA4). On
every device the loss decreases monotonically (8.3489 -> 3.3417 over 16 steps),
the generated text sample becomes word-like, and two EGG_FIXED_SEED runs
produce an identical Loss/Up+/Up- sequence -- the decisive check that the
32-lane reductions and shuffle masks are correct (a wrong-width partition would
diverge). The wave64 gfx90a result is the strongest evidence the logical-warp
fix is right. The CPU ternary-packing unit test (d-eggs/test_ternary.cpp) still
builds and passes.
This work was authored with assistance from Claude.
Builds additively on the integer-only trainer's compat shims so the
transformer trainers can compile with hipcc instead of nvcc, with every
HIP body behind #if defined(__HIP__) so the NVIDIA build is unchanged.
egg_warp_compat.cuh gains two width-32 shuffle helpers the attention and
RoPE kernels need: eggShflDownSync (head reduction, off from 16) and
eggShflXorSync (RoPE neighbor-pair swap, lane_mask 1). Both pin the
shuffle to EGG_WARP_SIZE (32) with a wavefront-wide mask so that on wave64
the two logical 32-lane warps sharing one physical wavefront stay
independent; on wave32 and NVIDIA the explicit width is a no-op.
egg_hip_compat.cuh aliases the two additional CUDA runtime symbols the
transformer trainers reference (cudaGetSymbolAddress, cudaMemcpyDeviceToDevice)
and adds a portable __dp4a definition. ROCm exposes the amdgcn sdot4 builtin
only under the dot1-insts target feature, absent on some supported arches
(e.g. gfx1100), so the shim computes the 8-bit packed dot product directly;
the compiler lowers it to the hardware dot instruction where present and to
plain multiplies otherwise, with the same numeric result on every arch.
Authored with assistance from Claude (Anthropic).
Test Plan:
hipcc -O3 --offload-arch=gfx1100 -x hip \
full_cuda_train_egg_transformer.cu -o /tmp/egg_xf_gfx1100
hipcc -O3 --offload-arch=gfx1100 -x hip \
full_cuda_train_egg_transformer_adam.cu -o /tmp/egg_xfadam_gfx1100
Both compile clean (only pre-existing nodiscard warnings on hipError_t
returns); llvm-objdump --offloading confirms a gfx1100 code object.
Compiles full_cuda_train_egg_transformer.cu and
full_cuda_train_egg_transformer_adam.cu (and their shared
egg_adaptive_normalize.h helper) with hipcc. The CUDA includes are guarded
#if !defined(__HIP__) and replaced by the compat headers under __HIP__, so
the NVIDIA nvcc build is untouched (every change is additive; the original
lines remain in #else branches).
The load-bearing fix is the warp width. WARP_SIZE here is a 32-lane
data-layout/tiling stride, not the physical wavefront, and must stay 32 on
wave32 and wave64 alike. The attention head reduction reduces one 64-element
head with a 32-lane shuffle-down (off from 16) and a lane-0 atomicAdd, so on
wave64 the two logical 32-lane warps that share a wavefront must not mix; the
RoPE neighbor swap is a width-1 partner exchange that likewise must stay
inside the logical warp. Both head reductions and the RoPE shfl_xor now route
through the explicit width-32 helpers, as does the logical-32 max reduction in
egg_adaptive_normalize.h. cub::BlockReduce is a block-wide collective over an
explicit BLOCK_THREADS, so it needs no warp-width change beyond the
namespace alias the compat header already provides.
Authored with assistance from Claude (Anthropic).
Test Plan:
hipcc -O3 --offload-arch=gfx1100 -x hip \
full_cuda_train_egg_transformer.cu -o /tmp/egg_xf_gfx1100
hipcc -O3 --offload-arch=gfx1100 -x hip \
full_cuda_train_egg_transformer_adam.cu -o /tmp/egg_xfadam_gfx1100
Both compile clean for gfx1100 (only pre-existing nodiscard warnings);
llvm-objdump --offloading confirms a gfx1100 code object in each. GPU
training validation (loss-decrease + pinned-seed determinism) runs next.
The two transformer trainers (full_cuda_train_egg_transformer.cu and full_cuda_train_egg_transformer_adam.cu) used time(NULL) ^ (step * 0x12345678) for the per-step perturbation seed, making runs wall-clock-dependent and non-reproducible. Add the same EGG_FIXED_SEED env override that PR d0rc#8 added to full_cuda_train_egg.cu: when EGG_FIXED_SEED is set, the seed is a pure function of (fixed_seed, step) so two runs with the same seed produce an identical loss sequence. Default behavior (unset EGG_FIXED_SEED) is unchanged. This enables the determinism gate needed to confirm that the width-32 warp shuffle helpers (eggShflDownSync / eggShflXorSync) and the portable __dp4a shim are correct: a wrong 64-lane partition or a wrong dot-product would break determinism and produce divergent loss values between fixed-seed runs. Tested on AMD Radeon Pro W7800 (gfx1100, wave32), ROCm 7.2.1, hipcc -O3 --offload-arch=gfx1100 -x hip. Both trainers build clean (pre-existing -Wunused-value warnings only). Two EGG_FIXED_SEED=12345 runs of the adam trainer produce bit-identical loss sequences across 90+ steps. The __dp4a portable shim and warp shuffles are validated as correct on gfx1100 by this determinism gate. Authored with Claude claude-sonnet-4-6 as AI assistant.
Adds the cuBLAS->hipBLAS and multi-GPU runtime surface the Muon optimizer
and the replicated data-parallel mgpu trainer need, additively on the
existing transformer-trainer shims, with every HIP body behind
#if defined(__HIP__) so the NVIDIA nvcc build is unchanged.
The hipBLAS aliases (handle/status/operation types, the OP_T/OP_N and
pointer-mode enums, and Create/Destroy/SetStream/SetPointerMode/Snrm2/Sgemm)
map 1:1 onto cuBLAS: hipBLAS (rocBLAS underneath) is column-major like
cuBLAS, so the Newton-Schulz leading-dimension and transpose logic carries
over without change. The multi-GPU aliases (the cudaStream_t typedef,
SetDevice, GetDeviceCount, the stream create/destroy/sync calls, host pinned
allocation, and the async memcpy/memset) cover the per-device stream and
model-replica enumeration plus the host pinned-memory fitness aggregation;
no NCCL/RCCL and no peer access are involved.
Authored with assistance from Claude (Anthropic).
Test Plan:
hipcc -O3 --offload-arch=gfx1100 -x hip -DUSE_MUON=1 \
full_cuda_train_transformer_adam_mgpu.cu -lhipblas -o egg_mgpu
# links libhipblas + librocblas; gfx1100 code object confirmed.
Compiles full_cuda_train_transformer_adam_mgpu.cu and its Muon
(muon_internal.cuh) and NTT (egg_ntt.cuh) headers with hipcc. The CUDA
includes are guarded #if !defined(__HIP__) and replaced by the compat
headers under __HIP__, so the NVIDIA nvcc build is untouched; every change
is additive and the original lines remain in #else branches.
The Muon Newton-Schulz iteration's cuBLAS calls (Snrm2 in device pointer
mode, two Sgemm orientations, the pointer-mode toggles) route through the
hipBLAS aliases; because hipBLAS shares the cuBLAS column-major convention
the OP_T/OP_N and leading-dimension mapping is unchanged. The multi-GPU
driver is replicated data-parallel: each device holds a full model copy and
a population shard on its own stream, fitness is aggregated through host
pinned memory and broadcast back per device. There is no NCCL/RCCL and no
peer access, so the device enumeration and stream/async calls map 1:1.
As in the single-GPU transformers, WARP_SIZE is a 32-lane data-layout stride
rather than the physical wavefront and must stay 32 on wave32 and wave64.
The RoPE neighbor swap and the two attention head reductions (a 32-lane
shuffle-down with a lane-0 atomicAdd) now run through the explicit width-32
helpers so the two logical warps sharing a wave64 wavefront stay
independent. egg_ntt.cuh fences every transform stage with a block-wide
__syncthreads(), so it is correct on both wavefront widths with only the
include guard. A reproducible EGG_FIXED_SEED override is added to the
training-loop seed, matching the other trainers and behavior-preserving when
unset.
Authored with assistance from Claude (Anthropic).
Test Plan:
hipcc -O3 --offload-arch=gfx1100 -x hip -DUSE_MUON=1 \
full_cuda_train_transformer_adam_mgpu.cu -lhipblas -o egg_mgpu
llvm-objdump --offloading egg_mgpu | grep -io 'gfx[0-9a-f]*' | sort -u # gfx1100
# also clean: USE_MUON=0 default, and USE_MUON=1 -DNTT_MODE=1.
The RoPE look-up table d_ROPE_LUT was sized for SEQ_LEN (32) positions, but generate_sequence_kernel runs for gen_seed_len + gen_output_len = 96 positions and apply_rope_integer() indexes the table at t * HEAD_DIM, which is out of bounds for any generated position t >= SEQ_LEN. On HIP this faults with hipErrorIllegalAddress immediately after the first training step; on CUDA the same access silently read out-of-bounds garbage, so the RoPE values applied to generated positions >= SEQ_LEN were already wrong there too. This is a pre-existing bug in the upstream source that the HIP build merely surfaced as a hard fault. The fix corrects both backends and is therefore unconditional: size the table to cover the maximum generation length the kernel can index (ROPE_LUT_MAX_LEN = max(MAX_GEN_LEN, SEQ_LEN)) and extend the init_tables() fill loop to populate the added positions. Authored with assistance from Claude. Test Plan: ``` hipcc -O3 --offload-arch=gfx1100 -x hip \ full_cuda_train_transformer_adam_mgpu.cu -o /tmp/egg_mgpu_ropefix -lhipblas ``` Builds clean (exit 0; only pre-existing -Wnodiscard warnings). The generate kernel no longer faults and produces correct generation.
Adds AMD ROCm support to the d-eggs Coordinator-Worker distributed EGGROLL trainer. The worker GPU translation unit is built with hipcc; the coordinator stays plain host C++. Every HIP-specific change sits behind #if defined(__HIP__) so the existing nvcc build is unchanged. Review order: start with the new include/utils/hip_compat.cuh, a self-contained CUDA-to-HIP shim for the d-eggs include tree (d-eggs carries its own headers, so it does not reach into the single-GPU trainers' compat headers). It maps the runtime, cuB and hipBLAS surface the worker, kernels and model/optimizer headers use, and provides the width-32 shuffle helpers and a portable __dp4a. The other headers gain a guarded include of it in place of the direct CUDA headers. The substantive device changes are the wave-width fixes. WARP_SIZE in config.h is a fixed 32-lane data-layout stride, not the physical wavefront; on wave64 (gfx90a) two logical warps share one wavefront. The attention head reductions in model/layers.cuh, the RoPE neighbor-pair exchange, and the adaptive-norm max-reduction in math/adaptive_norm.cuh now run at explicit width 32 with a wavefront-wide mask via eggShflDownSync/eggShflXorSync, keeping the two logical warps independent. The NTT/WHT path fences every stage with __syncthreads(), so it is correct on wave32 and wave64 unchanged. The Muon Newton-Schulz cuBLAS path maps 1:1 to hipBLAS (column-major, so OP_T/OP_N transfer directly). The worker's managed-memory allocations map to hipMallocManaged / hipMemAdvise / hipMemPrefetchAsync, and the captured optimizer-step graph maps to the hipGraph family. The captured region issues only async launches, memsets and copies on a single stream with no host sync (and no cuBLAS in the default build), so it is capture-safe; capture and replay build and link without an eager fallback. The single worker translation unit needs no relocatable device code, so the HIP Makefile arm omits -fgpu-rdc; HIPARCH selects the target(s) and defaults to gfx90a, caller-overridable. This work was authored with the assistance of Claude, an AI assistant by Anthropic. Test Plan: ``` make -C d-eggs USE_HIP=1 HIPARCH=gfx1100 coordinator worker print_arch llvm-objdump --offloading d-eggs/worker | grep -io 'gfx[0-9a-f]*' # gfx1100 g++ -O3 -Id-eggs/include -o test_ternary d-eggs/test_ternary.cpp && ./test_ternary ```
The d-eggs distributed trainer sizes the RoPE lookup table for the training window only (ROPE_LUT_SIZE = SEQ_LEN * HEAD_DIM), but generate_sequence_kernel runs for gen_seed_len + gen_output_len positions (96 by default), well past SEQ_LEN (64). apply_rope_integer indexes d_ROPE_LUT[t * HEAD_DIM + ...]; at t = SEQ_LEN the index reaches ROPE_LUT_SIZE, the first out-of-bounds element. On NVIDIA this read silently returns garbage; under HIP it faults, putting the context into an error state at step 0 so all later launches no-op (Loss stays 0, Updates stays 0). This is the same pre-existing upstream bug already fixed for the single-process transformer trainer; d-eggs carries an independent copy of the RoPE LUT definition that the earlier fix did not reach. Sizing the table to cover the max generation position corrects both the NVIDIA and the ROCm backends. The new MAX_GEN_LEN bound is overridable and defaults to the kernel's generation length; ROPE_LUT_SIZE and the LUT fill loop scale from it, so the allocations and cudaMemcpyToSymbol need no further change. Authored with assistance from Claude. Test Plan: ``` make -C projects/egg.c/src/d-eggs clean make -C projects/egg.c/src/d-eggs USE_HIP=1 HIPARCH=gfx1100 \ coordinator worker print_arch \ CFLAGS="-O3 -Iinclude -DVOCAB_SIZE=256" \ GPUFLAGS="-O3 -Iinclude -x hip --offload-arch=gfx1100 -lhipblas -DVOCAB_SIZE=256" llvm-objdump --offloading d-eggs/worker | grep -io 'gfx[0-9a-f]*' | sort -u ``` Builds clean (pre-existing nodiscard warnings only); worker code object is gfx1100. Distributed GPU run is re-validated separately.
jeffdaily
changed the title
The width-32 warp helpers in egg_warp_compat.cuh (eggShflDownSync, eggShflXorSync) are portable: their bodies expand to the same __shfl_*_sync intrinsics on both backends, and on NVIDIA EGG_FULL_MASK is 0xFFFFFFFFu with the explicit width 32 already the default warp width, so they are exactly equivalent to the raw __shfl_*_sync calls they wrap. The per-call #if defined(__HIP__) / #else duplication was therefore unnecessary; the integer trainer already calls its shim (eggWarpBroadcast) unconditionally. Collapse the eight guarded shuffle sites to the single portable call, and move the egg_warp_compat.cuh include out of the HIP-only block to an unconditional include placed after cub is in scope, so the shims are declared on the NVIDIA build too (they previously were not, which is why the guards had a raw-intrinsic CUDA branch). The header already handles both backends: egg_hip_compat.cuh aliases cub to hipcub before it, and EGG_FULL_MASK is backend-selected. Device code is unchanged on AMD: each affected translation unit compiles to a bit-identical gfx90a code object (disassembled ISA identical; only the non-semantic module fingerprint shifts with the source edit). The NVIDIA build is restored to compiling. Authored with the assistance of Claude. Test Plan: CUDA (nvcc 12.6, -arch=sm_86), all link cleanly: nvcc -O3 -arch=sm_86 full_cuda_train_egg_transformer.cu -o egg_transformer nvcc -O3 -arch=sm_86 full_cuda_train_egg_transformer_adam.cu -o egg_adam nvcc -O3 -arch=sm_86 full_cuda_train_transformer_adam_mgpu.cu -lcublas -o egg_mgpu nvcc -O3 -arch=sm_86 -DUSE_MUON=1 full_cuda_train_transformer_adam_mgpu.cu -lcublas -o egg_mgpu_muon HIP device-code equivalence (gfx90a), before vs after this change: hipcc -O3 --offload-arch=gfx90a -x hip <each>.cu [-lhipblas] -> egg_transformer / egg_transformer_adam: device ISA identical -> egg_mgpu: gfx90a code object disassembly identical (13736 insns)
egg_adaptive_normalize.h carried the same #if defined(__HIP__) / #else split as the transformer trainers: the HIP branch called the portable eggShflDownSync shim, the CUDA branch the raw __shfl_down_sync it wraps. The two are equivalent (the shim expands to the same intrinsic, with EGG_FULL_MASK = 0xFFFFFFFFu and the default width 32 on NVIDIA), so collapse both reduction sites to the unconditional shim call. This header is included only by the two adam/mgpu trainers, which already include egg_warp_compat.cuh unconditionally, so the shim is in scope on both backends. Device code is unchanged on AMD: the two affected translation units (adam, mgpu) compile to identical gfx90a/gfx1100/gfx1201 code objects before and after. Authored with the assistance of Claude. Test Plan: CUDA (nvcc 12.6, -arch=sm_86): all trainer variants + d-eggs make link cleanly. HIP device-code equivalence (gfx90a, gfx1100, gfx1201), before vs after: hipcc -O3 --offload-arch=<arch> -x hip full_cuda_train_egg_transformer_adam.cu hipcc -O3 --offload-arch=<arch> -x hip full_cuda_train_transformer_adam_mgpu.cu -lhipblas -> codeobj_diff verdict=identical for both, all three arches
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds AMD GPU support across EGGROLL via ROCm/HIP, building the existing CUDA sources with hipcc. Every HIP-specific change is guarded (on
__HIP__/USE_HIP), so the NVIDIA build is unchanged.Scope (the full repo)
full_cuda_train_egg.cu) - the base ternary-weight, single-GPU, no-cuBLAS trainer.full_cuda_train_egg_transformer.cu,..._adam.cu) - width-32 logical-warp head reductions and RoPE shuffles, plus a portable__dp4ashim (ROCm has no global__dp4a).full_cuda_train_transformer_adam_mgpu.cu,muon_internal.cuh) - the project's cuBLAS path mapped to hipBLAS (Snrm2/Sgemm Newton-Schulz, column-major 1:1), replicated data-parallel across GPUs (no RCCL).d-eggs/) - a self-contained HIP compat shim, hipBLAS Muon, managed memory, hipGraph capture/replay, and aUSE_HIPMakefile arm (caller-selectedHIPARCH).How it works
Two top-level compat headers (
egg_hip_compat.cuh,egg_warp_compat.cuh) and a d-eggs-local shim carry thecuda*->hip*runtime aliases, the cuBLAS->hipBLAS aliases, the width-32 shuffle helpers, and the__dp4ashim. EGGROLL maps one perturbation to a logical 32-lane warp; on a 64-lane wavefront (CDNA) the warp reductions are pinned to width 32 so two perturbations are not folded together (a no-op on wave32 and on NVIDIA). The matmul launch geometry is derived from the runtime warp size rather than a hardcoded width.This also fixes a pre-existing out-of-bounds in
generate_sequence_kernel(the RoPE LUT was sized for the training window but generation indexes past it) that the HIP build exposed - HIP faulted where CUDA silently read out of bounds; the fix corrects both backends (in the multi-GPU trainer and in d-eggs).The CUDA build is unchanged throughout: every HIP path is guarded, the standalone trainers build with nvcc exactly as before, and d-eggs keeps its original nvcc Makefile recipe. On AMD, build with
USE_HIP/--offload-arch=<arch>(and-DUSE_HIP=1 HIPARCH=<arch>for d-eggs); the README documents the AMD build alongside the CUDA one.Validation (real AMD GPUs)
The full port is validated on Linux across both wavefront widths - gfx90a (Instinct MI250X, CDNA2, wave64) and gfx1100 (Radeon Pro W7800, RDNA3, wave32) - covering every component:
__dp4ashim is confirmed by the adam trainer's identical fixed-seed loss across steps.The single-GPU trainers are additionally validated on Windows gfx1201 (RX 9070 XT, RDNA4). The multi-GPU Int8NativeFormer and the distributed d-eggs system are Linux-targeted (they use POSIX sockets /
unistd.h) and are validated on the Linux multi-GPU hosts above.Authored with the assistance of Claude.