Merge pull request #31 from ausimian/m16-mixed-precision

M16: mixed-precision training

Author: ausimian

PLAN.md

@@ -809,6 +809,31 @@ Promotes "mixed-precision master weights" out of v1 non-goals.
 0.5% of the f32 baseline; loss-scaling primitives documented with a
 worked example in the moduledoc.
 
+**Shipped.**
+
+- `Emily.MixedPrecision` (`lib/emily/mixed_precision.ex`) — `cast_params/2`,
+  `accumulate_grad/2`, `loss_scale/1`, `scale_loss/2`, `unscale/2`,
+  `update/2`, `has_overflow?/1`. Nested `LossScaler` struct with dynamic
+  scaling (halve on overflow, double every N successful steps, floor at
+  `min_scale`). Moduledoc includes a complete mixed-precision training
+  step as a worked example. Traversal covers maps, tuples, lists, and
+  `%Nx.Tensor{}` leaves; `Nx.Container` structs (e.g. `Axon.ModelState`)
+  are not traversed — documented.
+- bf16 grad equivalence (`test/emily/grad/bf16_grad_equivalence_test.exs`)
+  — all 8 zoo functions pass under Emily.Compiler vs BinaryBackend
+  Evaluator, both in bf16, within 1e-2 tolerance.
+- Mixed-precision MLP curve-matching
+  (`test/emily/training/bf16_mlp_curve_test.exs`) — 50-step training
+  loop with f32 master weights, bf16 forward pass, loss scaling, and f32
+  gradient accumulation. Emily vs BinaryBackend within rtol 5e-2.
+- bf16 MNIST convergence canary
+  (`test/emily/training/mnist_bf16_full_test.exs`) — `:training_full`,
+  uses `Axon.MixedPrecision.create_policy`. Target ≥ 95.5%.
+- Backend `coerce` fix: `Emily.Backend.wrap` now checks `Native.dtype(ref)`
+  and casts if the MLX buffer dtype disagrees with the declared Nx output
+  type. Previously only handled pred→u8; now handles all type mismatches
+  (e.g. bf16 buffer with f32 metadata from `Nx.Defn.grad` type promotion).
+
 ### M17 — Conv-pool training (was M10)
 
 Originally scoped as M10 in the pre-review plan. Re-prioritized below

RELEASE.md

@@ -8,6 +8,24 @@
 
 ## Added
 
+- M16 — Mixed-precision training. `Emily.MixedPrecision` delivers the
+  standard bf16 recipe: `cast_params/2` (downcast f32 → bf16 for the
+  forward pass), `accumulate_grad/2` (upcast bf16 grads → f32 for the
+  optimizer), `loss_scale/1` / `scale_loss/2` / `unscale/2` / `update/2`
+  (dynamic loss scaling with overflow detection). `LossScaler` struct
+  halves the scale on inf/nan overflow, doubles every N successful steps,
+  floors at a configurable minimum. Moduledoc includes a complete
+  worked example.
+  - **Backend `coerce` fix**: `Emily.Backend.wrap` now checks
+    `Native.dtype(ref)` and casts when the MLX buffer dtype disagrees
+    with the declared Nx output type. Fixes bf16 grads where
+    `Nx.Defn.grad` promotes the output type metadata to f32 but the
+    MLX buffer stays bf16.
+  - **Tests**: `mixed_precision_test.exs` (33 unit tests), bf16 grad
+    equivalence for all 8 zoo functions, bf16 mixed-precision MLP
+    curve-matching (50 steps, rtol 5e-2), bf16 MNIST convergence canary
+    (`:training_full`, target ≥ 95.5%).
+
 - M15 — Native linalg. `lu`, `svd`, `qr` (reduced), `cholesky`, `eigh`,
   `solve`, and `triangular_solve` now dispatch directly to
   `mx::linalg::*` instead of round-tripping through `Nx.BinaryBackend`.

lib/emily/backend.ex

@@ -55,8 +55,17 @@ defmodule Emily.Backend do
     %{out | data: %B{ref: coerce(ref, type, w)}}
   end
 
+  # Fast path: pred→u8 is the most common mismatch (MLX comparison/logical
+  # ops yield bool; Nx expects u8). The general path covers bf16→f32 type
+  # promotion from Nx.Defn.grad and any other MLX/Nx dtype disagreement.
   defp coerce(ref, {:u, 8}, w), do: Native.astype(w, ref, {:u, 8})
-  defp coerce(ref, _, _w), do: ref
+
+  defp coerce(ref, type, w) do
+    case Native.dtype(ref) do
+      ^type -> ref
+      _ -> Native.astype(w, ref, type)
+    end
+  end
 
   defp shape_list(shape) when is_tuple(shape), do: Tuple.to_list(shape)
 

lib/emily/mixed_precision.ex

