GPT-2 example: add a CUDA.jl version (#2765)
* GPT-2 example: add a CUDA.jl version next to the Reactant and PyTorch ones
Split examples/gpt2_fineweb into flux-reactant/ (the reference), flux-cuda/
and pytorch/, sharing data/ and described by a top-level README. The shared
training loop moves to flux-reactant/training.jl, and HellaSwag scoring no
longer depends on Reactant. flux-cuda/ has its own environment without
Reactant and includes the reference files; it runs eagerly on CuArrays with
Zygote and cuDNN's flash attention.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* GPT-2 example: CUDA.jl version at micro-batch 16 with Flux 0.16.14
Flux 0.16.14 (fused cross-entropy for labels, `NNlib.layernorm` in `LayerNorm`, no copy of
bf16 inputs), NNlib 0.9.55 (fused normalization kernels) and Zygote 0.7.14 (tied-embedding
gradient on the GPU) make the CUDA.jl step 4x faster and halve its memory: GPT-2 124M now
trains at ~135k tokens/s with the preset's micro-batch of 16, against ~31k at 8.
The Reactant numbers are those of Flux 0.16.13: with 0.16.14 the bf16 model trains in
Float32 under Reactant, to be fixed in NNlib.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* GPT-2 example: link the NNlib fix for the Reactant slowdown
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* GPT-2 example: PyTorch with torch.compile in the comparison
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
