You can deploy this model and start earning money today!

JosefAlbers/
retnphi

Name: retnphi
Availability: InStock

Views

No views yet

"""

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:

Byte-Level Processing: Operates directly on raw byte sequences, enabling universal application to any file type.
RetNet Integration: Incorporates RetNet's multi-scale exponential decay and group normalization for efficient long-range dependency modeling.
Dual-mode Processing: Supports parallel mode for efficient training and recurrent mode for inference.
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.
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

JosefAlbers/retnphi