@@ -0,0 +1,215 @@
+defmodule Emily.MixedPrecision do
+  @moduledoc """
+  Mixed-precision training utilities.
+
+  Standard recipe for memory-efficient training: bf16 activations with
+  f32 master weights and dynamic loss scaling. Keeps the full-precision
+  copy of parameters for numerically stable optimizer updates while
+  running the forward and backward pass in half precision.
+
+  ## Worked example
+
+      alias Emily.MixedPrecision, as: MP
+      alias Emily.MixedPrecision.LossScaler
+
+      # f32 master weights — the optimizer's ground truth.
+      master_params = init_params()
+      scaler = MP.loss_scale()
+
+      for {x, y} <- batches, reduce: {master_params, scaler} do
+        {params, scaler} ->
+          # Forward pass in bf16.
+          bf16_params = MP.cast_params(params, {:bf, 16})
+
+          # Backward pass: grad w.r.t. f32 master params, but the
+          # forward graph runs in bf16 thanks to the as_type casts
+          # inside the closure.
+          grads =
+            Nx.Defn.grad(params, fn p ->
+              p
+              |> MP.cast_params({:bf, 16})
+              |> forward(x, y)
+              |> MP.scale_loss(scaler)
+            end)
+
+          # Unscale, detect overflow, adjust scaler.
+          {grads, overflow?} = MP.unscale(grads, scaler)
+          scaler = MP.update(scaler, overflow?)
+
+          if overflow? do
+            {params, scaler}
+          else
+            f32_grads = MP.accumulate_grad(grads, {:f, 32})
+            {sgd_step(params, f32_grads, lr), scaler}
+          end
+      end
+
+  ## Container traversal
+
+  `cast_params/2`, `accumulate_grad/2`, and `has_overflow?/1` traverse
+  plain maps, tuples, and lists of `Nx.Tensor` leaves. For
+  `Axon.ModelState`, access the `.data` field first:
+
+      MP.cast_params(model_state.data, {:bf, 16})
+  """
+
+  defmodule LossScaler do
+    @moduledoc """
+    Dynamic loss-scaler state for mixed-precision training.
+
+    Tracks the current scale factor and the number of consecutive
+    successful (non-overflow) steps. On overflow the scale is halved;
+    after `growth_interval` successful steps it doubles.
+    """
+
+    @default_scale 65_536.0
+
+    @enforce_keys [:scale]
+    defstruct scale: @default_scale,
+              growth_factor: 2.0,
+              backoff_factor: 0.5,
+              growth_interval: 2000,
+              min_scale: 1.0,
+              counter: 0
+
+    @doc "Create a new loss scaler."
+    def new(opts \\ []) do
+      fields =
+        opts
+        |> Keyword.validate!([
+          :scale,
+          :growth_factor,
+          :backoff_factor,
+          :growth_interval,
+          :min_scale
+        ])
+        |> then(&Keyword.merge(Map.to_list(%__MODULE__{scale: @default_scale}), &1))
+
+      struct!(__MODULE__, fields)
+    end
+  end
+
+  @doc """
+  Downcast float tensors in a nested structure to `type`.
+
+  Integer and predicate tensors are left unchanged.
+  """
+  def cast_params(params, type), do: deep_apply(params, &Nx.as_type(&1, type))
+
+  @doc """
+  Upcast float tensors in a nested gradient structure to `type`.
+
+  Semantically identical to `cast_params/2` — exists for readability
+  at the call site (the direction of the cast is part of the name).
+  """
+  def accumulate_grad(grads, type), do: deep_apply(grads, &Nx.as_type(&1, type))
+
+  @doc """
+  Create a new dynamic loss scaler.
+
+  ## Options
+
+    * `:scale` — initial scale factor (default `65_536.0`)
+    * `:growth_factor` — multiply scale by this on growth (default `2.0`)
+    * `:backoff_factor` — multiply scale by this on overflow (default `0.5`)
+    * `:growth_interval` — successful steps before growing (default `2000`)
+    * `:min_scale` — floor for the scale (default `1.0`)
+  """
+  def loss_scale(opts \\ []), do: LossScaler.new(opts)
+
+  @doc """
+  Scale the loss by the scaler's current factor.
+
+  Call this inside the function passed to `Nx.Defn.grad` so that the
+  backward pass produces scaled gradients.
+  """
+  def scale_loss(loss, %LossScaler{scale: scale}) do
+    Nx.multiply(loss, scale)
+  end
+
+  @doc """
+  Unscale gradients and detect overflow.
+
+  Divides every float tensor in `grads` by `scaler.scale`, then checks
+  for inf/nan. Returns `{unscaled_grads, overflow?}`.
+  """
+  def unscale(grads, %LossScaler{scale: scale}) do
+    inv_scale = 1.0 / scale
+    unscaled = deep_apply(grads, &Nx.multiply(&1, inv_scale))
+    {unscaled, has_overflow?(unscaled)}
+  end
+
+  @doc """
+  Update the scaler after a training step.
+
+  On overflow: halves the scale (floored at `min_scale`), resets the
+  counter. On success: increments the counter; doubles the scale after
+  `growth_interval` consecutive successes.
+  """
+  def update(%LossScaler{} = scaler, true = _overflow) do
+    %{scaler | scale: max(scaler.min_scale, scaler.scale * scaler.backoff_factor), counter: 0}
+  end
+
+  def update(%LossScaler{} = scaler, false = _overflow) do
+    counter = scaler.counter + 1
+
+    if counter >= scaler.growth_interval do
+      %{scaler | scale: scaler.scale * scaler.growth_factor, counter: 0}
+    else
+      %{scaler | counter: counter}
+    end
+  end
+
+  @doc """
+  Check whether any tensor in a nested structure contains inf or nan.
+  """
+  def has_overflow?(structure), do: deep_overflow?(structure)
+
+  defp float_type?({kind, _}) when kind in [:f, :bf], do: true
+  defp float_type?(_), do: false
+
+  defp deep_apply(%Nx.Tensor{} = t, fun) do
+    if float_type?(t.type), do: fun.(t), else: t
+  end
+
+  defp deep_apply(map, fun) when is_map(map) and not is_struct(map) do
+    Map.new(map, fn {k, v} -> {k, deep_apply(v, fun)} end)
+  end
+
+  defp deep_apply(tuple, fun) when is_tuple(tuple) do
+    tuple |> Tuple.to_list() |> Enum.map(&deep_apply(&1, fun)) |> List.to_tuple()
+  end
+
+  defp deep_apply(list, fun) when is_list(list) do
+    Enum.map(list, &deep_apply(&1, fun))
+  end
+
+  defp deep_apply(other, _fun), do: other
+
+  defp deep_overflow?(%Nx.Tensor{} = t) do
+    if float_type?(t.type) do
+      t
+      |> Nx.is_nan()
+      |> Nx.logical_or(Nx.is_infinity(t))
+      |> Nx.any()
+      |> Nx.backend_transfer(Nx.BinaryBackend)
+      |> Nx.to_number() == 1
+    else
+      false
+    end
+  end
+
+  defp deep_overflow?(map) when is_map(map) and not is_struct(map) do
+    Enum.any?(map, fn {_k, v} -> deep_overflow?(v) end)
+  end
+
+  defp deep_overflow?(tuple) when is_tuple(tuple) do
+    tuple |> Tuple.to_list() |> Enum.any?(&deep_overflow?/1)
+  end
+
+  defp deep_overflow?(list) when is_list(list) do
+    Enum.any?(list, &deep_overflow?/1)
+  end
+
+  defp deep_overflow?(_), do: false
+end

