proper mac metal support for training

.githooks/prepare-commit-msg

@@ -0,0 +1,10 @@
+#!/bin/sh
+grep -qs "^Signed-off-by: " "$1" || {
+    name=$(git config user.name)
+    email=$(git config user.email)
+
+    if [ -n "$name" ] && [ -n "$email" ]; then
+        echo "" >> "$1"
+        echo "Signed-off-by: $name <$email>" >> "$1"
+    fi
+}

pyproject.toml

@@ -11,11 +11,13 @@ dependencies = [
     "transformers",
     "wandb",
     "sentencepiece",
+    "jax-metal; sys_platform == 'darwin'",
 ]
 
 [project.optional-dependencies]
 tpu = ["jax[tpu]"]
 gpu = ["jax[cuda12]"]
+metal = ["jax-metal"]
 
 [project.scripts]
 needle = "src.cli:main"

src/cli.py

@@ -78,7 +78,8 @@
     "num_dec_layers": 2,
     "max_enc_len": 128,
     "max_dec_len": 128,
-    "max_samples": 10000,
+    "max_samples": 10_000,
+    "max_eval_samples": 1024,
 }
 
 MAIN_CONFIG = {

src/model.py

@@ -1,10 +1,13 @@
 import math
 from dataclasses import dataclass
 
+import jax
 import jax.numpy as jnp
 import jax.nn.initializers as jinit
 import flax.linen as nn
 
+IS_MAC = jax.default_backend() == "METAL"
+
 
 def default_init():
     return jinit.normal(stddev=0.02)
@@ -213,6 +216,33 @@ def __call__(self, x, s, mask=None, rope=None):
         return x, s
 
 
+class DepthwiseConv1d(nn.Module):
+    """Depthwise 1D convolution with Metal fallback."""
+    features: int
+    kernel_size: int = 4
+    strides: int = 2
+    dtype: jnp.dtype = jnp.bfloat16
+
+    @nn.compact
+    def __call__(self, x):
+        if IS_MAC:
+            k = self.param("kernel", default_init(), (self.kernel_size, self.features))
+            k = k.astype(self.dtype)
+            T = x.shape[1]
+            pad_left = (self.kernel_size - 1) // 2
+            pad_right = self.kernel_size - 1 - pad_left
+            x_padded = jnp.pad(x, ((0, 0), (pad_left, pad_right), (0, 0)))
+            indices = jnp.arange(0, T, self.strides)[:, None] + jnp.arange(self.kernel_size)[None, :]
+            x_windows = x_padded[:, indices, :]
+            return jnp.sum(x_windows * k[None, None, :, :], axis=2)
+        return nn.Conv(
+            features=self.features, kernel_size=(self.kernel_size,),
+            strides=(self.strides,), padding='SAME',
+            feature_group_count=self.features,
+            dtype=self.dtype, use_bias=False,
+        )(x)
+
+
 class MemoryMixerEncoder(nn.Module):
     """Encoder using MemoryMixer blocks. Output is the final memory slots S."""
     config: TransformerConfig
@@ -228,17 +258,15 @@ def __call__(self, x, mask=None, rope=None):
         s_bias = nn.Dense(cfg.d_model, dtype=dt, use_bias=False, kernel_init=jinit.zeros, name="slot_init")(x_pool)
         s = jnp.broadcast_to(s_base.astype(dt), (x.shape[0], cfg.num_memory_slots, cfg.d_model)) + s_bias[:, None, :]
 
-        x = nn.Conv(features=cfg.d_model, kernel_size=(4,), strides=(2,),
-                    padding='SAME', feature_group_count=cfg.d_model,
-                    dtype=dt, use_bias=False, name="downsample_dw")(x)
+        x = DepthwiseConv1d(cfg.d_model, kernel_size=4, strides=2, dtype=dt, name="downsample_dw")(x)
         x = nn.Dense(cfg.d_model, dtype=dt, use_bias=False,
                      kernel_init=default_init(), name="downsample_pw")(x)
 
         if mask is not None:
             T_new = x.shape[1]
             if mask.shape[-1] % 2:
                 mask = jnp.pad(mask, ((0, 0), (0, 0), (0, 0), (0, 1)))
-            mask = mask.reshape(mask.shape[0], 1, 1, -1, 2).any(axis=-1)
+            mask = mask.reshape(mask.shape[0], 1, 1, -1, 2).any(axis=-1) if not IS_MAC else mask.reshape(mask.shape[0], 1, 1, -1, 2).max(axis=-1)
             mask = mask[..., :T_new]
 
         for i in range(cfg.num_encoder_layers):
@@ -361,10 +389,12 @@ def forward_with_aux(self, src, tgt, src_mask=None, tgt_mask=None, cross_mask=No
 
 
 def make_causal_mask(seq_len):
-    mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))
+    mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_ if not IS_MAC else jnp.int32))
     return mask[None, None, :, :]
 
 
 def make_padding_mask(tokens, pad_token_id):
     mask = tokens != pad_token_id
