272 lines
11 KiB
Python
272 lines
11 KiB
Python
"""
|
||
DeepSpeed ZeRO-2 全参 SFT 入口。
|
||
|
||
推荐在 trainer/ 目录启动：
|
||
deepspeed --num_gpus=7 train_full_sft_deepspeed.py --model_profile SearchLM-300M
|
||
"""
|
||
import os
|
||
import sys
|
||
|
||
__package__ = "trainer"
|
||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||
|
||
import argparse
|
||
import math
|
||
import time
|
||
import warnings
|
||
|
||
import torch
|
||
import torch.distributed as dist
|
||
from torch.utils.data import DataLoader, DistributedSampler
|
||
|
||
from dataset.lm_dataset import SFTDataset
|
||
from trainer.trainer_utils import (
|
||
Logger,
|
||
SkipBatchSampler,
|
||
add_deepspeed_args,
|
||
add_model_profile_args,
|
||
apply_model_profile,
|
||
build_lm_config,
|
||
assert_no_training_output_collision,
|
||
init_model,
|
||
is_main_process,
|
||
get_cuda_peak_memory_gb,
|
||
load_deepspeed_checkpoint,
|
||
load_deepspeed_config,
|
||
log_training_setup,
|
||
reduce_metrics,
|
||
resolve_attention_type,
|
||
save_deepspeed_checkpoint,
|
||
save_model_weights,
|
||
setup_seed,
|
||
)
|
||
|
||
warnings.filterwarnings("ignore")
|
||
|
||
|
||
def init_deepspeed_runtime(args, deepspeed):
|
||
deepspeed.init_distributed(dist_backend="nccl")
|
||
local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank if args.local_rank >= 0 else 0))
|
||
if torch.cuda.is_available():
|
||
torch.cuda.set_device(local_rank)
|
||
args.device = f"cuda:{local_rank}"
|
||
else:
|
||
args.device = "cpu"
|
||
args.training_backend = "DeepSpeed-ZeRO"
|
||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else local_rank))
|
||
return local_rank
|
||
|
||
|
||
def train_epoch(args, lm_config, model_engine, loader, iters, start_step=0, wandb=None):
|
||
start_time = time.time()
|
||
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
|
||
input_ids = input_ids.to(args.device, non_blocking=True)
|
||
labels = labels.to(args.device, non_blocking=True)
|
||
|
||
res = model_engine(input_ids, labels=labels)
|
||
aux_loss = res.aux_loss if res.aux_loss is not None else torch.tensor(0.0, device=args.device)
|
||
loss = res.loss + aux_loss
|
||
|
||
model_engine.backward(loss)
|
||
model_engine.step()
|
||
|
||
if step % args.log_interval == 0 or step == iters:
|
||
spend_time = time.time() - start_time
|
||
current_loss = loss.item()
|
||
current_aux_loss = aux_loss.item()
|
||
current_logits_loss = current_loss - current_aux_loss
|
||
lr_list = model_engine.get_lr()
|
||
current_lr = lr_list[0] if lr_list else args.learning_rate
|
||
eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
|
||
metrics = reduce_metrics({
|
||
"loss": current_loss,
|
||
"logits_loss": current_logits_loss,
|
||
"aux_loss": current_aux_loss,
|
||
})
|
||
peak_memory_gb = get_cuda_peak_memory_gb()
|
||
memory_text = f", peak_memory: {peak_memory_gb:.2f} GB" if peak_memory_gb is not None else ""
|
||
Logger(
|
||
f"Epoch:[{args.current_epoch + 1}/{args.epochs}]({step}/{iters}), "
|
||
f"loss: {metrics['loss']:.4f}, logits_loss: {metrics['logits_loss']:.4f}, "
|
||
f"aux_loss: {metrics['aux_loss']:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min"
|
||
f"{memory_text}"
|
||
)
|
||
if wandb and is_main_process():
|
||
wandb.log(
|
||
{
|
||
"loss": metrics["loss"],
|
||
"logits_loss": metrics["logits_loss"],
|
||
"aux_loss": metrics["aux_loss"],
|
||
"learning_rate": current_lr,
|
||
"epoch_time": eta_min,
|
||
}
|
||
)
|
||
|
||
if step % args.save_interval == 0 or step == iters:
|
||
model_engine.eval()
|
||
save_deepspeed_checkpoint(
|
||
model_engine,
|
||
lm_config,
|
||
weight=args.save_weight,
|
||
epoch=args.current_epoch,
|
||
step=step,
|
||
wandb=wandb,
|
||
save_dir="../checkpoints",
|
||
)
|
||
if is_main_process():
|
||
ckp = save_model_weights(lm_config, model_engine, save_dir=args.save_dir, weight=args.save_weight)
|
||
Logger(f"Lightweight model saved: {ckp}")
|
||
model_engine.train()
|
||
|
||
del input_ids, labels, res, loss, aux_loss
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description="Full SFT with DeepSpeed ZeRO")
|
||
parser.add_argument("--save_dir", type=str, default="../out")
|
||
parser.add_argument("--save_weight", default="full_sft", type=str)
|
||
parser.add_argument("--epochs", type=int, default=2)
|
||
parser.add_argument("--batch_size", type=int, default=48)
|
||
parser.add_argument("--learning_rate", type=float, default=1e-5)
|
||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||
parser.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16", "float32"])
|
||
parser.add_argument("--num_workers", type=int, default=8)
|
||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||
parser.add_argument("--log_interval", type=int, default=100)
|
||
parser.add_argument("--save_interval", type=int, default=1000)
|
||
parser.add_argument("--hidden_size", default=768, type=int)
|
||
parser.add_argument("--num_hidden_layers", default=8, type=int)
|
||
parser.add_argument("--num_attention_heads", default=8, type=int)
|
||
parser.add_argument("--num_key_value_heads", default=4, type=int)
|
||
parser.add_argument("--intermediate_size", default=None, type=int)
|
||
parser.add_argument("--max_seq_len", default=768, type=int)
|
||
parser.add_argument("--use_moe", default=0, type=int, choices=[0, 1])
|
||
parser.add_argument("--attention_type", default="gqa", choices=["gqa", "mha", "mqa", "mla"])
|
||
parser.add_argument("--use_mla", default=0, type=int, choices=[0, 1])
|
||
parser.add_argument("--kv_lora_rank", default=128, type=int)
|
||
parser.add_argument("--q_lora_rank", default=256, type=int)
|
||
parser.add_argument("--rope_dim", default=None, type=int)
|
||
parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl")
|
||
parser.add_argument("--from_weight", default="pretrain", type=str)
|
||
parser.add_argument("--from_resume", default=0, type=int, choices=[0, 1])
|
||
parser.add_argument("--overwrite_output", default=0, type=int, choices=[0, 1])
|
||
parser.add_argument("--use_wandb", action="store_true")
|
||
parser.add_argument("--wandb_project", type=str, default="SearchLM-Full-SFT-DS")
|
||
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1])
|
||
add_model_profile_args(parser)
|
||
add_deepspeed_args(parser)
|
||
return parser.parse_args()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
args = parse_args()
|
||
model_profile = apply_model_profile(args)
|
||
|
||
try:
|
||
import deepspeed
|
||
except ImportError as exc:
|
||
raise RuntimeError("DeepSpeed未安装。服务器环境执行: pip install deepspeed") from exc
|
||
|
||
local_rank = init_deepspeed_runtime(args, deepspeed)
|
||
os.makedirs(args.save_dir, exist_ok=True)
|
||
|
||
lm_config = build_lm_config(
|
||
hidden_size=args.hidden_size,
|
||
num_hidden_layers=args.num_hidden_layers,
|
||
use_moe=bool(args.use_moe),
|
||
attention_type=resolve_attention_type(args),
|
||
kv_lora_rank=args.kv_lora_rank,
|
||
q_lora_rank=args.q_lora_rank,
|
||
rope_dim=args.rope_dim,
|
||
num_attention_heads=args.num_attention_heads,
|
||
num_key_value_heads=args.num_key_value_heads,
|
||
intermediate_size=args.intermediate_size,
|
||
)
|
||
assert_no_training_output_collision(
|
||
lm_config,
|
||
weight=args.save_weight,
|
||
save_dir=args.save_dir,
|
||
checkpoint_dir="../checkpoints",
|
||
from_resume=args.from_resume,
|
||
overwrite=args.overwrite_output,
|
||
)
|
||
|
||
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
|
||
if args.use_compile == 1:
|
||
model = torch.compile(model)
|
||
Logger("torch.compile enabled before DeepSpeed initialize")
|
||
|
||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||
loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
|
||
iters = len(loader_for_count)
|
||
total_optimizer_steps = max(1, math.ceil(iters / args.accumulation_steps) * args.epochs)
|
||
ds_config = load_deepspeed_config(args, total_optimizer_steps)
|
||
|
||
model_engine, _, _, _ = deepspeed.initialize(
|
||
model=model,
|
||
model_parameters=model.parameters(),
|
||
config=ds_config,
|
||
)
|
||
|
||
ckp_data = None
|
||
if args.from_resume == 1:
|
||
ckp_data = load_deepspeed_checkpoint(
|
||
model_engine,
|
||
lm_config,
|
||
weight=args.save_weight,
|
||
save_dir="../checkpoints",
|
||
)
|
||
|
||
wandb = None
|
||
if args.use_wandb and is_main_process():
|
||
import swanlab as wandb
|
||
|
||
wandb_id = ckp_data.get("wandb_id") if ckp_data else None
|
||
wandb.init(
|
||
project=args.wandb_project,
|
||
name=f"SearchLM-Full-SFT-DS-{args.attention_type}-{args.hidden_size}",
|
||
id=wandb_id,
|
||
resume="must" if wandb_id else None,
|
||
)
|
||
|
||
extra = {
|
||
"model_profile": args.model_profile if model_profile else "",
|
||
"deepspeed_config": args.deepspeed_config,
|
||
"zero_stage": ds_config.get("zero_optimization", {}).get("stage", "none"),
|
||
"total_optimizer_steps": total_optimizer_steps,
|
||
}
|
||
log_training_setup(
|
||
args,
|
||
lm_config,
|
||
stage="full_sft_deepspeed",
|
||
dataset_len=len(train_ds),
|
||
iters=iters,
|
||
tokens_per_sample=args.max_seq_len,
|
||
extra=extra,
|
||
)
|
||
|
||
start_epoch = ckp_data.get("epoch", 0) if ckp_data else 0
|
||
start_step = ckp_data.get("step", 0) if ckp_data else 0
|
||
|
||
for epoch in range(start_epoch, args.epochs):
|
||
args.current_epoch = epoch
|
||
train_sampler and train_sampler.set_epoch(epoch)
|
||
setup_seed(42 + epoch + local_rank)
|
||
indices = torch.randperm(len(train_ds)).tolist()
|
||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||
loader = DataLoader(
|
||
train_ds,
|
||
batch_sampler=batch_sampler,
|
||
num_workers=args.num_workers,
|
||
pin_memory=True,
|
||
)
|
||
if skip > 0:
|
||
Logger(f"Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始")
|
||
train_epoch(args, lm_config, model_engine, loader, len(loader) + skip, skip, wandb)
|
||
|
||
if dist.is_initialized():
|
||
dist.destroy_process_group()
|