# Copyright © 2024 Apple Inc. import time from dataclasses import dataclass, field from functools import partial from pathlib import Path from typing import Optional import mlx.core as mx import mlx.nn as nn import numpy as np from mlx.nn.utils import average_gradients from mlx.utils import tree_flatten, tree_map from ..cli_ui import TrainUI, rprint from ..utils import maybe_set_recommended_wired_limit from .callbacks import TrainingCallback from .datasets import CacheDataset def _clear_cache(threshold: int): if mx.get_cache_memory() > threshold: mx.clear_cache() def grad_checkpoint(layer): """ Update all instances of type(layer) to use gradient checkpointing. """ fn = type(layer).__call__ def checkpointed_fn(model, *args, **kwargs): def inner_fn(params, *args, **kwargs): model.update(params) return fn(model, *args, **kwargs) return mx.checkpoint(inner_fn)(model.trainable_parameters(), *args, **kwargs) type(layer).__call__ = checkpointed_fn @dataclass class TrainingArgs: batch_size: int = field(default=4, metadata={"help": "Minibatch size."}) iters: int = field(default=100, metadata={"help": "Iterations to train for."}) val_batches: int = field( default=25, metadata={ "help": "Number of validation batches, -1 uses the entire validation set." }, ) steps_per_report: int = field( default=10, metadata={"help": "Number of training steps between loss reporting."}, ) steps_per_eval: int = field( default=200, metadata={"help": "Number of training steps between validations."} ) steps_per_save: int = field( default=100, metadata={"help": "Save the model every number steps"} ) max_seq_length: int = field( default=2048, metadata={"help": "Maximum sequence length."} ) adapter_file: str = field( default="adapters.safetensors", metadata={"help": "Save/load path for the trained adapter weights."}, ) grad_checkpoint: bool = field( default=False, metadata={"help": "Use gradient checkpointing to reduce memory use."}, ) grad_accumulation_steps: int = field( default=1, metadata={ "help": "Number of steps to accumulate gradients before applying an optimizer update." }, ) clear_cache_threshold: int = field( default=0, metadata={ "help": "Clear the allocator cache between steps if it grows too large." }, ) def default_loss(model, batch, lengths): inputs = batch[:, :-1] targets = batch[:, 1:] logits = model(inputs) steps = mx.arange(1, targets.shape[1] + 1) mask = mx.logical_and(steps >= lengths[:, 0:1], steps <= lengths[:, 1:]) ce = nn.losses.cross_entropy(logits, targets) * mask ntoks = mask.sum() ce = ce.astype(mx.float32).sum() / ntoks return ce, ntoks def iterate_batches( dataset, batch_size, max_seq_length, loop=False, seed=None, comm_group=None, ): # Sort by length: if isinstance(dataset, CacheDataset): len_fn = lambda idx: dataset.itemlen(idx) else: len_fn = lambda idx: len(dataset[idx][0]) idx = sorted(range(len(dataset)), key=len_fn) if len(dataset) < batch_size: raise ValueError( f"Dataset must have at least batch_size={batch_size}" f" examples but only has {len(dataset)}." ) # If running in distributed mode (N machines) then each one should skip N-1 # samples if comm_group is not None: offset = comm_group.rank() step = comm_group.size() else: offset = 0 step = 1 if batch_size % step != 0: raise ValueError("The batch size must be divisible by the number of workers") # Make the batches: batch_idx = [ idx[i + offset : i + offset + batch_size : step] for i in range(0, len(idx) - batch_size + 1, batch_size) ] if seed is not None: np.random.seed(seed) while True: indices = np.random.permutation(len(batch_idx)) for i in indices: batch = [dataset[j] for j in batch_idx[i]] if len(batch[0]) == 2: batch, offsets = zip(*batch) else: offsets = [0] * len(batch) lengths = [len(x) for x in batch] if max(lengths) > max_seq_length: rprint( f"[WARNING] Some sequences are longer than {max_seq_length} tokens. " f"The longest sentence {max(lengths)} will be truncated to {max_seq_length}. " "Consider pre-splitting your data to save memory." ) # Pad to one plus nearest multiple of pad_to or the maximum length pad_to = 32 max_length_in_batch = 1 + pad_to * ((max(lengths) + pad_to - 1) // pad_to) max_length_in_batch = min(max_length_in_batch, max_seq_length) batch_arr = np.zeros((batch_size // step, max_length_in_batch), np.int32) for j in range(batch_size // step): truncated_length = min(lengths[j], max_seq_length) batch_arr[j, :truncated_length] = batch[j][:truncated_length] lengths[j] = ( truncated_length # Update lengths to match truncated lengths ) batch = mx.array(batch_arr) yield batch, mx.array(list(zip(offsets, lengths))) if not loop: break def evaluate( model, dataset, batch_size, num_batches, max_seq_length=2048, loss: callable = default_loss, iterate_batches: callable = iterate_batches, clear_cache_threshold: int = 0, progress_callback: Optional[callable] = None, ): model.eval() all_losses = mx.array(0.0) ntokens = mx.array(0) index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1) batch_iter = zip( index_iterator, iterate_batches( dataset=dataset, batch_size=batch_size, max_seq_length=max_seq_length, comm_group=mx.distributed.init(), ), ) for _, batch in batch_iter: losses, toks = loss(model, *batch) all_losses += losses * toks ntokens += toks mx.eval(all_losses, ntokens) _clear_cache(clear_cache_threshold) if progress_callback is not None: progress_callback() all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu) ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu) avg_loss = (all_losses / ntokens).item() return avg_loss def train( model, optimizer, train_dataset, val_dataset=None, args: TrainingArgs = TrainingArgs(), loss: callable = default_loss, iterate_batches: callable = iterate_batches, training_callback: TrainingCallback = None, ): _ = maybe_set_recommended_wired_limit() world = mx.distributed.init() world_size = world.size() rank = world.rank() if args.grad_checkpoint: grad_checkpoint(model.layers[0]) loss_value_and_grad = nn.value_and_grad(model, loss) grad_accum_steps = args.grad_accumulation_steps if grad_accum_steps < 1: raise ValueError("grad_accumulation_steps must be at least 1") state = [model.state, optimizer.state, mx.random.state] @partial(mx.compile, inputs=state, outputs=state) def step(batch, prev_grad, do_update): (lvalue, toks), grad = loss_value_and_grad(model, *batch) if prev_grad is not None: grad = tree_map(lambda x, y: x + y, grad, prev_grad) if do_update: grad = average_gradients(grad) if grad_accum_steps > 1: grad = tree_map(lambda x: x / grad_accum_steps, grad) optimizer.update(model, grad) grad = None return lvalue, toks, grad model.train() losses = 0 n_tokens = 0 steps = 0 trained_tokens = 0 train_time = 0 grad_accum = None ui = TrainUI(args.iters, rank=rank) with ui: # Main training loop for it, batch in zip( range(1, args.iters + 1), iterate_batches( dataset=train_dataset, batch_size=args.batch_size, max_seq_length=args.max_seq_length, loop=True, comm_group=world, ), ): tic = time.perf_counter() if val_dataset and ( it == 1 or it % args.steps_per_eval == 0 or it == args.iters ): if args.val_batches == -1: val_total = len(val_dataset) // args.batch_size else: val_total = min( len(val_dataset) // args.batch_size, args.val_batches ) tic = time.perf_counter() with ui.val_task(val_total) as advance_val: val_loss = evaluate( model=model, dataset=val_dataset, loss=loss, batch_size=args.batch_size, num_batches=args.val_batches, max_seq_length=args.max_seq_length, iterate_batches=iterate_batches, progress_callback=advance_val, ) model.train() val_time = time.perf_counter() - tic ui.report_val(it, val_loss, val_time) if training_callback is not None: training_callback.on_val_loss_report( { "iteration": it - 1, "val_loss": val_loss, "val_time": val_time, } ) tic = time.perf_counter() lvalue, toks, grad_accum = step( batch, grad_accum, it % grad_accum_steps == 0, ) losses += lvalue n_tokens += toks steps += 1 mx.eval(state, losses, n_tokens, grad_accum) _clear_cache(args.clear_cache_threshold) train_time += time.perf_counter() - tic ui.advance() # Report training loss if needed if it % args.steps_per_report == 0 or it == args.iters: train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item() train_loss /= steps * world_size n_tokens = mx.distributed.all_sum(n_tokens, stream=mx.cpu).item() learning_rate = optimizer.learning_rate.item() it_sec = args.steps_per_report / train_time tokens_sec = float(n_tokens) / train_time trained_tokens += n_tokens peak_mem = mx.get_peak_memory() / 1e9 ui.report_train(it, train_loss, tokens_sec, trained_tokens) if training_callback is not None: training_callback.on_train_loss_report( { "iteration": it, "train_loss": train_loss, "learning_rate": learning_rate, "iterations_per_second": it_sec, "tokens_per_second": tokens_sec, "trained_tokens": trained_tokens, "peak_memory": peak_mem, } ) losses = 0 n_tokens = 0 steps = 0 train_time = 0 # Save adapter weights if it % args.steps_per_save == 0 and rank == 0: adapter_weights = dict(tree_flatten(model.trainable_parameters())) mx.save_safetensors(str(args.adapter_file), adapter_weights) checkpoint = ( Path(args.adapter_file).parent / f"{it:07d}_adapters.safetensors" ) mx.save_safetensors(str(checkpoint), adapter_weights) ui.report_save(checkpoint) # Save final weights if rank == 0: adapter_weights = dict(tree_flatten(model.trainable_parameters())) mx.save_safetensors(str(args.adapter_file), adapter_weights)