546 lines
15 KiB
Python
546 lines
15 KiB
Python
import argparse
|
|
import os
|
|
from functools import partial
|
|
from typing import Any, Dict
|
|
|
|
import torch
|
|
import torch.optim as optim
|
|
from torch import Tensor, nn
|
|
|
|
from astrai.config import AutoRegressiveLMConfig, TrainConfig
|
|
from astrai.dataset import DatasetFactory
|
|
from astrai.model import AutoRegressiveLM
|
|
from astrai.model.components.decoder_block import DecoderBlock
|
|
from astrai.trainer import SchedulerFactory, Trainer
|
|
|
|
|
|
class MuonMix(optim.Optimizer):
|
|
"""Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: nn.Module,
|
|
lr: float = 3e-4,
|
|
weight_decay: float = 0.1,
|
|
momentum: float = 0.95,
|
|
nesterov: bool = True,
|
|
ns_steps: int = 5,
|
|
adjust_lr_fn: str = "match_rms_adamw",
|
|
):
|
|
defaults = dict(
|
|
lr=lr,
|
|
weight_decay=weight_decay,
|
|
momentum=momentum,
|
|
nesterov=nesterov,
|
|
ns_steps=ns_steps,
|
|
adjust_lr_fn=adjust_lr_fn,
|
|
)
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
super().__init__(params, defaults)
|
|
|
|
matrix_params: list[Tensor] = []
|
|
other_params: list[Tensor] = []
|
|
for name, param in model.named_parameters():
|
|
if not param.requires_grad:
|
|
continue
|
|
if (
|
|
param.dim() >= 2
|
|
and "norm" not in name
|
|
and "bias" not in name
|
|
and "embed" not in name
|
|
and "lm_head" not in name
|
|
):
|
|
matrix_params.append(param)
|
|
else:
|
|
other_params.append(param)
|
|
|
|
self.muon = optim.Muon(
|
|
matrix_params,
|
|
lr=lr,
|
|
weight_decay=weight_decay,
|
|
momentum=momentum,
|
|
nesterov=nesterov,
|
|
ns_steps=ns_steps,
|
|
adjust_lr_fn=adjust_lr_fn,
|
|
)
|
|
self.adamw = optim.AdamW(
|
|
[{"params": other_params, "weight_decay": 0.0}],
|
|
lr=lr,
|
|
betas=(0.9, 0.95),
|
|
fused=True,
|
|
)
|
|
|
|
self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
|
|
|
|
@torch.no_grad()
|
|
def step(self, closure=None):
|
|
self.muon.step(closure)
|
|
self.adamw.step(closure)
|
|
|
|
def zero_grad(self, set_to_none: bool = True):
|
|
self.muon.zero_grad(set_to_none)
|
|
self.adamw.zero_grad(set_to_none)
|
|
|
|
def state_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"muon": self.muon.state_dict(),
|
|
"adamw": self.adamw.state_dict(),
|
|
}
|
|
|
|
def load_state_dict(self, state_dict: Dict[str, Any]):
|
|
self.muon.load_state_dict(state_dict["muon"])
|
|
self.adamw.load_state_dict(state_dict["adamw"])
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
|
|
parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.")
|
|
|
|
parser.add_argument(
|
|
"--train_type",
|
|
type=str,
|
|
required=True,
|
|
choices=["seq", "sft", "dpo", "grpo"],
|
|
help="Train type.",
|
|
)
|
|
parser.add_argument(
|
|
"--data_root_path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to the root directory of the dataset.",
|
|
)
|
|
parser.add_argument(
|
|
"--param_path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to the model parameters or resume checkpoint.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--n_epoch", type=int, default=1, help="Number of epochs to train."
|
|
)
|
|
parser.add_argument(
|
|
"--batch_per_device", type=int, default=1, help="Batch size per GPU."
|
|
)
|
|
parser.add_argument(
|
|
"--grad_accum_steps",
|
|
type=int,
|
|
default=1,
|
|
help="Number of iterations between each optimizer step.",
|
|
)
|
|
parser.add_argument(
|
|
"--warmup_ratio",
|
|
type=float,
|
|
default=0.05,
|
|
help="Fraction of total steps used for LR warmup.",
|
|
)
|
|
parser.add_argument(
|
|
"--max_lr", type=float, default=3e-4, help="Max learning rate for training."
|
|
)
|
|
parser.add_argument(
|
|
"--max_grad_norm",
|
|
type=float,
|
|
default=1.0,
|
|
help="Max gradient norm for clipping.",
|
|
)
|
|
parser.add_argument(
|
|
"--weight_decay",
|
|
type=float,
|
|
default=0.1,
|
|
help="Weight decay (applied to Muon matrix params; non-matrix use 0).",
|
|
)
|
|
parser.add_argument(
|
|
"--muon_momentum",
|
|
type=float,
|
|
default=0.95,
|
|
help="Momentum factor for Muon optimizer.",
|
|
)
|
|
parser.add_argument(
|
|
"--muon_nesterov",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Enable Nesterov momentum for Muon.",
|
|
)
|
|
parser.add_argument(
|
|
"--muon_ns_steps",
|
|
type=int,
|
|
default=5,
|
|
help="Newton-Schulz iteration steps for Muon.",
|
|
)
|
|
parser.add_argument(
|
|
"--muon_adjust_lr",
|
|
type=str,
|
|
default="match_rms_adamw",
|
|
choices=["original", "match_rms_adamw"],
|
|
help="Muon learning rate adjustment strategy.",
|
|
)
|
|
parser.add_argument(
|
|
"--random_seed", type=int, default=3407, help="Random seed for reproducibility."
|
|
)
|
|
parser.add_argument(
|
|
"--num_workers", type=int, default=4, help="Number of workers for data loading."
|
|
)
|
|
parser.add_argument(
|
|
"--no_pin_memory",
|
|
action="store_false",
|
|
dest="pin_memory",
|
|
help="Disable pin memory",
|
|
)
|
|
parser.add_argument(
|
|
"--window_size",
|
|
type=int,
|
|
default=None,
|
|
help="Max length of the input sequence.",
|
|
)
|
|
parser.add_argument(
|
|
"--stride", type=int, default=None, help="Step size of the input sequence."
|
|
)
|
|
parser.add_argument("--dpo_beta", type=float, default=0.1, help="DPO beta value.")
|
|
parser.add_argument("--group_size", type=int, default=4, help="GRPO group size.")
|
|
parser.add_argument(
|
|
"--grpo_clip_eps", type=float, default=0.2, help="GRPO clipping epsilon."
|
|
)
|
|
parser.add_argument(
|
|
"--grpo_kl_coef", type=float, default=0.01, help="GRPO KL penalty coefficient."
|
|
)
|
|
parser.add_argument(
|
|
"--label_smoothing",
|
|
type=float,
|
|
default=0.0,
|
|
help="cross_entropy function label smoothing parameter",
|
|
)
|
|
parser.add_argument(
|
|
"--gradient_checkpointing",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="Enable activation checkpointing for DecoderBlock modules.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--ckpt_interval",
|
|
type=int,
|
|
default=5000,
|
|
help="Number of iters between checkpoints.",
|
|
)
|
|
parser.add_argument(
|
|
"--ckpt_dir",
|
|
type=str,
|
|
default="checkpoint",
|
|
help="Directory to save checkpoints.",
|
|
)
|
|
parser.add_argument(
|
|
"--val_split",
|
|
type=float,
|
|
default=None,
|
|
help="Ratio to split from training dataset for validation (e.g. 0.05).",
|
|
)
|
|
parser.add_argument(
|
|
"--val_step",
|
|
type=int,
|
|
default=1000,
|
|
help="Number of optimizer steps between validation runs.",
|
|
)
|
|
parser.add_argument(
|
|
"--metrics",
|
|
nargs="*",
|
|
default=["loss", "lr", "grad_norm"],
|
|
help="Metrics to log (e.g. --metrics loss lr val_loss). Default: loss lr grad_norm.",
|
|
)
|
|
parser.add_argument(
|
|
"--log_dir",
|
|
type=str,
|
|
default="checkpoint/logs",
|
|
help="Directory for metric logs.",
|
|
)
|
|
parser.add_argument(
|
|
"--grpo_sync_interval",
|
|
type=int,
|
|
default=200,
|
|
help="GRPO ref model sync interval (steps).",
|
|
)
|
|
parser.add_argument(
|
|
"--start_epoch", type=int, default=0, help="Start epoch for training."
|
|
)
|
|
parser.add_argument(
|
|
"--start_samples",
|
|
type=int,
|
|
default=0,
|
|
help="Start samples (per rank) for training.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--master_addr",
|
|
type=str,
|
|
default="localhost",
|
|
help="Master node address for distributed training.",
|
|
)
|
|
parser.add_argument(
|
|
"--master_port",
|
|
type=str,
|
|
default="29500",
|
|
help="Master node port for distributed training.",
|
|
)
|
|
parser.add_argument(
|
|
"--backend",
|
|
type=str,
|
|
default="nccl",
|
|
help="Distributed training backend.",
|
|
)
|
|
parser.add_argument("--nprocs", type=int, default=1, help="Number of GPUs to use.")
|
|
parser.add_argument(
|
|
"--parallel_mode",
|
|
type=str,
|
|
default="none",
|
|
choices=["none", "ddp", "fsdp"],
|
|
help="Parallel training strategy (none, ddp, fsdp).",
|
|
)
|
|
parser.add_argument(
|
|
"--device_type", type=str, default="cuda", help="Device type to use."
|
|
)
|
|
parser.add_argument(
|
|
"--start_method",
|
|
type=str,
|
|
default="spawn",
|
|
choices=["spawn", "fork", "forkserver"],
|
|
help="Multiprocessing start method.",
|
|
)
|
|
parser.add_argument(
|
|
"--neftune_alpha",
|
|
type=float,
|
|
default=0.0,
|
|
help="NEFTune noise alpha (0=disabled, typical: 5.0).",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--schedule_type",
|
|
type=str,
|
|
default="cosine",
|
|
choices=["cosine", "sgdr", "wsd"],
|
|
help="Learning rate scheduler type.",
|
|
)
|
|
parser.add_argument(
|
|
"--min_rate",
|
|
type=float,
|
|
default=None,
|
|
help="Minimum LR as fraction of base LR. Uses scheduler default if not set (cosine/sgdr: 0.05, wsd: 0.0).",
|
|
)
|
|
parser.add_argument(
|
|
"--cycle_length",
|
|
type=int,
|
|
default=None,
|
|
help="SGDR first cycle length in steps. Defaults to total_steps - warmup_steps.",
|
|
)
|
|
parser.add_argument(
|
|
"--t_mult",
|
|
type=int,
|
|
default=2,
|
|
help="SGDR cycle length multiplier per restart.",
|
|
)
|
|
parser.add_argument(
|
|
"--stable_steps",
|
|
type=int,
|
|
default=None,
|
|
help="WSD stable plateau steps. Required when --schedule_type wsd.",
|
|
)
|
|
parser.add_argument(
|
|
"--decay_steps",
|
|
type=int,
|
|
default=None,
|
|
help="WSD decay steps. Defaults to total_steps - warmup_steps - stable_steps.",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
return args
|
|
|
|
|
|
def create_model(config):
|
|
return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
|
|
|
|
|
|
def create_optimizer(model, **kwargs) -> MuonMix:
|
|
return MuonMix(model, **kwargs)
|
|
|
|
|
|
def create_scheduler(
|
|
optimizer: optim.Optimizer, **kwargs
|
|
) -> optim.lr_scheduler.LRScheduler:
|
|
schedule_type = kwargs.pop("schedule_type")
|
|
return SchedulerFactory.create(schedule_type, optimizer, **kwargs)
|
|
|
|
|
|
def compute_total_steps(
|
|
dataset_len: int,
|
|
n_epoch: int,
|
|
batch_per_device: int,
|
|
nprocs: int,
|
|
grad_accum_steps: int,
|
|
) -> int:
|
|
|
|
def ceil_div(a: int, b: int) -> int:
|
|
return (a + b - 1) // b
|
|
|
|
samples_per_replica = ceil_div(dataset_len, nprocs)
|
|
batches_per_replica = ceil_div(samples_per_replica, batch_per_device)
|
|
total_steps = (batches_per_replica // grad_accum_steps) * n_epoch
|
|
return total_steps
|
|
|
|
|
|
def train(
|
|
train_type: str,
|
|
param_path: str,
|
|
data_root_path: str,
|
|
n_epoch: int,
|
|
batch_per_device: int,
|
|
start_epoch: int,
|
|
start_samples: int,
|
|
grad_accum_steps: int,
|
|
warmup_ratio: float,
|
|
ckpt_interval: int,
|
|
ckpt_dir: str,
|
|
val_split: float,
|
|
val_step: int,
|
|
metrics: list[str],
|
|
log_dir: str,
|
|
max_grad_norm: float,
|
|
random_seed: int,
|
|
num_workers: int,
|
|
pin_memory: bool,
|
|
gradient_checkpointing: bool,
|
|
window_size: int,
|
|
stride: int,
|
|
nprocs: int,
|
|
parallel_mode: str,
|
|
device_type: str,
|
|
backend: str,
|
|
master_addr: str,
|
|
master_port: str,
|
|
start_method: str,
|
|
neftune_alpha: float,
|
|
schedule_type: str,
|
|
min_rate: float,
|
|
cycle_length: int,
|
|
t_mult: int,
|
|
stable_steps: int,
|
|
decay_steps: int,
|
|
**kwargs,
|
|
):
|
|
assert train_type in ["seq", "sft", "dpo", "grpo"]
|
|
assert os.path.exists(param_path)
|
|
if nprocs > 1 and parallel_mode == "none":
|
|
raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")
|
|
|
|
# Load config
|
|
config_path = os.path.join(param_path, "config.json")
|
|
config = AutoRegressiveLMConfig.from_file(config_path)
|
|
config.neftune_alpha = neftune_alpha
|
|
|
|
if window_size is None:
|
|
window_size = config.max_len
|
|
|
|
strategy_kwargs = {
|
|
"beta": kwargs.pop("dpo_beta"),
|
|
"label_smoothing": kwargs.pop("label_smoothing"),
|
|
"clip_eps": kwargs.pop("grpo_clip_eps"),
|
|
"kl_coef": kwargs.pop("grpo_kl_coef"),
|
|
"group_size": kwargs.pop("group_size"),
|
|
"sync_interval": kwargs.pop("grpo_sync_interval"),
|
|
}
|
|
|
|
executor_kwargs = {
|
|
"gradient_as_bucket_view": True,
|
|
"broadcast_buffers": False,
|
|
}
|
|
|
|
model_fn = partial(create_model, config)
|
|
dataset = DatasetFactory.load(
|
|
train_type=train_type,
|
|
load_path=data_root_path,
|
|
window_size=window_size,
|
|
stride=stride,
|
|
)
|
|
|
|
optimizer_fn = partial(
|
|
create_optimizer,
|
|
lr=kwargs.pop("max_lr"),
|
|
weight_decay=kwargs.pop("weight_decay"),
|
|
momentum=kwargs.pop("muon_momentum"),
|
|
nesterov=kwargs.pop("muon_nesterov"),
|
|
ns_steps=kwargs.pop("muon_ns_steps"),
|
|
adjust_lr_fn=kwargs.pop("muon_adjust_lr"),
|
|
)
|
|
|
|
total_steps = compute_total_steps(
|
|
len(dataset), n_epoch, batch_per_device, nprocs, grad_accum_steps
|
|
)
|
|
warmup_steps = int(warmup_ratio * total_steps)
|
|
warmup_steps = min(warmup_steps, total_steps)
|
|
|
|
scheduler_kwargs = {"warmup_steps": warmup_steps}
|
|
|
|
if schedule_type == "cosine":
|
|
scheduler_kwargs["lr_decay_steps"] = total_steps - warmup_steps
|
|
elif schedule_type == "sgdr":
|
|
scheduler_kwargs["cycle_length"] = cycle_length or (total_steps - warmup_steps)
|
|
scheduler_kwargs["t_mult"] = t_mult
|
|
elif schedule_type == "wsd":
|
|
remaining = total_steps - warmup_steps
|
|
stable_steps_ = stable_steps or max(1, int(remaining * 0.8))
|
|
scheduler_kwargs["stable_steps"] = stable_steps_
|
|
scheduler_kwargs["decay_steps"] = max(
|
|
1, decay_steps or (remaining - stable_steps_)
|
|
)
|
|
|
|
if min_rate is not None:
|
|
scheduler_kwargs["min_rate"] = min_rate
|
|
|
|
scheduler_fn = partial(
|
|
create_scheduler,
|
|
schedule_type=schedule_type,
|
|
**scheduler_kwargs,
|
|
)
|
|
|
|
grad_ckpt_modules = [DecoderBlock] if gradient_checkpointing else []
|
|
|
|
train_config = TrainConfig(
|
|
model_fn=model_fn,
|
|
strategy=train_type,
|
|
dataset=dataset,
|
|
optimizer_fn=optimizer_fn,
|
|
scheduler_fn=scheduler_fn,
|
|
ckpt_dir=ckpt_dir,
|
|
n_epoch=n_epoch,
|
|
batch_per_device=batch_per_device,
|
|
start_epoch=start_epoch,
|
|
start_samples=start_samples,
|
|
ckpt_interval=ckpt_interval,
|
|
grad_accum_steps=grad_accum_steps,
|
|
max_grad_norm=max_grad_norm,
|
|
random_seed=random_seed,
|
|
num_workers=num_workers,
|
|
pin_memory=pin_memory,
|
|
nprocs=nprocs,
|
|
backend=backend,
|
|
master_addr=master_addr,
|
|
master_port=master_port,
|
|
parallel_mode=parallel_mode,
|
|
device_type=device_type,
|
|
start_method=start_method,
|
|
val_split=val_split,
|
|
val_step=val_step,
|
|
metrics=metrics,
|
|
log_dir=log_dir,
|
|
gradient_checkpointing_modules=grad_ckpt_modules,
|
|
executor_kwargs=executor_kwargs,
|
|
extra_kwargs=strategy_kwargs,
|
|
neftune_alpha=neftune_alpha,
|
|
)
|
|
|
|
trainer = Trainer(train_config)
|
|
trainer.train(resume_dir=param_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
train(**vars(args))
|