+    if IS_MAC:
+        mask = mask.astype(jnp.int32)
     return mask[:, None, None, :]

src/run.py

@@ -64,7 +64,7 @@ def decode_step(dec_buffer, encoder_out, src_mask):
     for i in range(max_gen_len - 1):
         next_logits = logits[0, i] / temperature
         rng, sample_rng = jax.random.split(rng)
-        next_token = jax.random.categorical(sample_rng, next_logits).item()
+        next_token = int(jax.random.categorical(sample_rng, next_logits))
 
         if next_token == eos_id:
             break

src/test.py

@@ -93,7 +93,7 @@ def decode_step(dec_buffer, encoder_out, src_mask):
         num_tokens = 0
         for i in range(max_gen_len - 1):
             rng, sample_rng = jax.random.split(rng)
-            next_token = jax.random.categorical(sample_rng, logits[0, i]).item()
+            next_token = int(jax.random.categorical(sample_rng, logits[0, i]))
 
             if next_token == eos_id:
                 break

src/train.py

@@ -726,16 +726,17 @@ def train(args):
 
         eval_params = jax_utils.unreplicate(ema_params)
         val_causal = make_causal_mask(args.max_dec_len)
+        num_val_batches = len(val_enc) // args.batch_size
         total_loss, total_toks = 0.0, 0.0
-        for vb in get_batches(val_enc, val_dec_in, val_dec_tgt, args.batch_size, shuffle=False):
+        for vb in tqdm(get_batches(val_enc, val_dec_in, val_dec_tgt, args.batch_size, shuffle=False), total=num_val_batches, desc="  Val", leave=False):
             vl, vt = val_loss_fn(eval_params, vb[0], vb[1], vb[2], val_causal)
             total_loss += float(vl)
             total_toks += float(vt)
         last_val_ppl = float(math.exp(total_loss / max(total_toks, 1)))
 
         q_params = _quantize_params(eval_params, group_size=_GROUP_SIZE)
         q_total_loss, q_total_toks = 0.0, 0.0
-        for vb in get_batches(val_enc, val_dec_in, val_dec_tgt, args.batch_size, shuffle=False):
+        for vb in tqdm(get_batches(val_enc, val_dec_in, val_dec_tgt, args.batch_size, shuffle=False), total=num_val_batches, desc="  Quant val", leave=False):
             vl, vt = val_loss_fn(q_params, vb[0], vb[1], vb[2], val_causal)
             q_total_loss += float(vl)
             q_total_toks += float(vt)
@@ -749,7 +750,7 @@ def train(args):
             for d_prime in _MRL_DIMS:
                 mrl_vl_fn = _make_mrl_val_loss_fn(apply_fn, d_prime)
                 mrl_total_loss, mrl_total_toks = 0.0, 0.0
-                for vb in get_batches(val_enc, val_dec_in, val_dec_tgt, args.batch_size, shuffle=False):
+                for vb in tqdm(get_batches(val_enc, val_dec_in, val_dec_tgt, args.batch_size, shuffle=False), total=num_val_batches, desc=f"  MRL d={d_prime}", leave=False):
                     vl, vt = mrl_vl_fn(eval_params, vb[0], vb[1], vb[2], val_causal)
                     mrl_total_loss += float(vl)
                     mrl_total_toks += float(vt)
@@ -769,13 +770,17 @@ def train(args):
         ckpt_name = f"needle_{args.num_layers}_{args.d_model}_{global_step}.pkl"
         ckpt_path = os.path.join(args.checkpoint_dir, ckpt_name)
 
+        print(f"  Saving checkpoint...", flush=True)
         with open(ckpt_path, "wb") as f:
             pickle.dump({"params": params_np, "config": config.__dict__}, f)
 
         from .test import measure_throughput, benchmark_generation_quality
         eval_params_jnp = jax.tree.map(jnp.array, params_np)
         del params_np  # free CPU memory
+        # Sync device before benchmarking to flush pending async operations
+        jax.tree.map(lambda x: x.block_until_ready(), eval_params_jnp)
         model = EncoderDecoderTransformer(config)
+        print(f"  Benchmarking...", flush=True)
         tp = measure_throughput(model, eval_params_jnp, tokenizer, num_runs=5)
         prompts = ["Once upon a time", "The little dog", "She was very happy because"]
         quality = benchmark_generation_quality(model, eval_params_jnp, tokenizer, prompts, max_gen_len=64, temperature=0.8)