1041 lines
36 KiB
Python
1041 lines
36 KiB
Python
from pathlib import Path
|
||
import sys
|
||
|
||
project_dir = Path(__file__).parent.parent
|
||
project_path = project_dir.as_posix()
|
||
if project_path not in sys.path:
|
||
sys.path.insert(0, project_path)
|
||
|
||
import argparse
|
||
from collections import Counter
|
||
from contextlib import nullcontext
|
||
from dataclasses import asdict, dataclass
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from typing import List, Literal, Union
|
||
import warnings
|
||
|
||
import datasets
|
||
import numpy as np
|
||
import pandas as pd
|
||
import torch
|
||
from torch import nn, optim
|
||
import torch.distributed as dist
|
||
from torch.nn import functional as F
|
||
from torch.nn.parallel import DistributedDataParallel
|
||
from torch.utils.data import DataLoader, Dataset, DistributedSampler
|
||
from tqdm import tqdm, trange
|
||
from transformers import (
|
||
AutoModel,
|
||
AutoModelForCausalLM,
|
||
AutoModelForSequenceClassification,
|
||
AutoTokenizer,
|
||
BatchEncoding,
|
||
BertModel,
|
||
BertTokenizerFast,
|
||
PreTrainedModel,
|
||
PreTrainedTokenizerBase,
|
||
TextStreamer,
|
||
)
|
||
import wandb
|
||
|
||
from model.model_lora import apply_lora, load_lora
|
||
from model.model_minimind import MiniMindConfig
|
||
from sentence_transformers import SentenceTransformer
|
||
from trainer.trainer_utils import (
|
||
Logger,
|
||
MiniMindForCausalLM,
|
||
SkipBatchSampler,
|
||
init_distributed_mode,
|
||
is_main_process,
|
||
lm_checkpoint,
|
||
setup_seed,
|
||
)
|
||
|
||
warnings.filterwarnings("ignore")
|
||
|
||
|
||
# ----- dataset utils -----
|
||
|
||
|
||
def save_jsonl(lines: Union[list[dict], pd.DataFrame], path, overwrite=False):
|
||
if isinstance(lines, pd.DataFrame):
|
||
lines = lines.to_dict(orient="records")
|
||
|
||
path = Path(path)
|
||
if not overwrite:
|
||
assert not path.exists()
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(path, "w") as f:
|
||
for i, line in enumerate(lines):
|
||
json_line = json.dumps(line, ensure_ascii=False, separators=(",", ":"))
|
||
if i == 0:
|
||
f.write(json_line)
|
||
else:
|
||
f.write("\n" + json_line)
|
||
print(f"Saved {len(lines)} lines to {path.as_posix()}")
|
||
|
||
|
||
def load_jsonl(path):
|
||
data = pd.read_json(path_or_buf=path, lines=True)
|
||
return data
|
||
|
||
|
||
def load_dataset(
|
||
dataset_name="ChineseSquad",
|
||
max_query_len=64,
|
||
max_passage_len=128,
|
||
split=True,
|
||
):
|
||
if dataset_name == "ChineseSquad":
|
||
dataset_name = "lighteval/ChineseSquad"
|
||
elif dataset_name == "natural-questions":
|
||
dataset_name = "sentence-transformers/natural-questions"
|
||
else:
|
||
raise NotImplementedError
|
||
|
||
dataset = datasets.load_dataset(dataset_name)
|
||
if dataset_name == "sentence-transformers/natural-questions":
|
||
dataset = dataset.filter(
|
||
lambda x: len(x["query"]) <= max_query_len
|
||
and len(x["answer"]) <= max_passage_len
|
||
)
|
||
if split:
|
||
dataset = dataset["train"].train_test_split(test_size=0.05, seed=42)
|
||
else:
|
||
dataset = dataset.filter(
|
||
lambda x: len(x["question"]) <= max_query_len
|
||
and len(x["context"]) <= max_passage_len
|
||
)
|
||
return dataset
|
||
|
||
|
||
class DPRDataset(Dataset):
|
||
def __init__(
|
||
self,
|
||
dataset: datasets.Dataset,
|
||
tokenizer: PreTrainedTokenizerBase,
|
||
max_query_len=64,
|
||
max_passage_len: int = 128,
|
||
):
|
||
super().__init__()
|
||
self.tokenizer = tokenizer
|
||
self.max_query_len = max_query_len
|
||
self.max_passage_len = max_passage_len
|
||
self.dataset = dataset
|
||
if "query" in self.dataset.features:
|
||
self.dataset_name = "nq"
|
||
else:
|
||
self.dataset_name = "cs"
|
||
|
||
def __len__(self):
|
||
return len(self.dataset)
|
||
|
||
def __getitem__(self, index):
|
||
sample = self.dataset[index]
|
||
|
||
def encode(key, max_length):
|
||
if isinstance(self.tokenizer, BertTokenizerFast):
|
||
text = sample[key]
|
||
else:
|
||
text = self.tokenizer.bos_token + sample[key] + self.tokenizer.eos_token
|
||
inputs = self.tokenizer(
|
||
text,
|
||
max_length=max_length,
|
||
truncation=True,
|
||
return_tensors="pt",
|
||
padding_side="right",
|
||
padding="max_length",
|
||
return_token_type_ids=False,
|
||
)
|
||
inputs["input_ids"] = inputs["input_ids"].view(-1)
|
||
inputs["attention_mask"] = inputs["attention_mask"].view(-1)
|
||
return inputs
|
||
|
||
if self.dataset_name == "cs": # ChineseSquad
|
||
query_inputs = encode("question", self.max_query_len)
|
||
passage_inputs = encode("context", self.max_passage_len)
|
||
else: # natural-questions
|
||
query_inputs = encode("query", self.max_query_len)
|
||
passage_inputs = encode("answer", self.max_passage_len)
|
||
return query_inputs, passage_inputs
|
||
|
||
def tolist(self):
|
||
if self.dataset_name == "cs": # ChineseSquad
|
||
queries = list(self.dataset["question"])
|
||
passages = list(self.dataset["context"])
|
||
else:
|
||
queries = list(self.dataset["query"])
|
||
passages = list(self.dataset["answer"])
|
||
return queries, passages
|
||
|
||
|
||
# ----- retriever -----
|
||
|
||
|
||
class Encoder(nn.Module):
|
||
"""
|
||
Adapted from CMU-11-667 (https://cmu-llms.org/) EncoderModel
|
||
"""
|
||
|
||
def __init__(self, encoder: PreTrainedModel, temperature: float = 1.0):
|
||
super().__init__()
|
||
self.encoder = encoder
|
||
self.register_buffer("temperature", torch.tensor(temperature))
|
||
|
||
def forward(
|
||
self,
|
||
query_inputs: BatchEncoding = None,
|
||
passage_inputs: BatchEncoding = None,
|
||
):
|
||
q_embeddings = (
|
||
self.encode(query_inputs) if query_inputs else None
|
||
) # (n_queries, hidden_dim)
|
||
p_embeddings = (
|
||
self.encode(passage_inputs) if passage_inputs else None
|
||
) # (n_passages, hidden_dim)
|
||
|
||
# for inference
|
||
if q_embeddings is None or p_embeddings is None:
|
||
return q_embeddings, p_embeddings
|
||
|
||
similarity = (
|
||
torch.matmul(q_embeddings, p_embeddings.T) / self.temperature
|
||
) # (n_queries, n_passages)
|
||
target = torch.arange(len(q_embeddings)).to(q_embeddings.device) # (n_queries,)
|
||
loss = F.cross_entropy(similarity, target)
|
||
return loss, similarity
|
||
|
||
def encode_text(
|
||
self,
|
||
texts: Union[str, List[str]],
|
||
tokenizer: PreTrainedTokenizerBase,
|
||
max_length=128,
|
||
):
|
||
if isinstance(texts, str):
|
||
texts = [texts]
|
||
texts = [tokenizer.bos_token + text + tokenizer.eos_token for text in texts]
|
||
inputs = tokenizer(
|
||
texts,
|
||
max_length=max_length,
|
||
truncation=True,
|
||
return_tensors="pt",
|
||
padding_side="right",
|
||
padding="max_length",
|
||
return_token_type_ids=False,
|
||
)
|
||
embeddings = self.encode(inputs.to(self.encoder.device))
|
||
return embeddings
|
||
|
||
def encode(self, inputs: BatchEncoding):
|
||
hidden_states = self.encoder(
|
||
**inputs
|
||
).last_hidden_state # (batch_size, seq_len, hidden_dim)
|
||
return self.pooling(hidden_states, inputs["attention_mask"])
|
||
|
||
def pooling(self, last_hidden_state, attention_mask):
|
||
"""
|
||
last_hidden_state: (batch_size, seq_len, hidden_dim)
|
||
attention_mask: (batch_size, seq_len)
|
||
embeddings: (batch_size, hidden_dim)
|
||
"""
|
||
if isinstance(self.encoder, BertModel):
|
||
last_hidden_state = last_hidden_state[:, 0, :] # (batch_size, hidden_dim)
|
||
else:
|
||
attention_mask = torch.where(
|
||
attention_mask == 1,
|
||
attention_mask
|
||
+ torch.arange(attention_mask.shape[1]).to(attention_mask.device),
|
||
0,
|
||
) # (batch_size, seq_len)
|
||
last_nonzero_indices = attention_mask.argmax(
|
||
dim=1, keepdim=True
|
||
) # (batch_size, 1)
|
||
last_token_indices = last_nonzero_indices.unsqueeze(-1).expand(
|
||
-1, -1, last_hidden_state.shape[-1]
|
||
) # (batch_size, 1, hidden_dim)
|
||
last_hidden_state = torch.gather(
|
||
last_hidden_state, 1, last_token_indices
|
||
) # (batch_size, 1, hidden_dim)
|
||
last_hidden_state = last_hidden_state.squeeze(1) # (batch_size, hidden_dim)
|
||
embeddings = F.normalize(last_hidden_state, p=2, dim=1)
|
||
return embeddings
|
||
|
||
|
||
@torch.inference_mode()
|
||
def similarity_bert(
|
||
model_bert: SentenceTransformer, queries: List[str], passages: List[str]
|
||
):
|
||
if isinstance(queries, str):
|
||
queries = [queries]
|
||
if isinstance(passages, str):
|
||
passages = [passages]
|
||
q_embeddings = model_bert.encode(queries) # (n_queries, hidden_dim)
|
||
p_embeddings = model_bert.encode(passages) # (n_passages, hidden_dim)
|
||
similarity = model_bert.similarity(
|
||
q_embeddings, p_embeddings
|
||
) # (n_queries, n_passages)
|
||
target = torch.arange(len(q_embeddings)).to(q_embeddings.device) # (n_queries,)
|
||
loss = F.cross_entropy(similarity, target)
|
||
return loss, similarity
|
||
|
||
|
||
@torch.inference_mode()
|
||
def calc_relative_advantage(similarity):
|
||
"""
|
||
similarity: (n_queries, n_passages)
|
||
"""
|
||
sorted_similarity = torch.sort(similarity, 1, descending=True)[0]
|
||
second_best = sorted_similarity[:, 1]
|
||
return (torch.diag(similarity) - second_best) / second_best
|
||
|
||
|
||
@torch.inference_mode()
|
||
def benchmark_retriever(
|
||
model: Union[Literal["bm25"], Encoder],
|
||
dataset: DPRDataset,
|
||
top_k=[5, 20, 40, 60, 80, 100],
|
||
batch_size=256,
|
||
):
|
||
if not np.iterable(top_k):
|
||
top_k = [top_k]
|
||
max_top_k = max(top_k)
|
||
top_k = sorted(top_k)
|
||
queries, passages = dataset.tolist()
|
||
unique_queries = []
|
||
unique_passages = []
|
||
for i in range(len(passages)):
|
||
if passages[i] not in unique_passages:
|
||
unique_queries.append(queries[i])
|
||
unique_passages.append(passages[i])
|
||
queries, passages = unique_queries, unique_passages
|
||
if isinstance(model, str):
|
||
if model == "bm25":
|
||
import bm25s
|
||
|
||
corpus_tokens = bm25s.tokenize(passages, stopwords="en")
|
||
retriever = bm25s.BM25(k1=0.9, b=0.4)
|
||
retriever.index(corpus_tokens)
|
||
query_tokens = bm25s.tokenize(queries)
|
||
indices = retriever.retrieve(query_tokens, k=max_top_k)[0]
|
||
else:
|
||
raise NotImplementedError
|
||
elif isinstance(model, Encoder):
|
||
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
|
||
all_q_embeddings = []
|
||
all_p_embeddings = []
|
||
device = model.encoder.device
|
||
for query_inputs, passage_inputs in tqdm(loader):
|
||
q_embeddings = model.encode(
|
||
query_inputs.to(device)
|
||
) # (batch_size, hidden_dim)
|
||
p_embeddings = model.encode(
|
||
passage_inputs.to(device)
|
||
) # (batch_size, hidden_dim)
|
||
all_q_embeddings.append(q_embeddings)
|
||
all_p_embeddings.append(p_embeddings)
|
||
q_embeddings = torch.cat(all_q_embeddings, dim=0) # (n_queries, hidden_dim)
|
||
p_embeddings = torch.cat(all_p_embeddings, dim=0) # (n_passages, hidden_dim)
|
||
similarity = q_embeddings @ p_embeddings.T # (n_queries, n_passages)
|
||
values, indices = torch.topk(similarity, k=max_top_k, dim=1)
|
||
indices = indices.cpu().numpy()
|
||
else: # sentence-transformers
|
||
similarity = similarity_bert(model, queries, passages)[1]
|
||
values, indices = torch.topk(similarity, k=max_top_k, dim=1)
|
||
indices = indices.cpu().numpy()
|
||
targets = np.arange(len(indices))[:, None]
|
||
accuracy = {}
|
||
for _top_k in top_k:
|
||
acc = np.any(indices[:, :_top_k] == targets, 1).mean().item()
|
||
accuracy[_top_k] = acc
|
||
# torch.cuda.empty_cache()
|
||
return accuracy
|
||
|
||
|
||
class BM25:
|
||
"""
|
||
bm25 = BM25(k1=0.9, b=0.4)
|
||
bm.index(corpus)
|
||
bm.retrieve(queries, top_k=100)
|
||
"""
|
||
|
||
def __init__(self, pattern=r"(?u)\b\w\w+\b", k1=0.9, b=0.4):
|
||
"""
|
||
The default `k1` and `b` are consistent with Karpukhin et al. (2020).
|
||
"""
|
||
self.pattern = re.compile(pattern)
|
||
self.k1 = k1
|
||
self.b = b
|
||
|
||
def split(self, text: str):
|
||
text = text.lower()
|
||
text = self.pattern.findall(text)
|
||
return text
|
||
|
||
def index(self, corpus: List[str]):
|
||
self.chunk_lens = []
|
||
self.freqs = []
|
||
token_to_num_chunks = Counter()
|
||
if isinstance(corpus, str):
|
||
corpus = [corpus]
|
||
for text in tqdm(corpus):
|
||
tokens = self.split(text)
|
||
chunk_len = len(tokens)
|
||
freqs = {}
|
||
for token in tokens:
|
||
if token not in freqs:
|
||
freqs[token] = 1
|
||
token_to_num_chunks[token] += 1
|
||
else:
|
||
freqs[token] += 1
|
||
self.chunk_lens.append(chunk_len)
|
||
self.freqs.append(freqs)
|
||
self.num_chunks = len(self.chunk_lens)
|
||
self.mean_chunk_len = np.mean(self.chunk_lens).item()
|
||
self.token_to_num_chunks = dict(token_to_num_chunks)
|
||
|
||
def retrieve(self, queries: List[str], top_k=100):
|
||
if isinstance(queries, str):
|
||
queries = [queries]
|
||
results = []
|
||
for query in tqdm(queries):
|
||
result = self._retrieve(query, top_k)
|
||
results.append(result)
|
||
indices = np.stack([result[0] for result in results])
|
||
scores = np.stack([result[1] for result in results])
|
||
return indices, scores
|
||
|
||
def _retrieve(self, query: str, top_k=100):
|
||
if top_k is None:
|
||
top_k = self.num_chunks
|
||
assert top_k <= self.num_chunks
|
||
query_tokens = self.split(query)
|
||
scores = np.zeros(self.num_chunks, dtype=float)
|
||
for chunk_index in range(self.num_chunks):
|
||
score = self.score(query_tokens, chunk_index)
|
||
scores[chunk_index] = score
|
||
indices = np.argsort(-scores)[:top_k]
|
||
scores = scores[indices]
|
||
return indices, scores
|
||
|
||
def score(self, query_tokens: List[str], chunk_index: int):
|
||
chunk_len = self.chunk_lens[chunk_index]
|
||
n = np.array(
|
||
list(
|
||
map(lambda token: self.token_to_num_chunks.get(token, 0), query_tokens)
|
||
)
|
||
)
|
||
freqs = np.array(
|
||
list(map(lambda token: self.freqs[chunk_index].get(token, 0), query_tokens))
|
||
)
|
||
IDF = self.calc_IDF(n)
|
||
weight = freqs / (
|
||
freqs + self.k1 * (1 - self.b + self.b * chunk_len / self.mean_chunk_len)
|
||
)
|
||
return np.sum(IDF * weight)
|
||
|
||
def calc_IDF(self, n: int):
|
||
"""
|
||
n: number of chunks containing the term
|
||
"""
|
||
return np.log((self.num_chunks - n + 0.5) / (n + 0.5) + 1)
|
||
|
||
|
||
# ----- train retriever -----
|
||
|
||
|
||
@dataclass
|
||
class TrainArgs:
|
||
save_dir: Union[str, Path] = project_dir / "checkpoints" # 模型保存目录
|
||
save_weight: str = "dpr" # 保存权重的前缀名
|
||
epochs: int = 1 # 训练轮数
|
||
batch_size: int = 32
|
||
learning_rate: float = 4e-6
|
||
weight_decay: float = 0.0
|
||
betas: tuple = (0.9, 0.999)
|
||
device: str = None
|
||
dtype: str = "bfloat16"
|
||
num_workers: int = 0 # 数据加载线程数
|
||
accumulation_steps: int = 1 # 梯度累积步数
|
||
grad_clip: float = 1.0 # 梯度裁剪阈值
|
||
log_interval: int = 100 # 日志打印间隔
|
||
save_interval: int = 100 # 模型保存间隔
|
||
hidden_size: int = 512 # 隐藏层维度
|
||
num_hidden_layers: int = 8 # 隐藏层数量
|
||
max_query_len: int = 64 # 问题最大长度
|
||
max_passage_len: int = 128 # 单个document的最大长度，对于`natural-questions`数据集，建议设置为512
|
||
use_amp: bool = None # 使用自动混合精度
|
||
temperature: float = 1.0
|
||
from_weight: str = "pretrain" # 基于哪个权重训练，为none则从头开始
|
||
from_resume: int = 0 # 是否自动检测&续训（0=否，1=是）
|
||
dataset: str = "ChineseSquad" # 训练数据集（同时也是检索数据集），`ChineseSquad`或`natural-questions`
|
||
use_wandb: bool = False
|
||
use_swanlab: bool = False # 使用swanlab时需要同时设置use_wandb=True, wandb会被替换为swanlab
|
||
train_bert: bool = False # 是否训练BERT-Mini，若是则加载BERT，则MiniMind模型和权重均失效
|
||
wandb_entity: str = None
|
||
wandb_project: str = "MiniMind"
|
||
|
||
def __post_init__(self):
|
||
assert self.dtype in ["bfloat16", "float16"]
|
||
assert self.dataset in ["ChineseSquad", "natural-questions"]
|
||
if self.device is None:
|
||
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||
if self.use_amp is None:
|
||
self.use_amp = "cuda" in self.device
|
||
self.save_dir = Path(self.save_dir).as_posix()
|
||
if self.train_bert:
|
||
self.hidden_size = 256 # BERT-mini
|
||
self.model_dir = project_dir / "out"
|
||
if not self.model_dir.exists():
|
||
self.model_dir = project_dir / "checkpoints"
|
||
self.model_dir = self.model_dir.as_posix()
|
||
self.wandb = wandb
|
||
if self.use_swanlab:
|
||
import swanlab
|
||
|
||
self.wandb = swanlab
|
||
|
||
|
||
def get_num_params(model):
|
||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
|
||
|
||
def init_model(
|
||
lm_config=None,
|
||
from_weight="pretrain",
|
||
tokenizer_path="../model",
|
||
save_dir="../out",
|
||
device="cuda",
|
||
bert=False,
|
||
eval_args=None,
|
||
):
|
||
if eval_args is not None:
|
||
tokenizer = AutoTokenizer.from_pretrained(eval_args.load_from)
|
||
if "model" in eval_args.load_from:
|
||
model = MiniMindForCausalLM(
|
||
MiniMindConfig(
|
||
hidden_size=eval_args.hidden_size,
|
||
num_hidden_layers=eval_args.num_hidden_layers,
|
||
use_moe=bool(eval_args.use_moe),
|
||
inference_rope_scaling=eval_args.inference_rope_scaling,
|
||
)
|
||
)
|
||
moe_suffix = "_moe" if eval_args.use_moe else ""
|
||
ckp = f"{eval_args.save_dir}/{eval_args.weight}_{eval_args.hidden_size}{moe_suffix}.pth"
|
||
model.load_state_dict(
|
||
torch.load(ckp, map_location=eval_args.device), strict=True
|
||
)
|
||
if eval_args.lora_weight != "None":
|
||
apply_lora(model)
|
||
load_lora(
|
||
model,
|
||
f"{eval_args.save_dir}/lora/{eval_args.lora_weight}_{eval_args.hidden_size}.pth",
|
||
)
|
||
else:
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
eval_args.load_from, trust_remote_code=True
|
||
)
|
||
else:
|
||
if bert:
|
||
model = AutoModel.from_pretrained("prajjwal1/bert-mini")
|
||
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-mini")
|
||
else:
|
||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||
model = MiniMindForCausalLM(lm_config)
|
||
model = Encoder(model, temperature=1.0)
|
||
|
||
if from_weight != "none":
|
||
moe_suffix = "_moe" if lm_config.use_moe else ""
|
||
weight_path = (
|
||
f"{save_dir}/{from_weight}_{lm_config.hidden_size}{moe_suffix}.pth"
|
||
)
|
||
weights = torch.load(weight_path, map_location=device)
|
||
model.load_state_dict(weights, strict=False)
|
||
|
||
Logger(
|
||
f"所加载Model可训练参数：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万"
|
||
)
|
||
return model.to(device), tokenizer
|
||
|
||
|
||
def train_epoch(
|
||
model: MiniMindForCausalLM,
|
||
args: TrainArgs,
|
||
lm_config: MiniMindConfig,
|
||
optimizer: optim.Optimizer,
|
||
scheduler: optim.lr_scheduler.LRScheduler,
|
||
scaler: torch.GradScaler,
|
||
epoch: int,
|
||
loader: DataLoader,
|
||
iters: int,
|
||
bar,
|
||
start_step=0,
|
||
use_wandb=False,
|
||
spend_time=0,
|
||
):
|
||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||
grad_norm = float("nan")
|
||
adv = float("nan")
|
||
wandb = args.wandb
|
||
start_time = time.time()
|
||
|
||
def get_spend_time():
|
||
return spend_time + time.time() - start_time
|
||
|
||
for step, (query_inputs, passage_inputs) in enumerate(loader, start=start_step + 1):
|
||
query_inputs = query_inputs.to(args.device)
|
||
passage_inputs = passage_inputs.to(args.device)
|
||
|
||
with torch.autocast(device_type=args.device, enabled=args.use_amp, dtype=dtype):
|
||
loss, similarity = model(query_inputs, passage_inputs)
|
||
loss = loss / args.accumulation_steps
|
||
|
||
scaler.scale(loss).backward()
|
||
loss = loss.detach().cpu().numpy().item() * args.accumulation_steps
|
||
lr = optimizer.param_groups[-1]["lr"]
|
||
|
||
if step % args.accumulation_steps == 0:
|
||
scaler.unscale_(optimizer)
|
||
grad_norm = torch.nn.utils.clip_grad_norm_(
|
||
model.parameters(), args.grad_clip
|
||
)
|
||
grad_norm = grad_norm.detach().cpu().numpy().item()
|
||
|
||
scaler.step(optimizer)
|
||
scaler.update()
|
||
|
||
optimizer.zero_grad()
|
||
torch.cuda.empty_cache()
|
||
|
||
if use_wandb:
|
||
adv = (
|
||
calc_relative_advantage(similarity)
|
||
.mean()
|
||
.detach()
|
||
.cpu()
|
||
.numpy()
|
||
.item()
|
||
)
|
||
wandb.log(
|
||
{
|
||
"train_step": (epoch * len(loader) + step)
|
||
// args.accumulation_steps,
|
||
"train/epoch": epoch + 1,
|
||
"train/loss": loss,
|
||
"train/lr": lr,
|
||
"train/grad_norm": grad_norm,
|
||
"train/adv": adv,
|
||
"train/time": get_spend_time(),
|
||
}
|
||
)
|
||
scheduler.step()
|
||
if bar is not None:
|
||
bar.update()
|
||
bar.set_postfix_str(
|
||
f"[{epoch+1}/{args.epochs}]loss={loss:.4f},adv={adv:.4f},grad_norm={grad_norm:.4f}"
|
||
)
|
||
|
||
if step % args.log_interval == 0 or step == iters - 1:
|
||
Logger(f"Epoch:[{epoch+1}/{args.epochs},{step}/{iters}] loss:{loss:.4f}")
|
||
|
||
if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
|
||
model.eval()
|
||
ckp = f"{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}.pth"
|
||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||
state_dict = model.module.state_dict()
|
||
else:
|
||
state_dict = model.state_dict()
|
||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||
torch.save(state_dict, ckp)
|
||
wandb_id = None
|
||
if args.use_wandb:
|
||
wandb_id = getattr(wandb.run, "id", None)
|
||
lm_checkpoint(
|
||
lm_config,
|
||
weight=args.save_weight,
|
||
model=model,
|
||
optimizer=optimizer,
|
||
scaler=scaler,
|
||
scheduler=scheduler,
|
||
epoch=epoch,
|
||
step=step,
|
||
wandb_id=wandb_id,
|
||
save_dir=args.save_dir,
|
||
spend_time=get_spend_time(),
|
||
)
|
||
model.train()
|
||
return get_spend_time()
|
||
|
||
|
||
def train(
|
||
args: TrainArgs,
|
||
model: MiniMindForCausalLM = None,
|
||
tokenizer: PreTrainedTokenizerBase = None,
|
||
train_ds: DPRDataset = None,
|
||
early_return=False,
|
||
):
|
||
# ========== 1. 初始化环境和随机种子 ==========
|
||
local_rank = init_distributed_mode()
|
||
if dist.is_initialized():
|
||
args.device = f"cuda:{local_rank}"
|
||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||
|
||
# ========== 2. 配置目录、模型参数、检查ckp ==========
|
||
os.makedirs(args.save_dir, exist_ok=True)
|
||
lm_config = MiniMindConfig(
|
||
hidden_size=args.hidden_size,
|
||
num_hidden_layers=args.num_hidden_layers,
|
||
use_moe=False,
|
||
)
|
||
ckp_data = (
|
||
lm_checkpoint(lm_config, weight=args.save_weight, save_dir=args.save_dir)
|
||
if args.from_resume == 1
|
||
else None
|
||
)
|
||
|
||
# ========== 3. 设置混合精度 ==========
|
||
# set in TrainArgs
|
||
|
||
# ========== 4. 定义模型、数据、优化器 ==========
|
||
if model is None:
|
||
model, tokenizer = init_model(
|
||
lm_config,
|
||
args.from_weight,
|
||
tokenizer_path=project_dir / "model",
|
||
save_dir=args.model_dir,
|
||
device=args.device,
|
||
bert=args.train_bert,
|
||
)
|
||
# model = Encoder(model, temperature=args.temperature)
|
||
num_params = get_num_params(model)
|
||
if train_ds is None:
|
||
dataset = load_dataset(args.dataset, args.max_query_len, args.max_passage_len)
|
||
train_ds = DPRDataset(
|
||
dataset["train"],
|
||
tokenizer,
|
||
max_query_len=args.max_query_len,
|
||
max_passage_len=args.max_passage_len,
|
||
)
|
||
key = "validation" if args.dataset == "ChineseSquad" else "test"
|
||
test_ds = DPRDataset(
|
||
dataset[key],
|
||
tokenizer,
|
||
max_query_len=args.max_query_len,
|
||
max_passage_len=args.max_passage_len,
|
||
)
|
||
|
||
if early_return: # for debugging
|
||
return model, tokenizer, train_ds, test_ds
|
||
|
||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||
scaler = torch.GradScaler(args.device, enabled=(args.dtype == "float16"))
|
||
optimizer = optim.AdamW(
|
||
model.parameters(),
|
||
lr=args.learning_rate,
|
||
weight_decay=args.weight_decay,
|
||
betas=args.betas,
|
||
)
|
||
|
||
# ========== 5. 从ckp恢复状态 ==========
|
||
start_epoch, start_step = 0, 0
|
||
spend_time = 0
|
||
if ckp_data:
|
||
model.load_state_dict(ckp_data["model"])
|
||
optimizer.load_state_dict(ckp_data["optimizer"])
|
||
scaler.load_state_dict(ckp_data["scaler"])
|
||
start_epoch = ckp_data["epoch"]
|
||
start_step = ckp_data.get("step", 0)
|
||
spend_time = ckp_data.get("spend_time", 0)
|
||
|
||
if start_step > 0: # 第一个epoch且存在检查点
|
||
batch_sampler = SkipBatchSampler(
|
||
train_sampler or range(len(train_ds)), args.batch_size, start_step + 1
|
||
)
|
||
loader = DataLoader(
|
||
train_ds,
|
||
batch_sampler=batch_sampler,
|
||
num_workers=args.num_workers,
|
||
pin_memory=True,
|
||
)
|
||
else: # 默认从头开始
|
||
loader = DataLoader(
|
||
train_ds,
|
||
batch_size=args.batch_size,
|
||
shuffle=(train_sampler is None),
|
||
sampler=train_sampler,
|
||
num_workers=args.num_workers,
|
||
pin_memory=True,
|
||
)
|
||
iters = len(loader)
|
||
max_step = args.epochs * iters
|
||
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_step)
|
||
if start_step > 0:
|
||
scheduler.load_state_dict(ckp_data["scheduler"])
|
||
|
||
# ========== 6. DDP包模型 ==========
|
||
if dist.is_initialized():
|
||
model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
|
||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||
|
||
# ========== 7. 配wandb ==========
|
||
run = nullcontext()
|
||
use_wandb = args.use_wandb and is_main_process()
|
||
if use_wandb:
|
||
wandb = args.wandb
|
||
wandb_id = ckp_data.get("wandb_id") if ckp_data else None
|
||
resume = "must" if wandb_id else None
|
||
wandb_run_name = f"MiniMind-{round(num_params/1e6)}M-{args.save_weight}-epoch-{args.epochs}-batchsize-{args.batch_size}-lr-{args.learning_rate}"
|
||
if wandb_id:
|
||
config = None
|
||
else:
|
||
config = asdict(args)
|
||
for key in [
|
||
"device",
|
||
"from_resume",
|
||
"use_wandb",
|
||
"wandb_project",
|
||
"save_dir",
|
||
]:
|
||
config.pop(key)
|
||
config["max_step"] = max_step
|
||
config["num_params"] = num_params
|
||
wandb_kwargs = dict(
|
||
entity=args.wandb_entity,
|
||
project=args.wandb_project,
|
||
name=wandb_run_name,
|
||
id=wandb_id,
|
||
resume=resume,
|
||
config=config,
|
||
)
|
||
if args.use_swanlab:
|
||
wandb_kwargs.pop("entity")
|
||
run = wandb.init(**wandb_kwargs)
|
||
if args.use_swanlab:
|
||
run = nullcontext()
|
||
else:
|
||
wandb.define_metric("train_step")
|
||
wandb.define_metric("eval_step")
|
||
wandb.define_metric("train/*", step_metric="train_step")
|
||
wandb.define_metric("eval/*", step_metric="eval_step")
|
||
|
||
# ========== 8. 开始训练 ==========
|
||
bar = tqdm(total=max_step, disable=not is_main_process())
|
||
with run:
|
||
for epoch in range(start_epoch, args.epochs):
|
||
train_sampler and train_sampler.set_epoch(epoch)
|
||
if epoch == start_epoch and start_step > 0: # 第一个epoch且存在检查点
|
||
Logger(
|
||
f"Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始"
|
||
)
|
||
# train_epoch(epoch, loader, len(loader) + start_step + 1, start_step, wandb)
|
||
_iters = iters + start_step + 1
|
||
bar.update(start_step)
|
||
else: # 默认从头开始
|
||
_iters = iters
|
||
spend_time = train_epoch(
|
||
model=model,
|
||
args=args,
|
||
lm_config=lm_config,
|
||
optimizer=optimizer,
|
||
scheduler=scheduler,
|
||
scaler=scaler,
|
||
epoch=epoch,
|
||
loader=loader,
|
||
iters=_iters,
|
||
start_step=start_step,
|
||
bar=bar,
|
||
use_wandb=use_wandb,
|
||
spend_time=spend_time,
|
||
)
|
||
bar.close()
|
||
|
||
|
||
# ----- RAG -----
|
||
|
||
|
||
@dataclass
|
||
class EvalArgs:
|
||
load_from: str = project_dir / "model"
|
||
save_dir: str = "out"
|
||
weight: str = "full_sft"
|
||
lora_weight: str = "None"
|
||
hidden_size: int = 768 # chat模型隐藏层维度(与retriever无关)
|
||
num_hidden_layers: int = 16 # chat模型隐藏层数量(与retriever无关)
|
||
use_moe: int = 0
|
||
inference_rope_scaling: bool = False
|
||
max_new_tokens: int = 8192
|
||
temperature: float = 0.85
|
||
top_p: float = 0.85
|
||
top_k: int = 20 # 检索时返回的文档数量，reranker会从中选择最终答案
|
||
rag: bool = True # 是否开启RAG
|
||
max_query_len: int = 64
|
||
max_passage_len: int = 128
|
||
historys: int = 0
|
||
use_sbert_retriever: bool = True # 使用sentence-transformer作为retriever，检索成功率会大幅提升
|
||
device: str = None
|
||
|
||
def __post_init__(self):
|
||
if self.device is None:
|
||
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||
self.save_dir = Path(self.save_dir)
|
||
if not self.save_dir.is_absolute():
|
||
self.save_dir = project_dir / self.save_dir
|
||
if not self.save_dir.exists():
|
||
self.save_dir = project_dir / "checkpoints"
|
||
self.save_dir = self.save_dir.as_posix()
|
||
if isinstance(self.load_from, Path):
|
||
self.load_from = self.load_from.as_posix()
|
||
|
||
|
||
@torch.inference_mode()
|
||
def init_rag(
|
||
device,
|
||
documents: Union[List[str], datasets.Dataset] = None,
|
||
retriever: Encoder = None,
|
||
tokenizer: PreTrainedTokenizerBase = None,
|
||
eval_args: EvalArgs = None,
|
||
):
|
||
print("loading models and encoding documents, this may take a while...")
|
||
if retriever is None or eval_args.use_sbert_retriever:
|
||
retriever = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
|
||
encode_fn = lambda retriever: retriever.encode(
|
||
documents, convert_to_tensor=True, show_progress_bar=True
|
||
)
|
||
else:
|
||
|
||
def encode_fn(retriever: Encoder):
|
||
all_p_embeddings = []
|
||
batch_size = 128
|
||
step = int(np.ceil(len(documents) / batch_size))
|
||
for i in trange(step):
|
||
docs = documents[i * batch_size : (i + 1) * batch_size]
|
||
p_embeddings = retriever.encode_text(
|
||
docs, tokenizer, max_length=eval_args.max_passage_len
|
||
)
|
||
all_p_embeddings.append(p_embeddings)
|
||
p_embeddings = torch.cat(
|
||
all_p_embeddings, dim=0
|
||
) # (n_passages, hidden_dim)
|
||
return p_embeddings
|
||
|
||
retriever.to(device)
|
||
retriever.eval()
|
||
reranker = AutoModelForSequenceClassification.from_pretrained(
|
||
"jinaai/jina-reranker-v2-base-multilingual", trust_remote_code=True
|
||
)
|
||
reranker.to(device)
|
||
reranker.eval()
|
||
# return retriever, reranker
|
||
if documents is None:
|
||
documents = load_dataset()
|
||
documents = documents["train"]["context"]
|
||
doc_embeddings = encode_fn(retriever)
|
||
return retriever, reranker, doc_embeddings
|
||
|
||
|
||
@torch.inference_mode()
|
||
def retrieve(
|
||
query: str,
|
||
retriever: Union[SentenceTransformer, Encoder],
|
||
reranker: PreTrainedModel,
|
||
documents: List[str],
|
||
doc_embeddings: torch.Tensor,
|
||
top_k=20,
|
||
tokenizer=None,
|
||
max_length=128,
|
||
):
|
||
"""
|
||
corpus_embeddings: (n_passages, hidden_dim)
|
||
"""
|
||
if isinstance(retriever, SentenceTransformer):
|
||
query_embeddings = retriever.encode(
|
||
query, convert_to_tensor=True
|
||
) # (hidden_dim,)
|
||
else:
|
||
query_embeddings = retriever.encode_text(
|
||
query, tokenizer, max_length=max_length
|
||
)[0]
|
||
similarity = doc_embeddings @ query_embeddings
|
||
values, indices = torch.topk(similarity, k=top_k)
|
||
indices = indices.cpu().numpy()
|
||
sentence_pairs = [[query, documents[i]] for i in indices]
|
||
scores = reranker.compute_score(sentence_pairs)
|
||
max_index = indices[np.argsort(scores)[-1]]
|
||
return max_index.item(), documents[max_index]
|
||
|
||
|
||
@torch.inference_mode()
|
||
def evaluate(
|
||
args: EvalArgs,
|
||
retriever: SentenceTransformer,
|
||
reranker: PreTrainedModel,
|
||
documents: List[str],
|
||
doc_embeddings: torch.Tensor,
|
||
):
|
||
prompts = [
|
||
"估计有4.88亿至5.35亿人信奉什么宗教",
|
||
"南安普敦机场在哪个城镇",
|
||
"路德什么时候死的",
|
||
"超级碗开幕之夜于何时在何地举行",
|
||
]
|
||
|
||
conversation = []
|
||
model, tokenizer = init_model(eval_args=args)
|
||
input_mode = int(input("[0] 自动测试\n[1] 手动输入\n"))
|
||
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||
|
||
prompt_iter = prompts if input_mode == 0 else iter(lambda: input("👶: "), "")
|
||
for prompt in prompt_iter:
|
||
setup_seed(2026) # or setup_seed(random.randint(0, 2048))
|
||
if input_mode == 0:
|
||
print(f"👶: {prompt}")
|
||
if args.rag:
|
||
doc = retrieve(
|
||
prompt,
|
||
retriever,
|
||
reranker,
|
||
documents,
|
||
doc_embeddings,
|
||
top_k=args.top_k,
|
||
tokenizer=tokenizer,
|
||
max_length=args.max_passage_len,
|
||
)[1]
|
||
prompt = f'你是一个智能问答助手，请严格按照以下要求回答问题： 如果资料中包含问题答案，请直接使用资料信息并注明"根据资料"，如果资料不相关、信息不足或未包含答案，请明确说明"资料中未包含相关信息"，然后可以基于常识进行补充\n资料：{doc}\n问题：{prompt}\n现在请开始回答：'
|
||
conversation = conversation[-args.historys :] if args.historys else []
|
||
conversation.append({"role": "user", "content": prompt})
|
||
|
||
templates = {
|
||
"conversation": conversation,
|
||
"tokenize": False,
|
||
"add_generation_prompt": True,
|
||
}
|
||
if args.weight == "reason":
|
||
templates["enable_thinking"] = True # 仅Reason模型使用
|
||
inputs = (
|
||
tokenizer.apply_chat_template(**templates)
|
||
if args.weight != "pretrain"
|
||
else (tokenizer.bos_token + prompt)
|
||
)
|
||
inputs = tokenizer(inputs, return_tensors="pt", truncation=True).to(args.device)
|
||
|
||
print("🤖️:", end="")
|
||
generated_ids = model.generate(
|
||
inputs=inputs["input_ids"],
|
||
attention_mask=inputs["attention_mask"],
|
||
max_new_tokens=args.max_new_tokens,
|
||
do_sample=True,
|
||
streamer=streamer,
|
||
pad_token_id=tokenizer.pad_token_id,
|
||
eos_token_id=tokenizer.eos_token_id,
|
||
top_p=args.top_p,
|
||
temperature=args.temperature,
|
||
repetition_penalty=1.0,
|
||
)
|
||
response = tokenizer.decode(
|
||
generated_ids[0][len(inputs["input_ids"][0]) :], skip_special_tokens=True
|
||
)
|
||
conversation.append({"role": "assistant", "content": response})
|
||
print("")
|