test/emily/grad/bf16_grad_equivalence_test.exs

@@ -0,0 +1,48 @@
+defmodule Emily.Grad.Bf16EquivalenceTest do
+  @moduledoc """
+  bf16 gradient equivalence tests (M16).
+
+  For each function in the grad zoo, casts the fixed inputs to bf16 and
+  asserts that `grad` computed under `Emily.Compiler` matches the same
+  grad computed under `Nx.Defn.Evaluator` on `Nx.BinaryBackend` — both
+  sides in bf16. The oracle is BinaryBackend's software bf16, not f32.
+
+  bf16 has ~3 decimal digits of precision, so tolerances are 1e-2
+  (matching `BackendGenerators.tol_for({:bf, _})`).
+
+  ## Note on grad type promotion
+
+  `Nx.Defn.grad` may promote the output type from bf16 to f32 in
+  certain backward ops. Both the Emily and BinaryBackend results are
+  cast to f32 before comparison to normalise the binary
+  representation — the bf16-level tolerance still applies because
+  the computation itself ran in bf16.
+  """
+
+  use ExUnit.Case, async: true
+
+  import Emily.BackendGenerators, only: [assert_close: 3, to_emily: 1]
+  import Emily.GradZoo
+
+  @bf16_tol 1.0e-2
+
+  for name <- Emily.GradZoo.all_functions() do
+    test "#{name} — bf16 grad matches BinaryBackend bf16 oracle" do
+      name = unquote(name)
+      fun = grad_function(name)
+      inputs = fixed_inputs_bf16(name)
+
+      emily_inputs = Enum.map(inputs, &to_emily/1)
+
+      emily =
+        Nx.Defn.jit_apply(fun, emily_inputs, compiler: Emily.Compiler)
+        |> Nx.as_type({:f, 32})
+
+      ref =
+        Nx.Defn.jit_apply(fun, inputs, compiler: Nx.Defn.Evaluator)
+        |> Nx.as_type({:f, 32})
+
+      assert_close(emily, ref, tol: @bf16_tol)
+    end
+  end
+end