github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge 1c884f47af into e369b33265

pull/240/merge
Eric 2025-04-01 13:47:59 +08:00 committed by GitHub
parent e369b33265 1c884f47af
commit ece40de732
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 7131 additions and 116 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
.gitignore vendored
View File

@ -1,3 +1,9 @@
/model/__pycache__
.idea/
__pycache__
/dataset
/out
/logs
/MiniMind2

61
eval_model.py
View File

@ -13,17 +13,24 @@ warnings.filterwarnings('ignore')
def init_model(args):
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
if args.load == 0:
tokenizer = AutoTokenizer.from_pretrained(f'./model/{args.vocab}_tokenizer')
moe_path = '_moe' if args.use_moe else ''
modes = {0: 'pretrain', 1: 'full_sft', 2: 'rlhf', 3: 'reason'}
ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}.pth'
suffix = f'_{args.model_suffix}' if args.model_suffix else ''
ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}{suffix}.pth'
model = MiniMindLM(LMConfig(
dim=args.dim,
n_layers=args.n_layers,
n_heads=args.n_heads,
n_kv_heads=args.n_kv_heads,
max_seq_len=args.max_seq_len,
use_moe=args.use_moe
use_moe=args.use_moe,
vocab_size=tokenizer.vocab_size,
use_mla=args.use_mla,
use_cache=True,
torch_dtype=args.dtype,
))
state_dict = torch.load(ckp, map_location=args.device)
@ -105,8 +112,11 @@ def main():
parser = argparse.ArgumentParser(description="Chat with MiniMind")
parser.add_argument('--lora_name', default='None', type=str)
parser.add_argument('--out_dir', default='out', type=str)
parser.add_argument('--model_suffix', default='', type=str)
parser.add_argument('--temperature', default=0.85, type=float)
parser.add_argument('--top_p', default=0.85, type=float)
parser.add_argument('--rp', default=1., type=float)
parser.add_argument('--dtype', default='bfloat16', type=str)
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
# 此处max_seq_len（最大允许输入长度）并不意味模型具有对应的长文本的性能，仅防止QA出现被截断的问题
# MiniMind2-moe (145M)：(dim=640, n_layers=8, use_moe=True)
@ -114,38 +124,50 @@ def main():
# MiniMind2 (104M)：(dim=768, n_layers=16)
parser.add_argument('--dim', default=512, type=int)
parser.add_argument('--n_layers', default=8, type=int)
parser.add_argument('--n_heads', default=8, type=int)
parser.add_argument('--n_kv_heads', default=2, type=int)
parser.add_argument('--max_seq_len', default=8192, type=int)
parser.add_argument('--use_moe', default=False, type=bool)
parser.add_argument('--vocab', default='minimind', type=str)
parser.add_argument('--use_mla', default=False, action='store_true')
# 携带历史对话上下文条数
# history_cnt需要设为偶数，即【用户问题, 模型回答】为1组；设置为0时，即当前query不携带历史上文
# 模型未经过外推微调时，在更长的上下文的chat_template时难免出现性能的明显退化，因此需要注意此处设置
parser.add_argument('--history_cnt', default=0, type=int)
parser.add_argument('--stream', default=True, type=bool)
parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
parser.add_argument('--model_mode', default=1, type=int,
parser.add_argument('--model_mode', default=0, type=int,
help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型")
parser.add_argument('--test_mode', default=0, type=int, choices=[0, 1], help="测试模式；0：自动测试，1：手动测试")
args = parser.parse_args()
if args.device == 'cpu':
# intel i5 12490f 发现 bfloat16 速度明显慢于 float32
args.dtype = 'float32'
model, tokenizer = init_model(args)
prompts = get_prompt_datas(args)
test_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
test_mode = args.test_mode
messages = []
for idx, prompt in enumerate(prompts if test_mode == 0 else iter(lambda: input('👶: '), '')):
setup_seed(random.randint(0, 2048))
# setup_seed(2025) # 如需固定每次输出则换成【固定】的随机种子
if test_mode == 0: print(f'👶: {prompt}')
messages = messages[-args.history_cnt:] if args.history_cnt else []
messages.append({"role": "user", "content": prompt})
if args.model_mode == 0:
new_prompt = tokenizer.bos_token + prompt
else:
messages = messages[-args.history_cnt:] if args.history_cnt else []
messages.append({"role": "user", "content": prompt})
new_prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)[-args.max_seq_len + 1:] if args.model_mode != 0 else (tokenizer.bos_token + prompt)
new_prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)[-args.max_seq_len + 1:]
answer = new_prompt
with torch.no_grad():
x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=args.device).unsqueeze(0)
outputs = model.generate(
@ -154,22 +176,25 @@ def main():
max_new_tokens=args.max_seq_len,
temperature=args.temperature,
top_p=args.top_p,
stream=True,
rp=args.rp,
stream=args.stream,
pad_token_id=tokenizer.pad_token_id
)
print('🤖️: ', end='')
try:
if not args.stream:
print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
else:
if args.stream:
history_idx = 0
for y in outputs:
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
lis = y[0].tolist()
answer = tokenizer.decode(lis, skip_special_tokens=True)
if (answer and answer[-1] == '<EFBFBD>') or not answer:
continue
print(answer[history_idx:], end='', flush=True)
# 此处兼容自定义的按字分词表的多余空格
print(answer[history_idx + 1 if args.vocab == 'char' and history_idx > 0 else history_idx:], end='', flush=True)
history_idx = len(answer)
else:
print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
except StopIteration:
print("No answer")
print('\n')

