diff --git a/.githooks/prepare-commit-msg b/.githooks/prepare-commit-msg new file mode 100755 index 0000000..7006ca5 --- /dev/null +++ b/.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 +} diff --git a/pyproject.toml b/pyproject.toml index 13ebd8e..7b55d41 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/cli.py b/src/cli.py index 5c32413..2c55009 100644 --- a/src/cli.py +++ b/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 = { diff --git a/src/model.py b/src/model.py index a60fa03..3a052d6 100644 --- a/src/model.py +++ b/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,9 +258,7 @@ 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) @@ -238,7 +266,7 @@ def __call__(self, x, mask=None, rope=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, :] diff --git a/src/run.py b/src/run.py index df60452..b186586 100644 --- a/src/run.py +++ b/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 diff --git a/src/test.py b/src/test.py index 86d3598..17a6cfe 100644 --- a/src/test.py +++ b/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 diff --git a/src/train.py b/src/train.py index ccf9a5a..36ff6af 100644 --- a/src/train.py +++ b/src/train.py @@ -726,8 +726,9 @@ 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) @@ -735,7 +736,7 @@ def train(args): 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)