diff --git a/benchmarks/python/cross_entropy_bench.py b/benchmarks/python/cross_entropy_bench.py new file mode 100644 index 0000000000..cc4667fe51 --- /dev/null +++ b/benchmarks/python/cross_entropy_bench.py @@ -0,0 +1,86 @@ +"""Fused Metal cross entropy vs the decomposed fallback.""" + +import time + +import mlx.core as mx + + +def fallback(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze(-1) + return mx.logsumexp(logits.astype(mx.float32), axis=-1) - score.astype(mx.float32) + + +def timeit(fn, *args, iters=50, warmup=10): + # Multi-GiB logits leave the buffer cache under pressure, which otherwise + # bleeds between measurements and makes later shapes look far slower. + mx.clear_cache() + for _ in range(warmup): + mx.eval(fn(*args)) + mx.synchronize() + start = time.perf_counter() + for _ in range(iters): + mx.eval(fn(*args)) + mx.synchronize() + return (time.perf_counter() - start) / iters * 1e3 + + +shapes = [ + (8192, 4096), + (4096, 32000), + (2048, 128256), + (1024, 151936), +] + +print( + f"{'rows x vocab':>20} {'dtype':>10} {'fallback':>10} {'fused':>10} {'speedup':>9}" +) +print("-" * 64) + +for dtype in [mx.float32, mx.bfloat16]: + for rows, V in shapes: + logits = mx.random.normal(shape=(rows, V), scale=2.0).astype(dtype) + targets = mx.random.randint(0, V, shape=(rows,)) + mx.eval(logits, targets) + + t_ref = timeit(lambda: fallback(logits, targets)) + t_fused = timeit(lambda: mx.fast.cross_entropy(logits, targets)) + print( + f"{rows:>7} x {V:<10} {str(dtype).split('.')[-1]:>10} " + f"{t_ref:>9.3f}ms {t_fused:>9.3f}ms {t_ref / t_fused:>8.2f}x" + ) + +print() +print("--- forward + backward ---") +print( + f"{'rows x vocab':>20} {'dtype':>10} {'fallback':>10} {'fused':>10} {'speedup':>9}" +) +print("-" * 64) + +for dtype in [mx.float32, mx.bfloat16]: + for rows, V in shapes: + logits = mx.random.normal(shape=(rows, V), scale=2.0).astype(dtype) + targets = mx.random.randint(0, V, shape=(rows,)) + mx.eval(logits, targets) + + g_ref = mx.grad(lambda x, y: fallback(x, y).sum(), argnums=0) + g_fused = mx.grad(lambda x, y: mx.fast.cross_entropy(x, y).sum(), argnums=0) + t_ref = timeit(g_ref, logits, targets, iters=20) + t_fused = timeit(g_fused, logits, targets, iters=20) + print( + f"{rows:>7} x {V:<10} {str(dtype).split('.')[-1]:>10} " + f"{t_ref:>9.3f}ms {t_fused:>9.3f}ms {t_ref / t_fused:>8.2f}x" + ) + +print() +print("--- peak memory, forward + backward, 4096 x 128256 bf16 ---") +for label, fn in [ + ("fallback", lambda x, y: fallback(x, y).sum()), + ("fused", lambda x, y: mx.fast.cross_entropy(x, y).sum()), +]: + logits = mx.random.normal(shape=(4096, 128256), scale=2.0).astype(mx.bfloat16) + targets = mx.random.randint(0, 128256, shape=(4096,)) + mx.eval(logits, targets) + mx.clear_cache() + mx.reset_peak_memory() + mx.eval(mx.grad(fn, argnums=0)(logits, targets)) + print(f" {label:>10}: {mx.get_peak_memory() / 2**30:.2f} GiB") diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index ea4a995ade..598b45fd75 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -39,6 +39,7 @@ if(MLX_METAL_JIT) target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jit_kernels.cpp) make_jit_source(arange) make_jit_source(copy) + make_jit_source(cross_entropy) make_jit_source(unary) make_jit_source(binary) make_jit_source(binary_two) @@ -124,6 +125,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy.cpp ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp diff --git a/mlx/backend/metal/cross_entropy.cpp b/mlx/backend/metal/cross_entropy.cpp new file mode 100644 index 0000000000..4b78cbd88f --- /dev/null +++ b/mlx/backend/metal/cross_entropy.cpp @@ -0,0 +1,157 @@ +// Copyright © 2026 Apple Inc. + +#include +#include + +#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/fast_primitives.h" + +namespace mlx::core::fast { + +namespace { + +// Above this the single pass kernel would need more than one read per thread, +// so switch to the looped variant. +constexpr int CROSS_ENTROPY_LOOPED_LIMIT = 4096; +constexpr int SIMD_SIZE = 32; +constexpr int N_READS = 4; + +size_t ceil_simd_multiple(size_t n) { + return SIMD_SIZE * ((n + SIMD_SIZE - 1) / SIMD_SIZE); +} + +} // namespace + +bool CrossEntropy::use_fallback(Stream s) { + return s.device == Device::cpu; +} + +void CrossEntropy::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + assert(inputs.size() == 2); // logits and targets + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + auto& out = outputs[0]; + + auto ensure_row_contiguous = [&s, &compute_encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } + array x_copy = contiguous_copy_gpu(x, s); + compute_encoder.add_temporary(x_copy); + return x_copy; + }; + + auto in = ensure_row_contiguous(inputs[0]); // [n_rows, V] + auto target = ensure_row_contiguous(inputs[1]); // [n_rows] + out.set_data(allocator::malloc(out.nbytes())); // [n_rows] in fp32 + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + auto get_kernel = [&d, &in](bool looped) { + std::string kernel_name = looped ? "looped_" : "block_"; + kernel_name += "cross_entropy_"; + kernel_name += type_to_name(in); + return get_cross_entropy_kernel(d, kernel_name, in); + }; + + bool looped = axis_size > CROSS_ENTROPY_LOOPED_LIMIT; + auto kernel = get_kernel(looped); + size_t threadgroup_size = 0; + if (!looped) { + threadgroup_size = ceil_simd_multiple((axis_size + N_READS - 1) / N_READS); + // The single pass kernel needs one thread per N_READS elements. If the + // pipeline cannot hold that many threads, fall back to the looped variant. + if (threadgroup_size > kernel->maxTotalThreadsPerThreadgroup()) { + looped = true; + kernel = get_kernel(true); + } + } + if (looped) { + threadgroup_size = kernel->maxTotalThreadsPerThreadgroup(); + } + + MTL::Size grid_dims(n_rows * threadgroup_size, 1, 1); + MTL::Size group_dims(threadgroup_size, 1, 1); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_input_array(target, 1); + compute_encoder.set_output_array(out, 2); + compute_encoder.set_bytes(axis_size, 3); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +void CrossEntropyVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + assert(inputs.size() == 4); // logits, targets, loss, cotangent + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + auto& out = outputs[0]; + + auto ensure_row_contiguous = [&s, &compute_encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } + array x_copy = contiguous_copy_gpu(x, s); + compute_encoder.add_temporary(x_copy); + return x_copy; + }; + + // The gradient has the same shape and type as the logits, so write it into + // the logits buffer whenever that buffer is ours to reuse. The kernel reads + // the target score into registers and barriers before overwriting anything. + auto set_output = [&s, &out](const array& x) { + if (x.flags().row_contiguous) { + if (x.is_donatable()) { + out.copy_shared_buffer(x); + } else { + out.set_data(allocator::malloc(out.nbytes())); + } + return x; + } + array x_copy = contiguous_copy_gpu(x, s); + out.copy_shared_buffer(x_copy); + return x_copy; + }; + + auto in = set_output(inputs[0]); // [n_rows, V] + auto target = ensure_row_contiguous(inputs[1]); // [n_rows] + auto loss = ensure_row_contiguous(inputs[2]); // [n_rows] fp32 + auto cotan = ensure_row_contiguous(inputs[3]); // [n_rows] fp32 + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + std::string kernel_name = "vjp_cross_entropy_"; + kernel_name += type_to_name(in); + auto kernel = get_cross_entropy_kernel(d, kernel_name, in); + + // The kernel loops over the row, so any threadgroup size is correct. Use + // just enough threads to cover the row without oversubscribing short rows. + size_t threadgroup_size = std::min( + static_cast(kernel->maxTotalThreadsPerThreadgroup()), + ceil_simd_multiple((axis_size + N_READS - 1) / N_READS)); + + MTL::Size grid_dims(n_rows * threadgroup_size, 1, 1); + MTL::Size group_dims(threadgroup_size, 1, 1); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_input_array(target, 1); + compute_encoder.set_input_array(loss, 2); + compute_encoder.set_input_array(cotan, 3); + compute_encoder.set_output_array(out, 4); + compute_encoder.set_bytes(axis_size, 5); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +} // namespace mlx::core::fast diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index 4fb1be1110..9fe1923e6f 100644 --- a/mlx/backend/metal/jit/includes.h +++ b/mlx/backend/metal/jit/includes.h @@ -18,6 +18,7 @@ const char* unary(); const char* binary(); const char* binary_two(); const char* copy(); +const char* cross_entropy(); const char* fft(); const char* gather_axis(); const char* gather_front(); diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 807a6550b7..6dd2cc163a 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -360,6 +360,28 @@ MTL::ComputePipelineState* get_logsumexp_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_cross_entropy_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in) { + std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1); + auto lib = d.get_library(lib_name, [&] { + // The library is keyed on the logits type: the loss is always float32. + auto t_str = get_type_string(in.dtype()); + std::string kernel_source; + kernel_source = metal::utils(); + kernel_source += metal::cross_entropy(); + kernel_source += + get_template_definition("block_" + lib_name, "cross_entropy", t_str); + kernel_source += get_template_definition( + "looped_" + lib_name, "cross_entropy_looped", t_str); + kernel_source += + get_template_definition("vjp_" + lib_name, "cross_entropy_vjp", t_str); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 42888f0a78..d9589e7fc8 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -64,6 +64,11 @@ MTL::ComputePipelineState* get_logsumexp_kernel( const std::string& kernel_name, const array& out); +MTL::ComputePipelineState* get_cross_entropy_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in); + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index ecabbda550..7cc90c803c 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -127,6 +127,7 @@ if(NOT MLX_METAL_JIT) build_kernel(binary binary.h binary_ops.h) build_kernel(binary_two binary_two.h) build_kernel(copy copy.h) + build_kernel(cross_entropy cross_entropy.h) build_kernel(fft fft.h fft/radix.h fft/readwrite.h) build_kernel( reduce diff --git a/mlx/backend/metal/kernels/cross_entropy.h b/mlx/backend/metal/kernels/cross_entropy.h new file mode 100644 index 0000000000..6cacc9cb3a --- /dev/null +++ b/mlx/backend/metal/kernels/cross_entropy.h @@ -0,0 +1,219 @@ +// Copyright © 2026 Apple Inc. + +// Fused logsumexp + gather. +// +// The accumulation is done in float32 regardless of the input type so that +// callers do not need to promote the logits with logits.astype(mx.float32). +// +// For each row: loss = logsumexp(x) - x[target] +// +// One threadgroup handles one row. The logsumexp is accumulated first, then a +// single thread does the gather and writes the loss. + +template +[[kernel]] void cross_entropy( + const device T* x, + const device int* y, + device float* loss, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint _lid [[thread_position_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + int lid = _lid; + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + AccT ld[N_READS]; + + const device T* row = x + gid * size_t(axis_size); + const device T* in = row + lid * N_READS; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + ld[i] = AccT(in[i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + ld[i] = + ((lid * N_READS + i) < axis_size) ? AccT(in[i]) : Limits::min; + } + } + if (simd_group_id == 0) { + local_max[simd_lane_id] = Limits::min; + local_normalizer[simd_lane_id] = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Get the max + AccT maxval = Limits::finite_min; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < ld[i]) ? ld[i] : maxval; + } + maxval = simd_max(maxval); + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group_id == 0) { + maxval = simd_max(local_max[simd_lane_id]); + if (simd_lane_id == 0) { + local_max[0] = maxval; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = local_max[0]; + + // Compute exp(x_i - maxval) and store the partial sums in local_normalizer + AccT normalizer = 0; + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(ld[i] - maxval); + } + normalizer = simd_sum(normalizer); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group_id == 0) { + normalizer = simd_sum(local_normalizer[simd_lane_id]); + if (simd_lane_id == 0) { + // Gather the score of the target class and subtract it from the lse. + AccT gap = maxval - AccT(row[y[gid]]); + loss[gid] = + static_cast(isinf(maxval) ? gap : log(normalizer) + gap); + } + } +} + +template +[[kernel]] void cross_entropy_looped( + const device T* x, + const device int* y, + device float* loss, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + const device T* row = x + gid * size_t(axis_size); + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + // The threadgroup may hold fewer than SIMD_SIZE simdgroups, so initialize + // every slot: the cross-simdgroup reduction below reads all of them. + if (simd_group_id == 0) { + local_max[simd_lane_id] = Limits::finite_min; + local_normalizer[simd_lane_id] = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Get the max and the normalizer in one go + AccT prevmax; + AccT maxval = Limits::finite_min; + AccT normalizer = 0; + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(row[offset + i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + vals[i] = (offset + i < axis_size) ? AccT(row[offset + i]) + : Limits::min; + } + } + prevmax = maxval; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < vals[i]) ? vals[i] : maxval; + } + normalizer *= fast::exp(prevmax - maxval); + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(vals[i] - maxval); + } + } + prevmax = maxval; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + normalizer = simd_sum(normalizer); + + prevmax = maxval; + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = simd_max(local_max[simd_lane_id]); + normalizer *= fast::exp(prevmax - maxval); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = simd_sum(local_normalizer[simd_lane_id]); + + if (lid == 0) { + AccT gap = maxval - AccT(row[y[gid]]); + loss[gid] = static_cast(isinf(maxval) ? gap : log(normalizer) + gap); + } +} + +// Gradient of the fused loss. The forward pass already folded the target score +// into the loss, so lse = loss + x[target] and the softmax probabilities come +// back as exp((x_i - x[target]) - loss). +template +[[kernel]] void cross_entropy_vjp( + const device T* x, + const device int* y, + const device float* loss, + const device float* cotan, + device T* grads, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]]) { + const device T* row = x + gid * size_t(axis_size); + device T* grads_row = grads + gid * size_t(axis_size); + + int y_n = y[gid]; + AccT g = AccT(cotan[gid]); + AccT loss_n = AccT(loss[gid]); + AccT x_t = AccT(row[y_n]); + + // grads aliases x when the logits buffer is donated, so hold every thread + // here until all of them have read the target score. + threadgroup_barrier(mem_flags::mem_device); + + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(row[offset + i]); + } + for (int i = 0; i < N_READS; i++) { + AccT p = fast::exp((vals[i] - x_t) - loss_n); + vals[i] = g * (p - ((offset + i) == y_n ? AccT(1) : AccT(0))); + } + for (int i = 0; i < N_READS; i++) { + grads_row[offset + i] = T(vals[i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + if (offset + i < axis_size) { + AccT v = AccT(row[offset + i]); + AccT p = fast::exp((v - x_t) - loss_n); + grads_row[offset + i] = + T(g * (p - ((offset + i) == y_n ? AccT(1) : AccT(0)))); + } + } + } + } +} diff --git a/mlx/backend/metal/kernels/cross_entropy.metal b/mlx/backend/metal/kernels/cross_entropy.metal new file mode 100644 index 0000000000..d3a2b63848 --- /dev/null +++ b/mlx/backend/metal/kernels/cross_entropy.metal @@ -0,0 +1,19 @@ +// Copyright © 2026 Apple Inc. + +#include +#include + +using namespace metal; + +// clang-format off +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/cross_entropy.h" + +#define instantiate_cross_entropy(name, itype) \ + instantiate_kernel("block_cross_entropy_" #name, cross_entropy, itype) \ + instantiate_kernel("looped_cross_entropy_" #name, cross_entropy_looped, itype) \ + instantiate_kernel("vjp_cross_entropy_" #name, cross_entropy_vjp, itype) \ + +instantiate_cross_entropy(float32, float) +instantiate_cross_entropy(float16, half) +instantiate_cross_entropy(bfloat16, bfloat16_t) // clang-format on diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index f1141e9792..e07a06f441 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -79,6 +79,13 @@ MTL::ComputePipelineState* get_logsumexp_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_cross_entropy_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index d1d0e781cc..68ba5fc2f5 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -215,26 +215,4 @@ void LUF::eval_gpu( throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI."); } -namespace fast { - -// There is no fused Metal cross entropy kernel yet -bool CrossEntropy::use_fallback(Stream s) { - return true; -} - -void CrossEntropy::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - throw std::runtime_error("[CrossEntropy::eval_gpu] Metal cross entropy NYI."); -} - -void CrossEntropyVJP::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - throw std::runtime_error( - "[CrossEntropyVJP::eval_gpu] Metal cross entropy NYI."); -} - -} // namespace fast - } // namespace mlx::core diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index b98d2765d6..8df7c29cd7 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -64,15 +64,15 @@ def cross_entropy( >>> nn.losses.cross_entropy(logits, targets) array([0.348587, 0.348587], dtype=float32) >>> - >>> # Half precision logits with class indices as targets. On CUDA a + >>> # Half precision logits with class indices as targets. On a GPU a >>> # fused kernel accumulates the reduction in float32: >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]], mx.bfloat16) >>> targets = mx.array([0, 1]) >>> nn.losses.cross_entropy(logits, targets) - array([0.0485873, 0.0485873], dtype=float32) + array([0.0485873, 0.0485873], dtype=bfloat16) >>> - >>> # Metal and the CPU reduce in the dtype of the logits, so upcast - >>> # them to get the same accuracy: + >>> # The cpu reduces in the dtype of the logits, so upcast them to get + >>> # the same accuracy: >>> nn.losses.cross_entropy(logits.astype(mx.float32), targets) array([0.0485873, 0.0485873], dtype=float32) """ @@ -95,9 +95,10 @@ def _drop_dim(shape, axis): f"Targets shape {targets.shape} does not match logits shape {logits.shape}." ) + # Both GPU backends implement the fused kernel, which shifts by the row max + # internally, so the explicit shift below is only needed off the fast path. use_fast = ( - mx.cuda.is_available() - and mx.default_device() == mx.gpu + mx.default_device() == mx.gpu and not targets_as_probs and label_smoothing == 0 and axis in (-1, logits.ndim - 1) diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 200781d372..0d48ef1be5 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -548,6 +548,95 @@ def cross_entropy_ref(logits, targets): self.assertEqual(out.shape, targets.shape) self.assertLess(mx.abs(out - expected).max().item(), tolerances[dtype]) + def test_cross_entropy_degenerate_rows(self): + def cross_entropy_ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits.astype(mx.float32), axis=-1) - score.astype( + mx.float32 + ) + + targets = mx.array([0, 1]) + for name, logits in [ + ("has +inf", mx.array([[1.0, float("inf"), 2.0], [1.0, 2.0, 3.0]])), + ( + "masked", + mx.array([[0.5, float("-inf"), 2.0], [float("-inf"), 1.0, 3.0]]), + ), + ]: + out = mx.fast.cross_entropy(logits, targets) + expected = cross_entropy_ref(logits, targets) + self.assertTrue(mx.allclose(out, expected).item(), msg=name) + + # An entirely masked row has no finite logsumexp, so the loss is nan + # both fused and unfused. + all_masked = mx.array([[float("-inf")] * 3, [1.0, 2.0, 3.0]]) + out = mx.fast.cross_entropy(all_masked, targets) + self.assertTrue(mx.isnan(out[0]).item()) + self.assertTrue(mx.isfinite(out[1]).item()) + + def test_cross_entropy_catastrophic_cancellation(self): + # Every logit in the row is equal, so the loss is exactly log(V). The + # fused kernel forms (max - x_target) before adding log(normalizer) and + # keeps that value, where computing logsumexp first would round it away. + # Only the fused kernel has this property, so there is nothing to check + # on the cpu, which always takes the decomposed fallback. + if mx.default_device() == mx.cpu: + self.skipTest("cross entropy is not fused on the cpu") + V = 4096 + logits = mx.full((2, V), 1e30) + targets = mx.array([0, 5]) + out = mx.fast.cross_entropy(logits, targets) + self.assertTrue(mx.allclose(out, mx.full((2,), math.log(V)), atol=1e-4).item()) + + def test_cross_entropy_non_contiguous(self): + def cross_entropy_ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits.astype(mx.float32), axis=-1) - score.astype( + mx.float32 + ) + + logits = mx.random.normal(shape=(1000, 8)).T + targets = mx.random.randint(0, 1000, shape=(8,)) + self.assertTrue( + mx.allclose( + mx.fast.cross_entropy(logits, targets), + cross_entropy_ref(logits, targets), + ).item() + ) + + logits = mx.random.normal(shape=(4, 2048))[:, ::2] + targets = mx.random.randint(0, 1024, shape=(4,)) + self.assertTrue( + mx.allclose( + mx.fast.cross_entropy(logits, targets), + cross_entropy_ref(logits, targets), + ).item() + ) + + def test_cross_entropy_grad_donated_input(self): + # The logits are produced inside the graph, so the vjp may write the + # gradient directly into that buffer. + def ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits, axis=-1) - score + + for V in [64, 4096, 32000]: + x = mx.random.normal(shape=(3, V), scale=2.0) + targets = mx.random.randint(0, V, shape=(3,)) + g1 = mx.grad(lambda a, y: ref(a * 2.0 + 1.0, y).sum(), argnums=0)( + x, targets + ) + g2 = mx.grad( + lambda a, y: mx.fast.cross_entropy(a * 2.0 + 1.0, y).sum(), argnums=0 + )(x, targets) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-5) + def test_cross_entropy_shape_checks(self): logits = mx.random.normal(shape=(4, 16)) with self.assertRaises(ValueError): diff --git a/python/tests/test_losses.py b/python/tests/test_losses.py index 96174f3c89..ffd86ea5c2 100644 --- a/python/tests/test_losses.py +++ b/python/tests/test_losses.py @@ -9,6 +9,68 @@ class TestLosses(mlx_tests.MLXTestCase): + def test_cross_entropy_gpu_matches_cpu(self): + # On a gpu the class index path is handed to the fused kernel, which + # shifts by the row max itself instead of using the explicit shift the + # fallback applies. The two have to agree, including for the arguments + # that fall back on either device. + if mx.default_device() == mx.cpu: + self.skipTest("needs a gpu to compare against the cpu") + + # cross_entropy picks the fast path off the default device, so switch + # the device rather than passing a stream. + def both_devices(*args, **kwargs): + gpu = nn.losses.cross_entropy(*args, **kwargs) + mx.set_default_device(mx.cpu) + try: + cpu = nn.losses.cross_entropy(*args, **kwargs) + mx.eval(cpu) + finally: + mx.set_default_device(mx.gpu) + mx.eval(gpu) + return gpu, cpu + + for V in [2, 7, 4096, 4097]: + logits = mx.random.normal(shape=(4, V), scale=3.0) + targets = mx.random.randint(0, V, shape=(4,)) + gpu, cpu = both_devices(logits, targets, reduction="none") + self.assertEqual(gpu.dtype, cpu.dtype) + self.assertTrue(mx.allclose(gpu, cpu, atol=1e-5), msg=f"V={V}") + + # In half precision the two are allowed to differ, because the + # fused kernel accumulates in float32 while the fallback reduces in + # the dtype of the logits. Hold it to the stronger property: the + # fused result is never further from the float32 answer. + for dtype in [mx.float16, mx.bfloat16]: + half = logits.astype(dtype) + reference = nn.losses.cross_entropy( + half.astype(mx.float32), targets, reduction="none" + ) + gpu, cpu = both_devices(half, targets, reduction="none") + self.assertEqual(gpu.dtype, dtype) + gpu_err = mx.abs(gpu.astype(mx.float32) - reference).max().item() + cpu_err = mx.abs(cpu.astype(mx.float32) - reference).max().item() + self.assertLessEqual(gpu_err, cpu_err + 1e-6, msg=f"V={V} {dtype}") + self.assertLess(gpu_err, 0.2, msg=f"V={V} {dtype}") + + # A large shared offset is where the fused and the decomposed paths + # could most easily disagree. + base = mx.array([[2.0, -1.0]]) + for offset in [0.0, 1e4, 1e6]: + gpu, cpu = both_devices(base + offset, mx.array([0]), reduction="none") + self.assertTrue(mx.allclose(gpu, cpu, atol=1e-5), msg=f"offset={offset}") + + # Arguments the fast path declines, so both devices decompose. + logits = mx.random.normal(shape=(4, 32)) + targets = mx.random.randint(0, 32, shape=(4,)) + for kwargs in [ + {"label_smoothing": 0.1}, + {"weights": mx.random.uniform(shape=(4,))}, + {"reduction": "mean"}, + ]: + gpu, cpu = both_devices(logits, targets, **kwargs) + self.assertTrue(mx.allclose(gpu, cpu, atol=1e-5), msg=str(kwargs)) + def test_cross_entropy(self): # No weights, no label smoothing logits = mx.array([[0.0, -float("inf")], [-float("inf"), 0.0]])