96
mem_utils.py Normal file
View File

@ -0,0 +1,96 @@
"""
一些显存调试的工具函数
"""
import torch
def get_model_memory_usage(model):
total_param_memory = 0
param_info = {}
for name, param in model.named_parameters():
# 计算单个参数的显存（单位：MB）
param_size = param.data.element_size() * param.data.nelement() / 1024 ** 2
total_param_memory += param_size
param_info[name] = {
"size_mb": f"{param_size:.2f} MB",
"shape": param.shape,
"dtype": param.dtype,
}
param_info["total_param_memory"] = f"{total_param_memory:.2f} MB"
return param_info
def get_optimizer_memory_usage(optimizer):
total_state_memory = 0
state_info = {}
for group in optimizer.param_groups:
for p in group["params"]:
if p in optimizer.state:
state = optimizer.state[p]
for key, value in state.items():
if isinstance(value, torch.Tensor):
# 计算优化器状态张量的显存
state_size = value.element_size() * value.nelement() / 1024 ** 2
total_state_memory += state_size
state_info[f"Param_{id(p)}_{key}"] = {
"size_mb": f"{state_size:.2f} MB",
"shape": value.shape,
"dtype": value.dtype,
}
state_info["total_optimizer_state_memory"] = f"{total_state_memory:.2f} MB"
return state_info
def log_memory_snapshot(model, optimizer, step):
print(f"--- Memory Snapshot (Step {step}) ---")
# 模型参数显存
model_mem = get_model_memory_usage(model)
print("Model Parameters:")
for name, info in model_mem.items():
print(f" {name}: {info}")
# 优化器状态显存
optimizer_mem = get_optimizer_memory_usage(optimizer)
print("\nOptimizer State:")
for name, info in optimizer_mem.items():
print(f" {name}: {info}")
# 总显存占用（模型 + 优化器）
total_memory = float(model_mem["total_param_memory"].split()[0]) + \
float(optimizer_mem["total_optimizer_state_memory"].split()[0])
print(f"\nTotal Memory: {total_memory:.2f} MB")
class ActivationMemoryTracker:
def __init__(self):
self.activation_info = {}
self.hooks = []
def track_activations(self, model):
for name, module in model.named_modules():
hook = module.register_forward_hook(self._hook_fn(name))
self.hooks.append(hook)
def _hook_fn(self, name):
def hook(_, __, output):
# 计算激活值的显存（单位：MB）
if isinstance(output, torch.Tensor):
mem = output.element_size() * output.nelement() / 1024 ** 2
self.activation_info[name] = {
"size_mb": f"{mem:.2f} MB",
"size": mem,
"shape": output.shape,
"dtype": output.dtype,
}
return hook
def clear(self):
for hook in self.hooks:
hook.remove()
self.activation_info = {}
self.hooks = []

21
model/LMConfig.py
View File

@ -19,6 +19,14 @@ class LMConfig(PretrainedConfig):
rope_theta: int = 1e6,
dropout: float = 0.0,
flash_attn: bool = True,
# MLA相关配置
use_mla: bool = False,
q_lora_rank: int = 720, # query压缩的维度
kv_lora_rank: int = 14,
qk_nope_head_dim: int = 38,
qk_rope_head_dim: int = 32,
v_head_dim: int = 64,
original_seq_len = 512,
####################################################
# Here are the specific configurations of MOE
# When use_moe is false, the following is invalid
@ -32,6 +40,7 @@ class LMConfig(PretrainedConfig):
aux_loss_alpha: float = 0.1,
seq_aux: bool = True,
norm_topk_prob: bool = True,
use_cache: bool = False,
**kwargs,
):
self.dim = dim
@ -46,6 +55,17 @@ class LMConfig(PretrainedConfig):
self.rope_theta = rope_theta
self.dropout = dropout
self.flash_attn = flash_attn
self.use_mla = use_mla
self.q_lora_rank = q_lora_rank
self.kv_lora_rank = kv_lora_rank
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
self.v_head_dim = v_head_dim
self.original_seq_len = original_seq_len
self.mscale: float = 1.
self.rope_factor = 20
self.max_batch_size = 1
####################################################
# Here are the specific configurations of MOE
# When use_moe is false, the following is invalid
@ -58,4 +78,5 @@ class LMConfig(PretrainedConfig):
self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数
self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失
self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率
self.use_cache = use_cache
super().__init__(**kwargs)

