mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-08-21 20:46:50 +08:00
399 lines
13 KiB
Python
399 lines
13 KiB
Python
# Copyright © 2024 Apple Inc.
|
|
|
|
import glob
|
|
import shutil
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Union
|
|
|
|
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
|
|
|
|
from trainer import TrainingArgs, TrainingCallback, grad_checkpoint
|
|
|
|
|
|
@dataclass
|
|
class PPOTrainingArgs(TrainingArgs):
|
|
vf_coef: float = field(default=0.5, metadata={"help": "Value function coefficient"})
|
|
cliprange: float = field(default=0.2, metadata={"help": "Policy gradient clipping range"})
|
|
cliprange_value: float = field(default=0.2, metadata={"help": "Value function clipping range"})
|
|
gamma: float = field(default=0.99, metadata={"help": "Discount factor"})
|
|
lambda_: float = field(default=0.95, metadata={"help": "GAE lambda"})
|
|
|
|
|
|
def compute_returns(
|
|
rewards: mx.array,
|
|
gamma: float = 0.99
|
|
) -> mx.array:
|
|
"""Compute returns with Generalized Advantage Estimation"""
|
|
returns = mx.zeros_like(rewards)
|
|
running_return = 0
|
|
|
|
for t in reversed(range(len(rewards))):
|
|
running_return = rewards[t] + gamma * running_return
|
|
returns = returns.at[t].set(running_return)
|
|
|
|
return returns
|
|
|
|
def compute_advantages(
|
|
values: mx.array,
|
|
returns: mx.array,
|
|
rewards: mx.array,
|
|
gamma: float = 0.99,
|
|
lambda_: float = 0.95
|
|
) -> mx.array:
|
|
"""Compute advantages using GAE"""
|
|
advantages = mx.zeros_like(returns)
|
|
running_advantage = 0
|
|
|
|
for t in reversed(range(len(returns))):
|
|
if t < len(returns) - 1:
|
|
delta = rewards[t] + gamma * values[t + 1] - values[t]
|
|
else:
|
|
delta = rewards[t] - values[t]
|
|
|
|
running_advantage = delta + gamma * lambda_ * running_advantage
|
|
advantages = advantages.at[t].set(running_advantage)
|
|
|
|
return (advantages - advantages.mean()) / (advantages.std() + 1e-8)
|
|
|
|
def make_predictions(model, x, mask):
|
|
inputs = x[:, :-1]
|
|
targets = x[:, 1:]
|
|
|
|
logits = model(inputs)
|
|
logits = logits.astype(mx.float32)
|
|
|
|
return -nn.losses.cross_entropy(logits, targets) * mask[:, :-1]
|
|
|
|
def compute_rewards(model, x, mask, reward_scale=1.0):
|
|
"""
|
|
Compute rewards based on model predictions and actual targets.
|
|
Basic implementation using log probabilities as rewards.
|
|
"""
|
|
logits = model(x[:, :-1])
|
|
targets = x[:, 1:]
|
|
|
|
log_probs = -nn.losses.cross_entropy(logits, targets, reduction='none')
|
|
rewards = log_probs * mask[:, :-1] * reward_scale
|
|
|
|
return rewards
|
|
|
|
def ppo_loss(
|
|
model,
|
|
inputs,
|
|
mask,
|
|
old_logprobs,
|
|
values,
|
|
old_values,
|
|
advantages,
|
|
returns,
|
|
vf_coef=0.5,
|
|
cliprange=0.2,
|
|
cliprange_value=0.2
|
|
):
|
|
# Get new log probs
|
|
new_logprobs = make_predictions(model, inputs, mask)
|
|
ntoks = mask[:, :-1].sum()
|
|
new_logprobs = new_logprobs.sum() / ntoks
|
|
|
|
# Value loss with clipping
|
|
vpred_clipped = mx.clip(values, old_values - cliprange_value, old_values + cliprange_value)
|
|
vf_loss = 0.5 * mx.maximum(
|
|
mx.square(values - returns),
|
|
mx.square(vpred_clipped - returns)
|
|
).mean()
|
|
|
|
# Policy loss with clipping
|
|
ratio = mx.exp(new_logprobs - old_logprobs)
|
|
pg_loss = mx.maximum(
|
|
-advantages * ratio,
|
|
-advantages * mx.clip(ratio, 1.0 - cliprange, 1.0 + cliprange)
|
|
).mean()
|
|
|
|
total_loss = pg_loss + vf_coef * vf_loss
|
|
return total_loss, pg_loss, vf_loss, ntoks
|
|
|
|
|
|
def iterate_ppo_batches(dataset, tokenizer, batch_size, max_seq_length, train=False):
|
|
# Sort by length
|
|
idx = sorted(range(len(dataset)), key=lambda idx: len(dataset[idx]))
|
|
if len(dataset) < batch_size:
|
|
raise ValueError(f"Dataset must have at least batch_size={batch_size} examples but only has {len(dataset)}.")
|
|
|
|
# Handle distributed training
|
|
step = mx.distributed.init().size()
|
|
if batch_size % step != 0:
|
|
raise ValueError("The batch size must be divisible by the number of workers")
|
|
|
|
# Make batches
|
|
batch_idx = [idx[i:i+batch_size:step] for i in range(0, len(idx)-batch_size+1, batch_size)]
|
|
|
|
while True:
|
|
indices = np.random.permutation(len(batch_idx))
|
|
for i in indices:
|
|
batch = [dataset[j] for j in batch_idx[i]]
|
|
lengths = [len(x) for x in batch]
|
|
|
|
# Handle sequence length
|
|
if max(lengths) > max_seq_length:
|
|
print(f"[WARNING] Truncating sequences longer than {max_seq_length}")
|
|
|
|
# Pad to multiple of 8
|
|
pad_to = 8
|
|
max_length_in_batch = pad_to * ((max(lengths) + pad_to - 1) // pad_to)
|
|
max_length_in_batch = min(max_length_in_batch, max_seq_length)
|
|
|
|
# Create batch array
|
|
batch_arr = np.zeros((batch_size // step, max_length_in_batch), np.int32)
|
|
mask = 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]
|
|
mask[j, :truncated_length] = 1
|
|
lengths[j] = truncated_length
|
|
|
|
batch = mx.array(batch_arr)
|
|
mask = mx.array(mask)
|
|
|
|
yield batch, mask
|
|
|
|
if not train:
|
|
break
|
|
|
|
|
|
def evaluate(
|
|
model,
|
|
dataset,
|
|
tokenizer,
|
|
batch_size,
|
|
num_batches,
|
|
max_seq_length=2048,
|
|
old_logprobs=None,
|
|
values=None,
|
|
old_values=None,
|
|
advantages=None,
|
|
returns=None,
|
|
vf_coef=0.5,
|
|
cliprange=0.2,
|
|
cliprange_value=0.2,
|
|
loss: callable = ppo_loss,
|
|
iterate_ppo_batches: callable = iterate_ppo_batches,
|
|
):
|
|
total_loss = 0
|
|
total_pg_loss = 0
|
|
total_vf_loss = 0
|
|
ntokens = 0
|
|
|
|
index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)
|
|
|
|
for _, batch in zip(
|
|
index_iterator,
|
|
iterate_ppo_batches(
|
|
dataset=dataset,
|
|
tokenizer=tokenizer,
|
|
batch_size=batch_size,
|
|
max_seq_length=max_seq_length,
|
|
),
|
|
):
|
|
losses, pg_loss, vf_loss, toks = loss(
|
|
model, *batch,
|
|
old_logprobs=old_logprobs,
|
|
values=values,
|
|
old_values=old_values,
|
|
advantages=advantages,
|
|
returns=returns,
|
|
vf_coef=vf_coef,
|
|
cliprange=cliprange,
|
|
cliprange_value=cliprange_value
|
|
)
|
|
|
|
total_loss += losses * toks
|
|
total_pg_loss += pg_loss * toks
|
|
total_vf_loss += vf_loss * toks
|
|
ntokens += toks
|
|
mx.eval(total_loss, total_pg_loss, total_vf_loss, ntokens)
|
|
|
|
total_loss = mx.distributed.all_sum(total_loss, stream=mx.cpu)
|
|
total_pg_loss = mx.distributed.all_sum(total_pg_loss, stream=mx.cpu)
|
|
total_vf_loss = mx.distributed.all_sum(total_vf_loss, stream=mx.cpu)
|
|
ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
|
|
|
|
return (total_loss / ntokens).item(), (total_pg_loss / ntokens).item(), (total_vf_loss / ntokens).item()
|
|
|
|
|
|
def train(
|
|
model,
|
|
tokenizer,
|
|
optimizer,
|
|
train_dataset,
|
|
val_dataset,
|
|
args: PPOTrainingArgs = PPOTrainingArgs(),
|
|
loss: callable = ppo_loss,
|
|
iterate_ppo_batches: callable = iterate_ppo_batches,
|
|
training_callback: TrainingCallback = None,
|
|
):
|
|
print(f"Starting PPO training..., iters: {args.iters}")
|
|
world = mx.distributed.init()
|
|
world_size = world.size()
|
|
rank = world.rank()
|
|
if world_size > 1:
|
|
print(f"Node {rank} of {world_size}")
|
|
|
|
if args.grad_checkpoint:
|
|
grad_checkpoint(model.layers[0])
|
|
|
|
state = [model.state, optimizer.state]
|
|
|
|
def step(batch):
|
|
x, mask = batch
|
|
|
|
# Initial forward pass
|
|
old_logprobs = make_predictions(model, x, mask)
|
|
values = model.value_head(x[:, :-1])
|
|
old_values = values.copy()
|
|
|
|
# Compute rewards (implement reward calculation based on your task)
|
|
rewards = compute_rewards(model, x, mask)
|
|
|
|
# Compute returns and advantages
|
|
returns = compute_returns(rewards, values, gamma=args.gamma)
|
|
advantages = compute_advantages(values, returns, rewards,
|
|
gamma=args.gamma,
|
|
lambda_=args.lambda_)
|
|
|
|
def loss_fn(model, x, mask):
|
|
total_loss, pg_loss, vf_loss, ntoks = ppo_loss(
|
|
model, x, mask,
|
|
old_logprobs, values, old_values,
|
|
advantages, returns,
|
|
vf_coef=args.vf_coef,
|
|
cliprange=args.cliprange,
|
|
cliprange_value=args.cliprange_value
|
|
)
|
|
return total_loss, ntoks, pg_loss, vf_loss
|
|
|
|
(loss_val, toks, pg_loss, vf_loss), grad = nn.value_and_grad(model, loss_fn)(x, mask)
|
|
grad = average_gradients(grad)
|
|
optimizer.update(model, grad)
|
|
|
|
return loss_val, toks, pg_loss, vf_loss
|
|
|
|
losses = 0
|
|
n_tokens = 0
|
|
steps = 0
|
|
trained_tokens = 0
|
|
# Main training loop
|
|
start = time.perf_counter()
|
|
for it, batch in zip(
|
|
range(1, args.iters + 1),
|
|
iterate_ppo_batches(
|
|
dataset=train_dataset,
|
|
tokenizer=tokenizer,
|
|
batch_size=args.batch_size,
|
|
max_seq_length=args.max_seq_length,
|
|
train=True,
|
|
),
|
|
):
|
|
# Report validation loss if needed, the first validation loss
|
|
# is always measured before any training.
|
|
if it == 1 or it % args.steps_per_eval == 0 or it == args.iters:
|
|
stop = time.perf_counter()
|
|
val_loss = evaluate(
|
|
model=model,
|
|
dataset=val_dataset,
|
|
loss=loss,
|
|
tokenizer=tokenizer,
|
|
batch_size=args.batch_size,
|
|
num_batches=args.val_batches,
|
|
max_seq_length=args.max_seq_length,
|
|
iterate_ppo_batches=iterate_ppo_batches,
|
|
)
|
|
val_time = time.perf_counter() - stop
|
|
if rank == 0:
|
|
print(
|
|
f"Iter {it}: "
|
|
f"Val loss {val_loss:.3f}, "
|
|
f"Val took {val_time:.3f}s",
|
|
flush=True,
|
|
)
|
|
|
|
if training_callback is not None:
|
|
val_info = {
|
|
"iteration": it,
|
|
"val_loss": val_loss,
|
|
"val_time": val_time,
|
|
}
|
|
training_callback.on_val_loss_report(val_info)
|
|
|
|
start = time.perf_counter()
|
|
|
|
lvalue, toks = step(batch)
|
|
losses += lvalue
|
|
n_tokens += toks
|
|
steps += 1
|
|
mx.eval(state, losses, n_tokens)
|
|
|
|
# Report training loss if needed
|
|
if it % args.steps_per_report == 0 or it == args.iters:
|
|
stop = time.perf_counter()
|
|
|
|
train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item()
|
|
train_loss /= steps * mx.distributed.init().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 / (stop - start)
|
|
tokens_sec = float(n_tokens) / (stop - start)
|
|
trained_tokens += n_tokens
|
|
peak_mem = mx.metal.get_peak_memory() / 1e9
|
|
if rank == 0:
|
|
print(
|
|
f"Iter {it}: Train loss {train_loss:.3f}, "
|
|
f"Learning Rate {learning_rate:.3e}, "
|
|
f"It/sec {it_sec:.3f}, "
|
|
f"Tokens/sec {tokens_sec:.3f}, "
|
|
f"Trained Tokens {trained_tokens}, "
|
|
f"Peak mem {peak_mem:.3f} GB",
|
|
flush=True,
|
|
)
|
|
|
|
if training_callback is not None:
|
|
train_info = {
|
|
"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,
|
|
}
|
|
training_callback.on_train_loss_report(train_info)
|
|
|
|
losses = 0
|
|
n_tokens = 0
|
|
steps = 0
|
|
start = time.perf_counter()
|
|
|
|
# Save adapter weights
|
|
if it % args.steps_per_save == 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)
|
|
print(
|
|
f"Iter {it}: Saved adapter weights to "
|
|
f"{args.adapter_file} and {checkpoint}."
|
|
)
|
|
|
|
# Save final weights
|
|
adapter_weights = dict(tree_flatten(model.trainable_parameters()))
|
|
mx.save_safetensors(str(args.adapter_file), adapter_weights)
|
|
print(f"Saved final weights to {args.adapter_file}.")
|