Create README.md

README.md

@@ -0,0 +1,594 @@
+"""
+# RetNPhi: Byte-Level Hybrid of Phi-3.5 and RetNet
+
+RetNPhi is an experimental architecture that transforms Phi-3.5 into a byte-level language model, incorporating RetNet-inspired mechanisms. This innovative approach enables the model to process raw byte sequences, allowing for universal file type handling.
+
+## Key Features:
+
+1. **Byte-Level Processing**: Operates directly on raw byte sequences, enabling universal application to any file type.
+2. **RetNet Integration**: Incorporates RetNet's multi-scale exponential decay and group normalization for efficient long-range dependency modeling.
+3. **Dual-mode Processing**: Supports parallel mode for efficient training and recurrent mode for inference.
+4. **Selective Fine-tuning**: Trains only specific layers (e.g., token embedding, post-attention layer normalizations) while keeping most of the original Phi-3.5 weights frozen.
+5. **Weight-Decomposed Low-Rank Adaptation (DoRA)**: Applies DoRA to self-attention output projections for efficient adaptation while preserving pretrained knowledge.
+
+## Implementation Strategy:
+
+- **Weight Reuse**: Utilizes frozen weights from the original Phi-3.5 model for most layers.
+- **Flexible DoRA Application**: Allows configuration of which layers and targets to apply DoRA.
+- **Configurable Architecture**: Supports both retention-based and original attention mechanisms.
+- **Untied Embeddings Option**: Provides the ability to use separate input and output embeddings.
+
+## Training and Inference:
+
+- Implements efficient training loops with customizable learning rate schedules.
+- Supports both training from scratch and fine-tuning from a checkpoint.
+- Provides a generation function for text completion tasks.
+
+## Goals:
+
+- Explore the potential of retention-like mechanisms in a byte-level Phi architecture.
+- Leverage dual-mode processing for efficient training and inference.
+- Develop a universal model capable of processing any file type.
+
+Note: This is a highly experimental implementation, designed for research and exploration rather than production use. It demonstrates the potential of combining pretrained models with novel architectures and efficient fine-tuning techniques.
+
+Author: Josef Albers
+Date: Aug 28, 2024
+"""
+
+import glob
+import json
+import math
+import time
+from datetime import datetime
+from types import SimpleNamespace
+
+import fire
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optim
+import numpy as np
+from huggingface_hub import snapshot_download
+from mlx.utils import tree_flatten, tree_unflatten
+
+from datasets import load_dataset
+
+class Tokenizer:
+    def __init__(self, file_path=None):
+        if file_path is None:
+            self.vocab = list(range(256))
+        else:
+            with open(file_path, 'r') as f:
+                content = f.read().lower().encode('utf-8')
+            self.vocab = sorted(set(content))
+        self.vocab_size = len(self.vocab)
+        self.byte_to_index = {byte: index for index, byte in enumerate(self.vocab)}
+        self.index_to_byte = {index: byte for index, byte in enumerate(self.vocab)}
+
+    def encode(self, text):
+        byte_seq = text.encode('utf-8')
+        return [self.byte_to_index[byte] for byte in byte_seq]
+
+    def decode(self, indices):
+        byte_seq = bytes(self.index_to_byte[index] for index in indices)
+        return byte_seq.decode('utf-8', errors='ignore')
+
+class SuRoPE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dim = config.hidden_size // config.num_attention_heads
+        self.original_max_position_embeddings = config.original_max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.scaling_factor = math.sqrt(1 + math.log(config.max_position_embeddings / config.original_max_position_embeddings) / math.log(config.original_max_position_embeddings))
+        self._long_factor = mx.array(config.rope_scaling["long_factor"], dtype=mx.float32)
+        self._short_factor = mx.array(config.rope_scaling["short_factor"], dtype=mx.float32)
+
+    def __call__(self, q, k, position_ids):
+        cos, sin = self._get_cos_sin(position_ids)
+        q = (q * cos) + (self._rotate_half(q) * sin)
+        k = (k * cos) + (self._rotate_half(k) * sin)
+        return q, k
+
+    def _get_cos_sin(self, position_ids):
+        su_factor = self._short_factor
+        position_ids_expanded = position_ids[:, None, :]
+        inv_freq = 1.0 / (su_factor * self.rope_theta**(mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim))
+        inv_freq_expanded = mx.repeat(inv_freq[None, :, None], position_ids.shape[0], axis=0)
+        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(0, 2, 1)
+        emb = mx.concatenate([freqs, freqs], axis=-1)
+        cos = mx.expand_dims(mx.cos(emb) * self.scaling_factor, axis=1)
+        sin = mx.expand_dims(mx.sin(emb) * self.scaling_factor, axis=1)
+        return cos, sin
+
+    def _rotate_half(self, x):
+        midpoint = x.shape[-1] // 2
+        x1, x2 = x[..., :midpoint], x[..., midpoint:]
+        return mx.concatenate([-x2, x1], axis=-1)
+
+class Phi3Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        dim = config.hidden_size
+        self.n_heads = n_heads = config.num_attention_heads
+        self.n_kv_heads = n_kv_heads = config.num_key_value_heads
+        self.num_hidden_layers = config.num_hidden_layers
+        self.head_dim = head_dim = config.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+        chop_1 = self.n_heads * self.head_dim
+        chop_2 = chop_1 + self.n_kv_heads * self.head_dim
+        self.chop = [chop_1, chop_2]
+        op_size = n_heads * head_dim + 2 * (n_kv_heads * head_dim)
+        self.qkv_proj = nn.Linear(dim, op_size, bias=False)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+        self.rope = SuRoPE(config)
+
+    def __call__(self, x, position_ids, attention_mask, cache, use_recurrent_mode):
+        B, L, _ = x.shape
+        qkv = self.qkv_proj(x)
+        q, k, v = mx.split(qkv, self.chop, axis=-1)
+        q = q.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        k = k.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        v = v.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        if cache is None:
+            position_ids = mx.arange(q.shape[2], dtype=mx.float32)[None] if position_ids is None else position_ids
+            q, k = self.rope(q,k,position_ids)
+            mask = mx.triu(mx.full((v.shape[2], v.shape[2]), -mx.inf), k=1)
+            if attention_mask is not None:
+                mask += mx.where(attention_mask[:, :, None]*attention_mask[:, None, :]==1, 0, -mx.inf)
+                mask = mx.expand_dims(mask, 1)
+            else:
+                mask = mask[None, None]
+        else:
+            past_k, past_v, past_p, past_m = cache
+            position_ids = past_p[:,-1:]+1
+            mask = mx.pad(past_m[:,:,-1:,:], ((0,0),(0,0),(0,0),(0,1)))
+            q, k = self.rope(q, k, position_ids)
+            k = mx.concatenate([past_k, k], axis=2)
+            v = mx.concatenate([past_v, v], axis=2)
+        cache = (k, v, position_ids, mask)
+        w = (q * self.scale) @ k.transpose(0, 1, 3, 2)
+        w += mask
+        w = mx.softmax(w, axis=-1)
+        o = w @ v
+        o = o.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(o).astype(x.dtype), cache
+
+class Phi3Retention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dim = dim = config.hidden_size
+        self.n_heads = n_heads = config.num_attention_heads
+        self.n_kv_heads = n_kv_heads = config.num_key_value_heads
+        self.num_hidden_layers = config.num_hidden_layers
+        self.head_dim = head_dim = config.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+        chop_1 = self.n_heads * self.head_dim
+        chop_2 = chop_1 + self.n_kv_heads * self.head_dim
+        self.chop = [chop_1, chop_2]
+        op_size = n_heads * head_dim + 2 * (n_kv_heads * head_dim)
+        self.qkv_proj = nn.Linear(dim, op_size, bias=False)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+        self.rope = SuRoPE(config)
+        xmin, xmax = math.log(1 / 32), math.log(1 / 512)
+        x = mx.linspace(xmin, xmax, num=n_heads)
+        self._gamma =  1 - x.exp()
+        self.gn = nn.GroupNorm(num_groups=head_dim, dims=-1, affine=False)
+
+    def __call__(self, x, position_ids, attention_mask, cache, use_recurrent_mode):
+        if use_recurrent_mode:
+            return self.recurrent_mode(x, cache)
+        B, L, _ = x.shape
+        qkv = self.qkv_proj(x)
+        q, k, v = mx.split(qkv, self.chop, axis=-1)
+        q = q.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        k = k.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        v = v.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        position_ids = mx.arange(q.shape[2], dtype=mx.float32)[None] if position_ids is None else position_ids
+        q, k = self.rope(q,k,position_ids)
+        cache = None
+        w = (q * self.scale) @ k.transpose(0, 1, 3, 2)
+        w = w * self._decay(L)
+        o = w @ v
+        o = o.transpose(0, 2, 1, 3).reshape(B*L, -1)
+        o = self.gn(o).reshape(B, L, -1)
+        return self.o_proj(o).astype(x.dtype), cache
+
+    def recurrent_mode(self, x, cache):
+        if cache is None:
+            s = mx.zeros((1, 32, 96, 96))
+            n = 0
+        else:
+            s, n = cache
+        qkv = self.qkv_proj(x)
+        q, k, v = mx.split(qkv, self.chop, axis=-1)
+        q = q.reshape(1, 1, self.n_heads, -1).transpose(0, 2, 1, 3)
+        k = k.reshape(1, 1, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        v = v.reshape(1, 1, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        position_ids = mx.array([[n]])
+        q, k = self.rope(q,k,position_ids)
+        k = k * self.scale
+        s = self._gamma[None, :, None, None] * s + (k.transpose(0, 1, 3, 2) @ v)
+        o = q @ s
+        o = o.transpose(0, 2, 1, 3).reshape(1, -1)
+        o = self.gn(o).reshape(1, 1, -1)
+        o = self.o_proj(o).astype(x.dtype)
+        return o, (s, n+1)
+
+    def _decay(self, sequence_length):
+        n = mx.arange(sequence_length)[:,None]
+        m = mx.arange(sequence_length)[None]
+        D = (self._gamma[:, None, None] ** (n-m)) * (n >= m)
+        return D
+
+class Phi3MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def __call__(self, x):
+        x = self.gate_up_proj(x)
+        gate, x = mx.split(x, 2, axis=-1)
+        return self.down_proj(nn.silu(gate) * x)
+
+class Phi3DecoderLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        if config.use_retention:
+            self.self_attn = Phi3Retention(config)
+        else:
+            self.self_attn = Phi3Attention(config)
+        self.mlp = Phi3MLP(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(self, x, position_ids, attention_mask, cache, use_recurrent_mode):
+        r, cache = self.self_attn(self.input_layernorm(x), position_ids, attention_mask, cache, use_recurrent_mode)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r, cache
+
+class Phi3Model(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embed_new = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [Phi3DecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(self, input_ids, pixel_values, image_sizes, position_ids, attention_mask, cache, use_recurrent_mode):
+        x = self.embed_new(input_ids)
+        cache = [None]*len(self.layers) if cache is None else cache
+        for i, l in enumerate(self.layers):
+            x, cache[i] = l(x, position_ids, attention_mask, cache[i], use_recurrent_mode)
+        return self.norm(x), cache
+
+class Phi3ForCausalLM(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.model = Phi3Model(config)
+        if config.untie_embedding:
+            self.lm_new = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+            self.untie = True
+        else:
+            self.untie = False
+
+    def __call__(self, input_ids, pixel_values=None, image_sizes=None, position_ids=None, attention_mask=None, cache=None, use_recurrent_mode=False):
+        x, cache = self.model(input_ids, pixel_values, image_sizes, position_ids, attention_mask, cache, use_recurrent_mode)
+        if self.untie:
+            return self.lm_new(x), cache
+        return self.model.embed_new.as_linear(x), cache
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+class DoRALinear(nn.Module):
+    @staticmethod
+    def from_linear(linear, r, alpha, scale, dropout):
+        output_dims, input_dims = linear.weight.shape
+        if isinstance(linear, nn.QuantizedLinear):
+            input_dims *= 32 // linear.bits
+        lora_lin = DoRALinear(input_dims=input_dims, output_dims=output_dims, r=r, alpha=alpha, scale=scale, dropout=dropout)
+        lora_lin.linear = linear
+        return lora_lin
+
+    def __init__(self, input_dims, output_dims, r, alpha, scale, dropout, bias=False):
+        super().__init__()
+        self.linear = nn.Linear(input_dims, output_dims, bias=bias)
+        self.dropout = nn.Dropout(p=dropout)
+        self.scale = scale * (alpha / r)
+        scale = 1 / math.sqrt(input_dims)
+        self.lora_a = mx.random.uniform(low=-scale, high=scale, shape=(input_dims, r))
+        self.lora_b = mx.zeros(shape=(r, output_dims))
+        self.m = mx.linalg.norm(self._dequantized_weight(), axis=1).astype(mx.float32)
+
+    def _dequantized_weight(self):
+        weight = self.linear.weight
+        if isinstance(self.linear, nn.QuantizedLinear):
+            weight = mx.dequantize(weight, self.linear.scales, self.linear.biases, self.linear.group_size, self.linear.bits)
+        return weight
+
+    def __call__(self, x):
+        y = self.linear(x)
+        z = (self.dropout(x) @ self.lora_a) @ self.lora_b
+        z = y + (self.scale * z)
+        adapted = self._dequantized_weight() + (self.scale * self.lora_b.T) @ self.lora_a.T
+        denom = mx.stop_gradient(mx.linalg.norm(adapted, axis=1))
+        z = (self.m / denom) * z
+        return z.astype(x.dtype)
+
+def linear_to_lora_layers(model, lora_layers, lora_targets, lora_rank, lora_scale, lora_dropout):
+    if lora_layers == 'all':
+        lora_layers = model.layers
+    elif isinstance(lora_layers, int):
+        lora_layers = model.layers[-lora_layers:]
+    elif isinstance(lora_layers, list):
+        lora_layers = [model.layers[i] for i in lora_layers]
+    else:
+        raise ValueError("Invalid type for lora_layers. Expected int (number of layers) or list (layer indices or names).")
+    def to_lora(layer):
+        return DoRALinear.from_linear(layer, r=lora_rank, alpha=lora_rank, scale=lora_scale, dropout=lora_dropout)
+    for l in lora_layers:
+        lora_layers = [(k, to_lora(m)) for k, m in l.named_modules() if k in lora_targets]
+        l.update_modules(tree_unflatten(lora_layers))
+
+def load_base_model(model_cfg, init=False):
+    model_id='microsoft/Phi-3.5-mini-instruct'
+    model_path = snapshot_download(model_id, allow_patterns=["*.safetensors", "config.json"])
+    with open(f"{model_path}/config.json", "r") as f:
+        config = json.load(f)
+    config = config|model_cfg
+    model_config = SimpleNamespace(**config)
+    model = Phi3ForCausalLM(model_config)
+    model_weight = [(k, v) for wf in glob.glob(f"{model_path}/*.safetensors") for k, v in mx.load(wf).items()]
+    model.load_weights(model_weight, strict=False)
+    model.set_dtype(mx.float32)
+    if init:
+        init_fn_embed = nn.init.normal(mean=-0.000453949, std=0.0344238)
+        model.apply_to_modules(lambda k, v: v.apply(init_fn_embed) if k.endswith('embed_new') else None)
+        if model_config.untie_embedding:
+            init_fn_lm = nn.init.normal(mean=-0.000231743, std=0.043457)
+            model.apply_to_modules(lambda k, v: v.apply(init_fn_lm) if k.endswith('lm_new') else None)
+    class_predicate = lambda k, m: hasattr(m, "to_quantized") and not k.endswith('new')
+    nn.quantize(model, 64, 4, class_predicate)
+    mx.eval(model.parameters())
+    return model
+
+def load_model_for_training(lora_cfg, model_cfg, thaws, from_path=None):
+    model = load_base_model(model_cfg, init=False)
+    if from_path:
+        model.load_weights(from_path, strict=False)
+    model.freeze()
+    if len(lora_cfg['targets']) > 1:
+        linear_to_lora_layers(model, lora_layers=lora_cfg['layers'], lora_targets=lora_cfg['targets'], lora_rank=lora_cfg['rank'], lora_scale=lora_cfg['scale'], lora_dropout=lora_cfg['dropout'])
+    model.apply_to_modules(lambda k, v: v.unfreeze() if any(k.endswith(t) for t in thaws) else None)
+    mx.eval(model.parameters())
+    # print("Trainable parameters:", [i[0] for i in tree_flatten(model.trainable_parameters())])
+    model.train()
+    return model
+
+def load_model_for_inference(lora_cfg, model_cfg):
+    model = load_base_model(model_cfg, init=False)
+    if len(lora_cfg['targets']) > 1:
+        linear_to_lora_layers(model, lora_layers=lora_cfg['layers'], lora_targets=lora_cfg['targets'], lora_rank=lora_cfg['rank'], lora_scale=lora_cfg['scale'], lora_dropout=lora_cfg['dropout'])
+    _path = 'trained_retnphi.safetensors' if model_cfg['use_retention'] else 'trained_orgnphi.safetensors'
+    model.load_weights(_path, strict=False)
+    mx.eval(model.parameters())
+    model.eval()
+    return model
+
+def generate(prompt, lora_cfg, model_cfg, max_tokens=50, verbose = True):
+    model = load_model_for_inference(lora_cfg=lora_cfg, model_cfg=model_cfg)
+    input_ids = mx.array(tokenizer.encode(prompt))
+    if model_cfg['use_retention']:
+        cache = None
+        for i in input_ids:
+            logits, cache = model(i[None, None], cache=cache, use_recurrent_mode=True)
+    else:
+        logits, cache = model(input_ids[None])
+    token = mx.argmax(logits[:,-1,:], axis=-1)
+    mx.eval(token, cache)
+    list_tokens = token.tolist()
+    for i in range(max_tokens):
+        logits, cache = model(token[None], cache=cache, use_recurrent_mode=True)
+        token = mx.argmax(logits[:,-1,:], axis=-1)
+        mx.eval(token, cache)
+        list_tokens += token.tolist()
+        if tokenizer.decode(list_tokens[-2:]) == '\n\n':
+            break
+    output = tokenizer.decode(list_tokens)
+    if verbose:
+        print(f'{prompt=} + {output=}\n-> {prompt+output}')
+    del model
+    return output
+
+def train_gsm(learning_rates, num_epochs, batch_size, seq_length, lora_cfg, model_cfg, thaws, take, from_path=None):
+    def load_gsm_data(tokenizer, is_tiny=True):
+        if is_tiny:
+            data = load_dataset("TinyGSM/TinyGSM")["train"]
+            if take:
+                data = data.take(take)
+            data = data.filter(lambda x: len(x['question']) < 100 and ':' not in x['question'] and '-' not in x['question'] and "'" not in x['code'] and '\n    result =' in x['code'])
+            split_point = int(len(data) * 0.8)
+            train_data = data.select(range(split_point))
+            eval_data = data.select(range(split_point, len(data)))
+            def format_example(example):
+                code_raw = example['code']
+                start = code_raw.rfind('\n    """')
+                if start == -1:
+                    print('Wrong format to start')
+                    return code_raw.strip()
+                start = start + 8
+                end = code_raw.rfind('\n    result =')
+                if end == -1:
+                    print('Wrong format to end')
+                    end = len(code_raw)
+                code_block = code_raw[start:end]
+                code_lines = code_block.split('\n    ')
+                formatted_code = '\n'.join(line.rstrip() for line in code_lines if line.strip())
+                formatted_code = '\n' + formatted_code.strip() + '\n\n'
+                result = (example['question'].strip(), formatted_code)
+                return result
+        else:
+            dataset = load_dataset("openai/gsm8k", "main")
+            train_data = dataset["train"]
+            eval_data = dataset["test"]
+            def format_example(example):
+                return (example['question'].strip(), '\n'+example['answer'].strip()+'\n\n')
+        train_formatted = [format_example(ex) for ex in train_data]
+        eval_formatted = [format_example(ex) for ex in eval_data]
+        return train_formatted, eval_formatted
+
+    def create_batches(data, tokenizer, batch_size, seq_length):
+        def _encode(x):
+            return [tokenizer.encode(i) for i in x]
+        encoded_data = [_encode(x) for x in data]
+        encoded_data = [x for x in encoded_data if len(x[0]+x[1]) <= seq_length+1]
+        if batch_size is None:
+            batch_size = min(len(encoded_data), 64)
+        else:
+            encoded_data = encoded_data[:(len(encoded_data) // batch_size) * batch_size]
+            np.random.shuffle(encoded_data)
+        for i in range(0, len(encoded_data), batch_size):
+            batch = encoded_data[i:i+batch_size]
+            max_len = min(max(len(q+a)-1 for q, a in batch), seq_length)
+            x_batch = []
+            y_batch = []
+            mask_batch = []
+            for q, a in batch:
+                combined = (q+a)[:max_len+1]
+                x = combined[:-1]
+                y = combined[1:]
+                pad_length = max_len - len(x)
+                x = x + [0] * pad_length
+                y = y + [0] * pad_length
+                mask = [False] * (len(q)-1) + [True] * (len(a)) + [False] * (pad_length)
+                x_batch.append(x)
+                y_batch.append(y)
+                mask_batch.append(mask)
+            yield mx.array(x_batch), mx.array(y_batch), mx.array(mask_batch)
+
+    def loss_fn(model, X, y, mask):
+        logits, _ = model(X)
+        logits = logits.astype(mx.float32)
+        ce = nn.losses.cross_entropy(logits, y, reduction='none')
+        masked_loss = ce * mask
+        return masked_loss.sum(), mask.sum()
+
+    def evaluate(model, data, tokenizer, seq_length):
+        model.eval()
+        total_loss = 0
+        total_samples = 0
+        for X, y, mask in create_batches(data, tokenizer, None, seq_length):
+            loss, ntoks = loss_fn(model, X, y, mask)
+            total_loss += loss.item()
+            total_samples += ntoks.item()
+        return total_loss / total_samples if total_samples > 0 else -1
+
+    def get_optimizer(train_data):
+        num_batches_per_epoch = len(list(create_batches(train_data, tokenizer, batch_size, seq_length)))
+        print(f'{num_batches_per_epoch=}')
+        num_steps = num_epochs * num_batches_per_epoch
+        num_warmup = num_steps // 10
+        max_lr, min_lr = learning_rates
+        if num_warmup > 2:
+            warmup = optim.linear_schedule(min_lr*0.1, max_lr, steps=num_warmup)
+            cosine = optim.cosine_decay(max_lr, num_steps - num_warmup, min_lr)
+            lr_schedule = optim.join_schedules([warmup, cosine], [num_warmup])
+        else:
+            lr_schedule = optim.cosine_decay(max_lr, num_steps, min_lr)
+        return optim.Lion(learning_rate=lr_schedule), num_steps
+
+    for arg_name in sorted(locals()):
+        if arg_name != 'self':
+            arg_value = locals()[arg_name]
+            if not callable(arg_value):
+                print(f"{arg_name}: {arg_value}")
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    print(f'--- {timestamp} ---')
+    train_data, eval_data = load_gsm_data(tokenizer=tokenizer)
+    model = load_model_for_training(lora_cfg=lora_cfg, model_cfg=model_cfg, thaws=thaws)
+    optimizer, num_steps = get_optimizer(train_data)
+    loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
+    mx.eval(model, optimizer)
+    metrics = {
+        'steps': [],
+        'learning_rates': [],
+        'all_train_losses': [],
+        'avg_train_losses': [],
+        'val_losses': [],
+        'trained_toks': [],
+    }
+    step = 0
+    trained_toks = 0
+    losses = []
+    tic = time.perf_counter()
+    for epoch in range(num_epochs):
+        for X, y, loss_mask in create_batches(data=train_data, tokenizer=tokenizer, batch_size=batch_size, seq_length=seq_length):
+            model.train()
+            (loss, ntoks), grads = loss_and_grad_fn(model, X, y, loss_mask)
+            optimizer.update(model, grads)
+            mx.eval(loss, ntoks, model, optimizer)
+            losses.append(loss.item())
+            trained_toks += ntoks.item()
+            step += 1
+            if (step % (num_steps // 30) == 0):
+                avg_train_loss = np.mean(losses)
+                lr = optimizer.learning_rate.item()
+                val_loss = evaluate(model=model, data=eval_data, tokenizer=tokenizer, seq_length=seq_length)
+                print(f"{avg_train_loss:8.4f} ({val_loss:6.4f}) @ {step//(num_steps//30):2}/30 w/ {lr:.2e} ({time.perf_counter() - tic:.2f} sec)")
+                metrics['val_losses'].append(val_loss)
+                # print(f"{avg_train_loss:8.4f} @ {step//(num_steps//30):2}/30 w/ {lr:.2e} ({time.perf_counter() - tic:.2f} sec)")
+                tic = time.perf_counter()
+                metrics['steps'].append(step)
+                metrics['learning_rates'].append(lr)
+                metrics['all_train_losses'].extend(losses)
+                metrics['avg_train_losses'].append(avg_train_loss)
+                metrics['trained_toks'].append(trained_toks)
+                losses = []
+                trained_toks = 0
+    _path = f'trained_retnphi.safetensors' if model_cfg['use_retention'] else f'trained_orgnphi.safetensors'
+    mx.save_safetensors(_path, dict(tree_flatten(model.trainable_parameters())))
+    log = {
+        'args': {
+            'learning_rates': learning_rates,
+            'num_epochs': num_epochs,
+            'batch_size': batch_size,
+            'seq_length': seq_length,
+            'lora_cfg': lora_cfg,
+            'model_cfg': model_cfg,
+            'thaws': thaws,
+            'from_path': from_path
+        },
+        'metrics': metrics
+    }
+    with open(f'train_log_{timestamp}.json', 'w') as f:
+        json.dump(log, f, indent=2)
+    del model
+
+tokenizer = Tokenizer()
+
+def main(take=1000, layers='all', targets=["self_attn.o_proj"], thaws=['new', 'post_attention_layernorm'], rank=32, scale=0.1, dropout=0.0, lr_max=1e-4, lr_min=1e-5, num_epochs=90, batch_size=1, seq_length=256, vocab_size=256, use_retention=True, untie_embedding=True, prompt='There are 8 candies in a carton. How many candies will be in 5 cartons?'):
+    lora_cfg = dict(layers=layers, targets=targets, rank=rank, scale=scale, dropout=dropout)
+    model_cfg = dict(vocab_size=vocab_size, use_retention=use_retention, untie_embedding=untie_embedding)
+    train_gsm(learning_rates=(lr_max, lr_min), num_epochs=num_epochs, batch_size=batch_size, seq_length=seq_length, lora_cfg=lora_cfg, model_cfg=model_cfg, thaws=thaws, take=take)
+    generate(prompt=prompt, lora_cfg=lora_cfg, model_cfg=model_cfg, max_tokens=(seq_length-len(prompt)))
+
+if __name__ == "__main__":
+    main(take=None, num_epochs=3) # -> 240916
+    main(take=None, num_epochs=3, untie_embedding=False)
+
+    main(take=None, num_epochs=3, use_retention=False)
+    main(take=None, num_epochs=3, untie_embedding=False, use_retention=False)
+    # fire.Fire(main)
+
+# Output:
+# 388.7268 @  1/30 w/ 3.36e-05 (64.73 sec)
+# ...
+#   4.3768 @ 30/30 w/ 1.00e-05 (64.36 sec)
+# prompt='There are 8 candies in a carton. How many candies will be in 5 cartons?' + output='\ncandies_in_carton = 8 \nnumber_of_cartons = 5\ntotal_no_of_candies = candies_in_carton * number_of_cartons\n\n'
+# -> There are 8 candies in a carton. How many candies will be in 5 cartons?
+# candies_in_carton = 8
+# number_of_cartons = 5
+# total_no_of_candies = candies_in_carton * number_of_cartons