191
model/attention.py Normal file
View File

@ -0,0 +1,191 @@
import math
from typing import Optional, Literal, Tuple
import torch
import torch.nn.functional as F
from torch import nn
from .LMConfig import LMConfig
block_size = 128
attn_impl: Literal["naive", "absorb"] = "absorb"
class Linear(nn.Module):
"""
这是deepseek的线性层实现，暂时没有使用
Custom linear layer with support for quantized weights and optional bias.
Args:
in_features (int): Number of input features.
out_features (int): Number of output features.
bias (bool): Whether to include a bias term. Defaults to False.
dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
"""
dtype = torch.bfloat16
def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or Linear.dtype))
nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='linear')
self.register_parameter("scale", None)
if bias:
self.bias = nn.Parameter(torch.empty(out_features))
else:
self.register_parameter("bias", None)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass for the custom linear layer.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Transformed tensor after linear computation.
"""
return F.linear(x, self.weight, self.bias)
def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
"""
Applies rotary positional embeddings to the input tensor.
Args:
x (torch.Tensor): Input tensor with positional embeddings to be applied.
freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings.
Returns:
torch.Tensor: Tensor with rotary embeddings applied.
"""
dtype = x.dtype
x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
y = torch.view_as_real(x * freqs_cis).flatten(3)
return y.to(dtype)
class MLA(nn.Module):
"""
Multi-Headed Attention Layer (MLA).
Attributes:
dim (int): Dimensionality of the input features.
n_heads (int): Number of attention heads.
n_local_heads (int): Number of local attention heads for distributed systems.
q_lora_rank (int): Rank for low-rank query projection.
kv_lora_rank (int): Rank for low-rank key/value projection.
qk_nope_head_dim (int): Dimensionality of non-positional query/key projections.
qk_rope_head_dim (int): Dimensionality of rotary-positional query/key projections.
qk_head_dim (int): Total dimensionality of query/key projections.
v_head_dim (int): Dimensionality of value projections.
softmax_scale (float): Scaling factor for softmax in attention computation.
"""
def __init__(self, args: LMConfig):
super().__init__()
self.dim = args.dim
self.n_heads = args.n_heads
self.n_local_heads = args.n_heads
self.q_lora_rank = args.q_lora_rank
self.kv_lora_rank = args.kv_lora_rank
self.qk_nope_head_dim = args.qk_nope_head_dim
self.qk_rope_head_dim = args.qk_rope_head_dim
self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
self.v_head_dim = args.v_head_dim
dtype = args.torch_dtype
if self.q_lora_rank == 0:
self.wq = nn.Linear(self.dim, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype)
else:
self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=dtype)
self.q_norm = nn.RMSNorm(self.q_lora_rank, dtype=dtype)
self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype)
self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False, dtype=dtype)
self.kv_norm = nn.RMSNorm(self.kv_lora_rank, dtype=dtype)
self.wkv_b = nn.Linear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, dtype=dtype)
self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False, dtype=dtype)
self.softmax_scale = self.qk_head_dim ** -0.5
if args.max_seq_len > args.original_seq_len:
mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
self.softmax_scale = self.softmax_scale * mscale * mscale
self.use_cache = args.use_cache
if self.use_cache:
if attn_impl == "naive":
self.register_buffer("k_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.qk_head_dim, dtype=dtype), persistent=False)
self.register_buffer("v_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.v_head_dim, dtype=dtype), persistent=False)
else:
self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank, dtype=dtype), persistent=False)
self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim, dtype=dtype), persistent=False)
def forward(self,
x: torch.Tensor,
start_pos: int,
freqs_cis: torch.Tensor,
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
mask: Optional[torch.Tensor] = None):
"""
Forward pass for the Multi-Headed Attention Layer (MLA).
Args:
x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).
start_pos (int): Starting position in the sequence for caching.
freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.
Returns:
torch.Tensor: Output tensor with the same shape as the input.
"""
bsz, seqlen, _ = x.size()
end_pos = start_pos + seqlen
if self.q_lora_rank == 0:
q = self.wq(x)
else:
q = self.wq_b(self.q_norm(self.wq_a(x)))
q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
q_pe = apply_rotary_emb(q_pe, freqs_cis)
kv = self.wkv_a(x)
kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)
v = None
if attn_impl == "naive":
q = torch.cat([q_nope, q_pe], dim=-1)
kv = self.wkv_b(self.kv_norm(kv))
kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)
k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)
if self.use_cache:
self.k_cache[:bsz, start_pos:end_pos] = k
self.v_cache[:bsz, start_pos:end_pos] = v
k = self.k_cache[:bsz, :end_pos]
v = self.k_cache[:bsz, :end_pos]
scores = torch.einsum("bshd,bthd->bsht", q, k) * self.softmax_scale
else:
# wkv_b = self.wkv_b.weight if self.wkv_b.scale is None else weight_dequant(self.wkv_b.weight, self.wkv_b.scale, block_size)
wkv_b = self.wkv_b.weight
wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
kv = self.kv_norm(kv)
k_pe = k_pe.squeeze(2)
if self.use_cache:
self.kv_cache[:bsz, start_pos:end_pos] = kv
self.pe_cache[:bsz, start_pos:end_pos] = k_pe
kv = self.kv_cache[:bsz, :end_pos]
k_pe = self.pe_cache[:bsz, :end_pos]
scores = (torch.einsum("bshc,btc->bsht", q_nope, kv) +
torch.einsum("bshr,btr->bsht", q_pe, k_pe)) * self.softmax_scale
if mask is not None:
scores += mask.unsqueeze(1)
scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x)
if attn_impl == "naive":
x = torch.einsum("bsht,bthd->bshd", scores, v)
else:
x = torch.einsum("bsht,btc->bshc", scores, kv)
x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])
x = self.wo(x.flatten(2))
return x, None

6
model/char_tokenizer/special_tokens_map.json Normal file
View File

@ -0,0 +1,6 @@
{
"bos_token": "<s>",
"eos_token": "</s>",
"pad_token": "</s>",
"unk_token": "<unk>"
}

6451
model/char_tokenizer/tokenizer.json Normal file
View File

File diff suppressed because it is too large Load Diff

36
model/char_tokenizer/tokenizer_config.json Normal file
View File

@ -0,0 +1,36 @@
{
"added_tokens_decoder": {
"0": {
"content": "<s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"1": {
"content": "</s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"2": {
"content": "<unk>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
}
},
"bos_token": "<s>",
"clean_up_tokenization_spaces": false,
"eos_token": "</s>",
"extra_special_tokens": {},
"model_max_length": 32768,
"pad_token": "</s>",
"tokenizer_class": "PreTrainedTokenizerFast",
"unk_token": "<unk>"
}

98
model/dataset.py
View File

@ -1,31 +1,87 @@
import json
import random
import re
import subprocess
from typing import List, Dict
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torch.utils.data import Dataset
import torch
from sklearn.model_selection import train_test_split
import os
import ast
os.environ["TOKENIZERS_PARALLELISM"] = "false"
def count_lines(file_name):
out = subprocess.getoutput("wc -l %s" % file_name)
return int(out.split()[0])
def texts_to_bin(input_path, tokenizer, content_key="content"):
max_buffered_length = 1 * 1024 * 1024
output_path = f'{input_path}.bin'
with open(input_path, "r", encoding="utf-8") as reader:
with open(output_path, "wb") as writer:
buffered_ids = []
i = 0
while True:
line = reader.readline()
if not line:
break
content = json.loads(line).get(content_key, "")
if not content:
continue
# 将数据序列化为二进制格式
tokenized = tokenizer(content)
buffered_ids += tokenized["input_ids"]
if len(buffered_ids) >= max_buffered_length:
arr = np.array(buffered_ids, dtype=np.uint16)
writer.write(arr.tobytes())
buffered_ids.clear()
i += 1
print(f"write {i}m bytes") if i % 100 == 0 else None
# 处理最后一段不满max_buffer_length的token序列
if len(buffered_ids) > 0:
arr = np.array(buffered_ids, dtype=np.uint16)
writer.write(arr.tobytes())
print(f"write arr: {len(arr)}")
return output_path
class PretrainBinDataset(Dataset):
def __init__(self, data_path, max_tokens):
with open(data_path) as f:
f.seek(0, 2)
self.total_tokens = f.tell() // np.dtype("uint16").itemsize
print(f"total_tokens: {self.total_tokens}")
self.data = np.memmap(data_path, dtype=np.uint16, shape=(self.total_tokens//max_tokens, max_tokens))
def __len__(self):
return self.data.shape[0]
def __getitem__(self, index):
assert isinstance(index, int)
item = self.data[index]
input = item[:-1].astype(np.int64)
target = item[1:].astype(np.int64) # 在计算交叉熵损失时要求目标输出为长整型
return torch.from_numpy(input), torch.from_numpy(target)
class PretrainDataset(Dataset):
def __init__(self, data_path, tokenizer, max_length=512):
def __init__(self, data_paths, tokenizer, max_length=512):
super().__init__()
self.tokenizer = tokenizer
self.max_length = max_length
self.samples = self.load_data(data_path)
self.samples = self.load_data(data_paths)
def load_data(self, path):
@staticmethod
def load_data(paths) -> List[Dict]:
samples = []
with open(path, 'r', encoding='utf-8') as f:
for line_num, line in enumerate(f, 1):
data = json.loads(line.strip())
samples.append(data)
for path in paths:
with open(path, 'r', encoding='utf-8') as f:
for line in f:
data = json.loads(line.strip())
samples.append(data)
return samples
def __len__(self):
@ -35,21 +91,16 @@ class PretrainDataset(Dataset):
sample = self.samples[index]
# 构建输入文本
text = f"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}"
text = sample['text']
encoding = self.tokenizer(
text,
max_length=self.max_length,
padding='max_length',
padding=False,
truncation=True,
return_tensors='pt'
)
input_ids = encoding.input_ids.squeeze()
loss_mask = (input_ids != self.tokenizer.pad_token_id)
X = torch.tensor(input_ids[:-1], dtype=torch.long)
Y = torch.tensor(input_ids[1:], dtype=torch.long)
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
return X, Y, loss_mask
return torch.tensor(input_ids, dtype=torch.long)
class SFTDataset(Dataset):
@ -196,5 +247,6 @@ class DPODataset(Dataset):
return loss_mask
if __name__ == "__main__":
pass
if __name__ == '__main__':
c = count_lines("test")
print(c)

86
model/model.py
View File

@ -12,12 +12,14 @@ from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from .attention import MLA
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float):
def __init__(self, dim: int, eps: float, dtype: torch.dtype = torch.float32):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
self.weight = nn.Parameter(torch.ones(dim, dtype = dtype))
def forward(self, x):
return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)
@ -68,13 +70,15 @@ class Attention(nn.Module):
self.n_local_kv_heads = self.n_kv_heads
self.n_rep = self.n_local_heads // self.n_local_kv_heads
self.head_dim = args.dim // args.n_heads
self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
dtype = args.torch_dtype
self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False, dtype=dtype)
self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False, dtype=dtype)
self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False, dtype=dtype)
self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False, dtype=dtype)
self.attn_dropout = nn.Dropout(args.dropout)
self.resid_dropout = nn.Dropout(args.dropout)
self.dropout = args.dropout
self.use_cache = args.use_cache
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
@ -83,9 +87,10 @@ class Attention(nn.Module):
def forward(self,
x: torch.Tensor,
start_pos: int,
pos_cis: torch.Tensor,
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
use_cache=False):
mask: Optional[torch.Tensor] = None):
bsz, seq_len, _ = x.shape
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
@ -97,7 +102,7 @@ class Attention(nn.Module):
if past_key_value is not None:
xk = torch.cat([past_key_value[0], xk], dim=1)
xv = torch.cat([past_key_value[1], xv], dim=1)
past_kv = (xk, xv) if use_cache else None
past_kv = (xk, xv) if self.use_cache else None
xq, xk, xv = (
xq.transpose(1, 2),
@ -131,9 +136,10 @@ class FeedForward(nn.Module):
hidden_dim = 4 * config.dim
hidden_dim = int(2 * hidden_dim / 3)
config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)
dtype = config.torch_dtype
self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False, dtype=dtype)
self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False, dtype=dtype)
self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False, dtype=dtype)
self.dropout = nn.Dropout(config.dropout)
def forward(self, x):
@ -153,7 +159,7 @@ class MoEGate(nn.Module):
self.norm_topk_prob = config.norm_topk_prob
self.gating_dim = config.dim
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim), dtype=config.torch_dtype))
self.reset_parameters()
def reset_parameters(self) -> None:
@ -181,7 +187,7 @@ class MoEGate(nn.Module):
topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
if self.seq_aux:
scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device, dtype=self.config.torch_dtype)
ce.scatter_add_(1, topk_idx_for_aux_loss,
torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
seq_len * aux_topk / self.n_routed_experts)
@ -235,7 +241,7 @@ class MOEFeedForward(nn.Module):
@torch.no_grad()
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
expert_cache = torch.zeros_like(x)
expert_cache = torch.zeros_like(x, dtype=self.config.torch_dtype)
idxs = flat_expert_indices.argsort()
tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
token_idxs = idxs // self.config.num_experts_per_tok
@ -263,19 +269,20 @@ class MiniMindBlock(nn.Module):
self.n_heads = config.n_heads
self.dim = config.dim
self.head_dim = config.dim // config.n_heads
self.attention = Attention(config)
self.attention = MLA(config) if config.use_mla else Attention(config)
self.layer_id = layer_id
self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps, dtype=config.torch_dtype)
self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps, dtype=config.torch_dtype)
self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
def forward(self, x, start_pos, pos_cis, past_key_value=None, mask=None):
h_attn, past_kv = self.attention(
self.attention_norm(x),
start_pos,
pos_cis,
past_key_value=past_key_value,
use_cache=use_cache
mask=mask
)
h = x + h_attn
out = h + self.feed_forward(self.ffn_norm(h))
@ -289,32 +296,42 @@ class MiniMindLM(PreTrainedModel):
self.params = params or LMConfig()
super().__init__(self.params)
self.vocab_size, self.n_layers = params.vocab_size, params.n_layers
self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
dtype = params.torch_dtype
self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim, dtype=dtype)
self.dropout = nn.Dropout(params.dropout)
self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])
self.norm = RMSNorm(params.dim, eps=params.norm_eps)
self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
self.norm = RMSNorm(params.dim, eps=params.norm_eps, dtype=dtype)
self.output = nn.Linear(params.dim, params.vocab_size, bias=False, dtype=dtype)
self.tok_embeddings.weight = self.output.weight
self.use_cache = params.use_cache
self.register_buffer("pos_cis",
precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
precompute_pos_cis(
dim=params.qk_rope_head_dim if params.use_mla else params.dim // params.n_heads,
theta=params.rope_theta
),
persistent=False)
self.OUT = CausalLMOutputWithPast()
def forward(self,
input_ids: Optional[torch.Tensor] = None,
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
use_cache: bool = False,
**args):
seqlen = input_ids.size(1)
past_key_values = past_key_values or [None] * len(self.layers)
start_pos = args.get('start_pos', 0)
h = self.dropout(self.tok_embeddings(input_ids))
pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]
pos_cis = self.pos_cis[start_pos:start_pos + seqlen]
past_kvs = []
mask = None
if seqlen > 1:
mask = torch.full((seqlen, seqlen), float("-inf"), device=input_ids.device, dtype=self.params.torch_dtype).triu_(1)
for l, layer in enumerate(self.layers):
h, past_kv = layer(
h, pos_cis,
h,
start_pos=start_pos,
pos_cis=pos_cis,
past_key_value=past_key_values[l],
use_cache=use_cache
mask=mask
)
past_kvs.append(past_kv)
logits = self.output(self.norm(h))
@ -326,16 +343,16 @@ class MiniMindLM(PreTrainedModel):
@torch.inference_mode()
def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
stream=False, rp=1., use_cache=True, pad_token_id=0, **args):
stream=False, rp=1., pad_token_id=0, **args):
# 流式生成
if stream:
return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, **args)
# 直接生成
generated = []
for i in range(input_ids.size(0)):
non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0)
out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, **args)
tokens_list = [tokens[:, -1:] for tokens in out]
gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad
full_sequence = torch.cat([non_pad, gen], dim=-1)
@ -349,14 +366,13 @@ class MiniMindLM(PreTrainedModel):
]
return torch.cat(generated, dim=0)
def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):
def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, **args):
start, first_seq, past_kvs = input_ids.shape[1], True, None
while input_ids.shape[1] < max_new_tokens - 1:
if first_seq or not use_cache:
out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False
if first_seq or not self.use_cache:
out, first_seq = self(input_ids, past_key_values=past_kvs, **args), False
else:
out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,
start_pos=input_ids.shape[1] - 1, **args)
out = self(input_ids[:, -1:], past_key_values=past_kvs, start_pos=input_ids.shape[1] - 1, **args)
logits, past_kvs = out.logits[:, -1, :], out.past_key_values
logits[:, list(set(input_ids.tolist()[0]))] /= rp
logits /= (temperature + 1e-9)

11
requirements.txt
View File

@ -19,12 +19,15 @@ scikit_learn==1.5.1
sentence_transformers==2.3.1
simhash==2.1.2
tiktoken==0.5.1
transformers==4.48.0
transformers==4.48.3
jinja2==3.1.2
jsonlines==4.0.0
trl==0.13.0
ujson==5.1.0
ujson==5.4.0
wandb==0.18.3
streamlit==1.30.0
torch==2.2.2
torchvision==0.17.2
torch==2.5.1
requests~=2.32.3
langchain==0.3.19
opencc==1.1.9
tokenizers~=0.21.0

1
scripts/plot.py Normal file
View File

@ -0,0 +1 @@
# TODO loss图生成

13
scripts/train_tokenizer.py
View File

@ -16,6 +16,12 @@ import os
random.seed(42)
base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
print(base_dir)
tokenizer_dir = f"{base_dir}/model/my_tokenizer"
def train_tokenizer():
# 读取JSONL文件并提取文本数据
@ -25,7 +31,7 @@ def train_tokenizer():
data = json.loads(line)
yield data['text']
data_path = '../dataset/pretrain_hq.jsonl'
data_path = f'{base_dir}/dataset/pretrain_hq_s.jsonl'
# 初始化tokenizer
tokenizer = Tokenizer(models.BPE())
@ -57,10 +63,9 @@ def train_tokenizer():
assert tokenizer.token_to_id("</s>") == 2
# 保存tokenizer
tokenizer_dir = "../model/minimind_tokenizer"
os.makedirs(tokenizer_dir, exist_ok=True)
tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
tokenizer.model.save("../model/minimind_tokenizer")
tokenizer.model.save(tokenizer_dir)
# 手动创建配置文件
config = {
@ -118,7 +123,7 @@ def eval_tokenizer():
from transformers import AutoTokenizer
# 加载预训练的tokenizer
tokenizer = AutoTokenizer.from_pretrained("../model/minimind_tokenizer")
tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
messages = [
{"role": "system", "content": "你是一个优秀的聊天机器人，总是给我正确的回应！"},

166
train_pretrain.py
View File

@ -1,6 +1,9 @@
import io
import os
import sys
import platform
import argparse
import threading
import time
import math
import warnings
@ -8,13 +11,16 @@ import pandas as pd
import torch
import torch.distributed as dist
from torch import optim, nn
from torch._C._profiler import ActiveProfilerType
from torch.nn.parallel import DistributedDataParallel
from torch.nn.utils.rnn import pad_sequence
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, DistributedSampler
from contextlib import nullcontext
from transformers import AutoTokenizer
from transformers import AutoTokenizer, PreTrainedTokenizerFast
from mem_utils import log_memory_snapshot, ActivationMemoryTracker
from model.model import MiniMindLM
from model.LMConfig import LMConfig
from model.dataset import PretrainDataset
@ -27,22 +33,43 @@ def Logger(content):
print(content)
def get_dynamic_lr(cur_step, warmup_steps, decay_steps, lr):
min_lr = lr / 10
if cur_step < warmup_steps:
return lr * (cur_step / warmup_steps)
if cur_step > decay_steps:
return min_lr
step_ratio = (cur_step - warmup_steps)/(decay_steps-warmup_steps)
cos_scope = 0.5 * (1 + math.cos(math.pi * step_ratio))
return min_lr + (lr - min_lr) * cos_scope
def get_lr(current_step, total_steps, lr):
return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
def train_epoch(epoch, wandb):
# 初始化跟踪器
# tracker = ActivationMemoryTracker()
# tracker.track_activations(model)
loss_fct = nn.CrossEntropyLoss(reduction='none')
step_base = epoch * iter_per_epoch
total_steps = iter_per_epoch * args.epochs
warmup_steps = int(total_steps * 0.005)
decay_steps = int(total_steps * 0.9)
if epoch == 0:
for param_group in optimizer.param_groups: # 在此处设置学习率
param_group['lr'] = 5e-6
start_time = time.time()
for step, (X, Y, loss_mask) in enumerate(train_loader):
X = X.to(args.device)
Y = Y.to(args.device)
loss_mask = loss_mask.to(args.device)
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
with ctx:
res = model(X)
loss = loss_fct(
@ -56,6 +83,11 @@ def train_epoch(epoch, wandb):
scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0:
current_step = step_base + step # 计算当前总步数
lr = get_dynamic_lr(current_step, warmup_steps, decay_steps, args.learning_rate)
for param_group in optimizer.param_groups: # 在此处设置学习率
param_group['lr'] = lr
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
@ -67,19 +99,20 @@ def train_epoch(epoch, wandb):
if step % args.log_interval == 0:
spend_time = time.time() - start_time
Logger(
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} spend:{}min remain:{}min'.format(
epoch + 1,
args.epochs,
step,
iter_per_epoch,
loss.item() * args.accumulation_steps,
optimizer.param_groups[-1]['lr'],
spend_time // 60,
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
wandb.log({"loss": loss.item() * args.accumulation_steps,
"lr": optimizer.param_groups[-1]['lr'],
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
"remain": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
model.eval()
@ -91,15 +124,26 @@ def train_epoch(epoch, wandb):
else:
state_dict = model.state_dict()
torch.save(state_dict, ckp)
buffer = io.BytesIO()
torch.save(state_dict, buffer)
threading.Thread(target=save_checkpoint, args=(buffer, ckp)).start()
model.train()
# 内存调试
# log_memory_snapshot(model, optimizer, step)
# # 输出激活值显存统计
# print("中间激活值显存占用：")
# total = 0
# for name, info in tracker.activation_info.items():
# print(f"Layer: {name} | {info}")
# total += info["size"]
# print(f"中间激活值显存总占用: {total:.2f} MB")
# print(torch.cuda.memory_summary())
# return
def init_model(lm_config):
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
model = MiniMindLM(lm_config).to(args.device)
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
return model, tokenizer
def save_checkpoint(buffer: io.BytesIO, path):
with open(path, 'wb') as f:
f.write(buffer.getvalue())
def init_distributed_mode():
@ -116,12 +160,21 @@ def init_distributed_mode():
# torchrun --nproc_per_node 2 1-pretrain.py
if __name__ == "__main__":
# 启用动态显存段扩展
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
parser = argparse.ArgumentParser(description="MiniMind Pretraining")
parser.add_argument("--out_dir", type=str, default="out")
parser.add_argument("--save_dir", type=str, default="out")
# 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
parser.add_argument("--epochs", type=int, default=1)
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--learning_rate", type=float, default=5e-4)
parser.add_argument('--dim', default=512, type=int)
parser.add_argument('--n_layers', default=8, type=int)
parser.add_argument('--n_heads', default=8, type=int)
parser.add_argument('--n_kv_heads', default=2, type=int)
parser.add_argument('--max_seq_len', default=512, type=int)
parser.add_argument('--use_mla', default=False, action='store_true')
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
parser.add_argument("--dtype", type=str, default="bfloat16")
parser.add_argument("--use_wandb", action="store_true")
@ -131,27 +184,30 @@ if __name__ == "__main__":
parser.add_argument("--accumulation_steps", type=int, default=8)
parser.add_argument("--grad_clip", type=float, default=1.0)
parser.add_argument("--warmup_iters", type=int, default=0)
parser.add_argument("--log_interval", type=int, default=100)
parser.add_argument("--save_interval", type=int, default=100)
parser.add_argument("--log_interval", type=int, default=500)
parser.add_argument("--save_interval", type=int, default=200)
parser.add_argument('--local_rank', type=int, default=-1)
parser.add_argument('--dim', default=512, type=int)
parser.add_argument('--n_layers', default=8, type=int)
parser.add_argument('--max_seq_len', default=512, type=int)
parser.add_argument('--vocab', default='minimind', type=str)
parser.add_argument('--use_moe', default=False, type=bool)
parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl")
parser.add_argument("--data_paths", type=str, nargs="+", default=["./dataset/pretrain_hq.jsonl"])
parser.add_argument("--ckpt_path", type=str, default=None)
args = parser.parse_args()
lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
args.save_dir = os.path.join(args.out_dir)
args.log_interval = args.save_interval
os.makedirs(args.save_dir, exist_ok=True)
os.makedirs(args.out_dir, exist_ok=True)
tokens_per_iter = args.batch_size * lm_config.max_seq_len
torch.manual_seed(1337)
device_type = "cuda" if "cuda" in args.device else "cpu"
args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
if device_type == "cpu":
ctx = nullcontext()
else:
# 映射dtype字符串到torch类型
amp_dtype = torch.bfloat16 if args.dtype == 'bfloat16' else torch.float16
ctx = torch.cuda.amp.autocast(dtype=amp_dtype)
ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run?
ddp_local_rank, DEVICE = 0, "cuda:0"
@ -167,20 +223,70 @@ if __name__ == "__main__":
else:
wandb = None
model, tokenizer = init_model(lm_config)
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
lm_config = LMConfig(
dim=args.dim,
n_layers=args.n_layers,
n_heads=args.n_heads,
n_kv_heads=args.n_kv_heads,
max_seq_len=args.max_seq_len,
use_moe=args.use_moe,
use_mla=args.use_mla,
torch_dtype=args.dtype,
)
tokenizer = AutoTokenizer.from_pretrained(f'./model/{args.vocab}_tokenizer')
lm_config.vocab_size = tokenizer.vocab_size
model = MiniMindLM(lm_config)
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
if args.ckpt_path is not None:
# 接着训练
state_dict = torch.load(args.ckpt_path, map_location=args.device)
model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=True)
model.to(args.device)
train_ds = PretrainDataset(args.data_paths, tokenizer, max_length=args.max_seq_len)
def collate_fn(batch):
# 对 batch 中的每个样本进行 padding，动态padding到每个batch的最大长度
batch = pad_sequence(batch, batch_first=True, padding_value=tokenizer.pad_token_id)
# 以下是padding到最大长度，方便调试显存占用情况
# max_len = args.max_seq_len
# processed = []
# for input_ids in batch:
# # 截断超过 max_len 的部分
# if len(input_ids) > max_len:
# processed.append(input_ids[:max_len])
# else:
# # 不足则填充到 max_len
# padded = torch.full(
# (max_len,),
# tokenizer.pad_token_id,
# dtype=input_ids.dtype
# )
# padded[:len(input_ids)] = input_ids
# processed.append(padded)
# batch = torch.stack(processed) # 直接堆叠已等长的样本
# 直接通过切片生成输入和目标序列
bx = batch[:, :-1] # 输入序列（去掉最后一个token）
by = batch[:, 1:] # 目标序列（去掉第一个token）
# 生成attention mask（去掉第一个token的位置）
bl = (batch != tokenizer.pad_token_id)[:, 1:]
return bx, by, bl
train_sampler = DistributedSampler(train_ds) if ddp else None
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size,
pin_memory=True,
drop_last=False,
shuffle=False,
shuffle=True,
num_workers=args.num_workers,
sampler=train_sampler
sampler=train_sampler,
collate_fn=collate_fn,
)
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
if ddp: