Repository navigation
Add mixed-precision FP16 inference to the Metal backend - #3
Conversation
There was a problem hiding this comment.
🔵 Needs a closer look
Mixed-precision GPU numerics rely on MPS FP16 GEMM accumulation behavior, buffer alignment, and hazard tracking that cannot be fully verified by static review, and full-checkpoint/real-data qualification is explicitly deferred.
No open issues 🎉
What changed in this PR
This PR adds a mixed-precision FP16 path to the existing Apple Silicon Metal/MPS backend, enabling --device metal --precision fp16 across inference, host rollout, and the ./varda forecast wrapper. It is stacked on #2 (which introduced the FP32 Metal backend) and carries only the incremental FP16 support. The precision policy mirrors CUDA/autocast: FP16 storage for weights and activations, but FP32 for GEMM results/accumulation, bias addition, LayerNorm statistics/affine, attention reductions/softmax, and all physical output operations. A key design constraint is that MPS is never asked for an all-half GEMM (its intermediate range can overflow before cancellation); instead, FP16 operands feed a GEMM that writes an FP32 scratch result, then linear_store_half rounds once after bias.
Changes:
- Remove the Metal FP16 rejection (in
model.cppand thevardawrapper) and wire precision selection through the backend, selecting half vs. float storage via anActivationtypedef and a per-tensorelement_size. - Add FP16 machinery: padded half weight rows, reusable packing/scratch buffers,
pack_half_from_float/pack_half_from_half/linear_store_halfkernels, half-aware arenas, chunked FP16 GEMMs, and hook upcasting; all kernels gain FP32-no-op casts so the FP32 path stays bitwise-identical to #2. - Extend tests (synthetic FP16 runs, per-precision parity/chunk loops, a dense-cancellation regression, CPU-FP16 rejection, FP16 LayerNorm shader cases) and update docs/CI accordingly.
| File | Description |
|---|---|
src/varda/metal/backend_metal.mm |
Core FP16 support: element_size-aware Tensor, half_data, half arenas, padded-row linear loads, FP16 GEMM path with packing/FP32 scratch/store, hook upcasting. |
src/varda/metal/kernels.metal |
Activation typedef, pack/store-half kernels, and FP32-no-op casts across elementwise/attention/LayerNorm/residual kernels. |
src/varda/model.cpp |
Removes the Metal FP32-only guard so the Metal backend accepts FP16. |
varda |
Drops the wrapper's Metal FP16 rejection; updates the setup note. |
python/tests/test_metal.py |
Adds FP16 to repeated inference/parity/rollout/empty-graph tests, a dense-cancellation regression, and switches the rejection test to CPU FP16. |
tests/metal/test_metal_kernels.mm |
Runs the direct LayerNorm shader test for both FP32 and FP16 with a half-aware tolerance/guard. |
tests/run_synthetic.sh |
Adds Metal FP16 default/chunk-37 synthetic runs. |
.github/workflows/cpu-portability.yml |
Globs all a.metal-fp*.log files on failure. |
docs/metal.md, docs/apple-silicon.md, README.md, CMakeLists.txt |
Document and describe the new mixed-FP16 capability. |
Review effort: Balanced
5358459 to
55b1eef
Compare
Adds
--device metal --precision fp16for inference, host rollout and the forecast wrapper. Metal previously rejected FP16. This PR is stacked on #2 and targetscodex/apple-silicon-metal; it contains only the additional mixed-precision support.The precision policy follows the CUDA/autocast approach: FP16 weights and activation storage, FP32 GEMM results/accumulation and bias addition, FP32 LayerNorm statistics/affine and attention reductions/softmax, and FP32 physical output operations. Each dense result is rounded once to FP16 after bias. Edge projection uses half operands, FP32 sums and an explicit half projection rounding. This is not a bitwise CUDA implementation.
MPS never receives an all-half GEMM. Padded weights and reusable input packing handle odd widths and unaligned chunks; compatible half activations are used directly. FP32 result scratch and embedding chunks keep temporary storage bounded. Public input/output remain FP32, and hooks upcast the same stored production values after synchronization. FP32 retains its existing path; CPU stays FP32-only and CUDA is unchanged.
Validation on Apple M1, without changing any tolerance files:
proc.3.outin e: rel_absmax 1.395e-3, max_abs 2.488e-2 against the existing 3e-3 tolerance. Input stages are bitwise exact, and hooks are inert in both modes.The full local GPU suite passes with Metal API and shader validation enabled. Hosted CI retains API validation for MPS and shader validation for the direct kernel tests, following #2's documented Paravirtual-GPU instrumentation limitation.
Two alternating 20-call FP32/FP16 benchmark rounds on M1 (first three discarded) measured warm mean times of 27.1 -> 23.0 ms for synthetic e and 32.6 -> 27.8 ms for d, about 1.18x faster. Repeats are bitwise stable. These small-model timings do not establish full-Varda performance. Half storage reduces selected buffers, not the entire footprint.
Full-checkpoint and real-data forecast/rollout qualification still require a larger-memory Mac and real reference assets. CUDA execution was not retested. No Metal graph capture or GPU-resident rollout window is added.
Linux x86_64 and macOS arm64 CI pass at 5358459. The macOS job executed FP32/FP16 GPU parity and all seven integration tests on an Apple Paravirtual device; GPU checks were not skipped.