Merge b59f54cfc5 into 512eed0b65
Longfang
GitHub
commit
e38f16c912
|
|
@ -4,3 +4,4 @@ __pycache__/
|
|||
out
|
||||
website/
|
||||
docs-minimind/
|
||||
.venv/
|
||||
|
|
|
|||
|
|
@ -212,10 +212,14 @@ class MiniMindModel(nn.Module):
|
|||
past_key_values = past_key_values or [None] * len(self.layers)
|
||||
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
||||
hidden_states = self.dropout(self.embed_tokens(input_ids))
|
||||
# Recompute RoPE buffers lost during meta-device init (transformers>=5.x)
|
||||
if self.freqs_cos[0, 0] == 0:
|
||||
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
|
||||
self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
|
||||
# Lazy recompute of RoPE buffers if they were dropped by transformers>=5.x meta-device init.
|
||||
# Hot path: skipped via cheap python bool (no GPU sync). The original `freqs_cos[0,0]==0`
|
||||
# tensor compare forced a device-to-host sync on every forward.
|
||||
if getattr(self, '_rope_ready', False) is False:
|
||||
if self.freqs_cos.numel() == 0 or bool(self.freqs_cos[0, 0].eq(0).cpu()):
|
||||
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
|
||||
self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
|
||||
self._rope_ready = True
|
||||
position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
|
||||
presents = []
|
||||
for layer, past_key_value in zip(self.layers, past_key_values):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,177 @@
|
|||
"""
|
||||
W8A8 QAT fake-quantization utilities (pure PyTorch, STE).
|
||||
|
||||
- WeightFakeQuant: per-output-channel symmetric int8 on Linear weights.
|
||||
- ActFakeQuant: per-tensor symmetric int8 on activations, with EMA observer
|
||||
and a warmup window (collect stats only, no quant in forward).
|
||||
- QATLinear: drop-in nn.Linear replacement that fake-quantizes its input
|
||||
activation and its weight on every forward.
|
||||
- prepare_qat: in-place swap every nn.Linear in a model with QATLinear,
|
||||
skipping any module whose qualified name matches `skip_patterns`
|
||||
(default: lm_head, MoE router `mlp.gate`).
|
||||
|
||||
The STE used here zeros gradients for values that saturate (clip-aware STE),
|
||||
which is more stable than plain pass-through for low-bit training.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class _FakeQuantize(torch.autograd.Function):
|
||||
"""round(x/scale + zp) -> clamp -> dequant. Backward: STE inside clip, 0 outside."""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, scale, zero_point, qmin, qmax):
|
||||
q = torch.round(x / scale + zero_point)
|
||||
q_clamped = q.clamp(qmin, qmax)
|
||||
ctx.save_for_backward((q >= qmin) & (q <= qmax))
|
||||
return (q_clamped - zero_point) * scale
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
(mask,) = ctx.saved_tensors
|
||||
return grad_output * mask.to(grad_output.dtype), None, None, None, None
|
||||
|
||||
|
||||
def fake_quantize(x, scale, zero_point, qmin, qmax):
|
||||
return _FakeQuantize.apply(x, scale, zero_point, qmin, qmax)
|
||||
|
||||
|
||||
class WeightFakeQuant(nn.Module):
|
||||
"""Per-output-channel symmetric int8 fake-quant for Linear weight [out, in]."""
|
||||
|
||||
def __init__(self, bits: int = 8):
|
||||
super().__init__()
|
||||
self.bits = bits
|
||||
self.qmax = (1 << (bits - 1)) - 1
|
||||
self.qmin = -(1 << (bits - 1))
|
||||
|
||||
def forward(self, w: torch.Tensor) -> torch.Tensor:
|
||||
max_abs = w.detach().abs().amax(dim=1, keepdim=True).clamp_min(1e-8)
|
||||
scale = (max_abs / self.qmax).to(w.dtype)
|
||||
zp = torch.zeros_like(scale)
|
||||
return fake_quantize(w, scale, zp, self.qmin, self.qmax)
|
||||
|
||||
|
||||
class ActFakeQuant(nn.Module):
|
||||
"""Per-tensor symmetric int8 fake-quant with EMA min/max observer + warmup."""
|
||||
|
||||
def __init__(self, bits: int = 8, momentum: float = 0.99, observer_steps: int = 200):
|
||||
super().__init__()
|
||||
self.bits = bits
|
||||
self.qmax = (1 << (bits - 1)) - 1
|
||||
self.qmin = -(1 << (bits - 1))
|
||||
self.momentum = momentum
|
||||
self.observer_steps = observer_steps
|
||||
self.register_buffer("running_max", torch.zeros(1))
|
||||
self.register_buffer("step", torch.zeros(1, dtype=torch.long))
|
||||
self.enabled = True
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"bits={self.bits}, observer_steps={self.observer_steps}, momentum={self.momentum}"
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if not self.enabled:
|
||||
return x
|
||||
if self.training:
|
||||
cur_max = x.detach().abs().amax().to(self.running_max.dtype)
|
||||
if self.step.item() == 0:
|
||||
self.running_max.copy_(cur_max.unsqueeze(0))
|
||||
else:
|
||||
self.running_max.mul_(self.momentum).add_(cur_max.unsqueeze(0) * (1.0 - self.momentum))
|
||||
self.step += 1
|
||||
# Warmup: collect stats only; passing x through unquantized lets the
|
||||
# observer stabilize before fake-quant starts perturbing gradients.
|
||||
if self.step.item() < self.observer_steps:
|
||||
return x
|
||||
max_abs = self.running_max.clamp_min(1e-8)
|
||||
scale = (max_abs / self.qmax).to(x.dtype)
|
||||
zp = torch.zeros_like(scale)
|
||||
return fake_quantize(x, scale, zp, self.qmin, self.qmax)
|
||||
|
||||
|
||||
class QATLinear(nn.Linear):
|
||||
"""nn.Linear with fake-quant on input activation and on weight."""
|
||||
|
||||
def __init__(self, in_features, out_features, bias=True, *, w_bits=8, a_bits=8,
|
||||
a_momentum=0.99, a_observer_steps=200):
|
||||
super().__init__(in_features, out_features, bias=bias)
|
||||
self.weight_fq = WeightFakeQuant(bits=w_bits)
|
||||
self.act_fq = ActFakeQuant(bits=a_bits, momentum=a_momentum, observer_steps=a_observer_steps)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, linear: nn.Linear, **kw) -> "QATLinear":
|
||||
qm = cls(linear.in_features, linear.out_features, bias=linear.bias is not None, **kw)
|
||||
# Share existing tensors so optimizer state stays attached after swap.
|
||||
qm.weight = linear.weight
|
||||
if linear.bias is not None:
|
||||
qm.bias = linear.bias
|
||||
return qm
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x_q = self.act_fq(x)
|
||||
w_q = self.weight_fq(self.weight)
|
||||
return F.linear(x_q, w_q, self.bias)
|
||||
|
||||
|
||||
DEFAULT_SKIP_PATTERNS = ("lm_head", "mlp.gate")
|
||||
|
||||
|
||||
def _matches_skip(full_name: str, pattern: str) -> bool:
|
||||
"""Pattern matches at a dot boundary, so 'mlp.gate' does NOT skip 'mlp.gate_proj'."""
|
||||
return full_name == pattern or full_name.endswith("." + pattern)
|
||||
|
||||
|
||||
def prepare_qat(model: nn.Module, w_bits: int = 8, a_bits: int = 8,
|
||||
skip_patterns=DEFAULT_SKIP_PATTERNS,
|
||||
a_momentum: float = 0.99, a_observer_steps: int = 200) -> int:
|
||||
"""Swap every nn.Linear in `model` with QATLinear in place.
|
||||
|
||||
A Linear is skipped iff its qualified name equals a skip pattern or ends
|
||||
with `'.' + pattern` (segment-aware, not substring). Defaults skip the LM
|
||||
head (precision-sensitive, often tied to embeddings) and the MoE router
|
||||
`mlp.gate` (small, routing-sensitive) without touching the dense
|
||||
FeedForward's `mlp.gate_proj`.
|
||||
|
||||
Returns the number of linears replaced.
|
||||
"""
|
||||
replaced = 0
|
||||
for parent_name, parent in model.named_modules():
|
||||
for child_name, child in list(parent.named_children()):
|
||||
if not isinstance(child, nn.Linear) or isinstance(child, QATLinear):
|
||||
continue
|
||||
full_name = f"{parent_name}.{child_name}" if parent_name else child_name
|
||||
if any(_matches_skip(full_name, p) for p in skip_patterns):
|
||||
continue
|
||||
new_mod = QATLinear.from_float(
|
||||
child, w_bits=w_bits, a_bits=a_bits,
|
||||
a_momentum=a_momentum, a_observer_steps=a_observer_steps,
|
||||
)
|
||||
# Move buffers to weight's device. dtype on buffers stays fp32 on purpose.
|
||||
new_mod.act_fq.running_max = new_mod.act_fq.running_max.to(child.weight.device)
|
||||
new_mod.act_fq.step = new_mod.act_fq.step.to(child.weight.device)
|
||||
setattr(parent, child_name, new_mod)
|
||||
replaced += 1
|
||||
return replaced
|
||||
|
||||
|
||||
def set_act_quant_enabled(model: nn.Module, enabled: bool) -> None:
|
||||
"""Toggle activation fake-quant (e.g. disable for eval-on-fp comparison)."""
|
||||
for m in model.modules():
|
||||
if isinstance(m, ActFakeQuant):
|
||||
m.enabled = enabled
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def bake_quantized_weights(model: nn.Module) -> None:
|
||||
"""Replace each QATLinear's weight with its fake-quantized value, in place.
|
||||
|
||||
After this, the model can run as plain QATLinears without further weight
|
||||
updates and the saved fp weights already reflect the int8 grid; useful for
|
||||
final eval / export.
|
||||
"""
|
||||
for m in model.modules():
|
||||
if isinstance(m, QATLinear):
|
||||
m.weight.copy_(m.weight_fq(m.weight))
|
||||
|
|
@ -0,0 +1,163 @@
|
|||
"""
|
||||
Streaming 下载 SkyPile-150B 子采样，转成 minimind PretrainDataset 兼容格式（{"text": ...}）。
|
||||
|
||||
特性：
|
||||
- Streaming：不全下，边读边写，磁盘只占目标量
|
||||
- Token 计数：精确控制下载量（按 budget 而非样本数停）
|
||||
- 断点续传：进程被杀也能从最后写到的行接着下
|
||||
- 进度日志：每 N 个样本打印进度
|
||||
- 简易过滤：太短 / 重复字符多的样本跳过
|
||||
|
||||
Usage:
|
||||
# 默认下 25B tokens（1B 模型 Chinchilla 最优）
|
||||
python download_skypile.py
|
||||
|
||||
# 改下载量
|
||||
python download_skypile.py --target_tokens 10_000_000_000
|
||||
|
||||
# 改输出路径
|
||||
python download_skypile.py --output ../dataset/skypile_25b.jsonl
|
||||
|
||||
需要 fwdproxy（HuggingFace 在 allowlist 但 OD 无直接 DNS）：
|
||||
export https_proxy=http://fwdproxy:8080
|
||||
export http_proxy=http://fwdproxy:8080
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "scripts"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
import warnings
|
||||
from datasets import load_dataset
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def count_tokens_fast(text: str) -> int:
|
||||
"""快速估算 token 数：中文按字符 ≈ token，英文按词 × 1.3。
|
||||
用于在线计数（精确 tokenize 太慢）。"""
|
||||
n_chars = len(text)
|
||||
# 中文 token ≈ 1 char；英文 1 token ≈ 4 char
|
||||
# 简单估算：取个折中
|
||||
return int(n_chars / 1.8)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Stream-download SkyPile-150B subset")
|
||||
parser.add_argument('--repo', default='Skywork/SkyPile-150B', type=str)
|
||||
parser.add_argument('--split', default='train', type=str)
|
||||
parser.add_argument('--output', default='../dataset/skypile_25b.jsonl', type=str)
|
||||
parser.add_argument('--target_tokens', type=int, default=25_000_000_000,
|
||||
help='目标 token 数（estimate），到达后停止')
|
||||
parser.add_argument('--min_length', type=int, default=200,
|
||||
help='样本最小字符数，太短跳过')
|
||||
parser.add_argument('--log_every', type=int, default=5000)
|
||||
parser.add_argument('--tokenizer_path', type=str, default='../model',
|
||||
help='用 minimind tokenizer 精确计算 token 数（每 N 个样本核对一次）')
|
||||
parser.add_argument('--exact_count_every', type=int, default=50000,
|
||||
help='每 N 个样本精确 tokenize 一次校准估算误差')
|
||||
args = parser.parse_args()
|
||||
|
||||
output_path = os.path.abspath(args.output)
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# 断点续传：count 已有文件的行数
|
||||
n_written, est_tokens_written = 0, 0
|
||||
if os.path.exists(output_path):
|
||||
print(f'[resume] found existing {output_path}, counting...')
|
||||
with open(output_path) as f:
|
||||
for line in f:
|
||||
n_written += 1
|
||||
try:
|
||||
est_tokens_written += count_tokens_fast(json.loads(line)['text'])
|
||||
except Exception:
|
||||
pass
|
||||
print(f'[resume] {n_written:,} samples already written, ~{est_tokens_written/1e9:.2f}B tokens')
|
||||
if est_tokens_written >= args.target_tokens:
|
||||
print('[resume] target already reached, exit')
|
||||
return
|
||||
|
||||
# 校准 tokenizer（可选；下载完后跑 eval_perplexity 时会用到）
|
||||
tokenizer = None
|
||||
if args.tokenizer_path:
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
||||
print(f'[tokenizer] loaded for calibration: vocab={tokenizer.vocab_size}')
|
||||
except Exception as e:
|
||||
print(f'[tokenizer] failed to load ({e}), will skip calibration')
|
||||
|
||||
print(f'[start] streaming {args.repo} split={args.split} → {output_path}')
|
||||
print(f'[target] {args.target_tokens/1e9:.1f}B tokens (estimated)')
|
||||
|
||||
ds = load_dataset(args.repo, split=args.split, streaming=True)
|
||||
|
||||
skipped_short = 0
|
||||
n_seen = 0
|
||||
t0 = time.time()
|
||||
calibration_ratio = 1.0 # actual_tokens / estimated_tokens，每次校准更新
|
||||
|
||||
with open(output_path, 'a') as fout:
|
||||
for sample in ds:
|
||||
n_seen += 1
|
||||
# 跳过已写过的（断点续传）
|
||||
if n_seen <= n_written:
|
||||
continue
|
||||
|
||||
text = sample.get('text', '')
|
||||
if len(text) < args.min_length:
|
||||
skipped_short += 1
|
||||
continue
|
||||
|
||||
est_tok = count_tokens_fast(text)
|
||||
actual_est = int(est_tok * calibration_ratio)
|
||||
|
||||
# 写入 minimind 格式
|
||||
fout.write(json.dumps({'text': text}, ensure_ascii=False) + '\n')
|
||||
n_written += 1
|
||||
est_tokens_written += actual_est
|
||||
|
||||
# 进度日志
|
||||
if n_written % args.log_every == 0:
|
||||
elapsed = time.time() - t0
|
||||
rate = (n_written / elapsed) if elapsed > 0 else 0
|
||||
progress = est_tokens_written / args.target_tokens * 100
|
||||
eta_h = (args.target_tokens - est_tokens_written) / max(est_tokens_written / elapsed, 1) / 3600
|
||||
print(f'[{n_written:>9,} samples] {est_tokens_written/1e9:>5.2f}B tok '
|
||||
f'({progress:>5.1f}%) | {rate:.0f} samples/s | ETA {eta_h:.1f}h | '
|
||||
f'skipped_short={skipped_short:,}')
|
||||
fout.flush()
|
||||
|
||||
# 周期性精确 tokenize 几条校准估算
|
||||
if tokenizer and n_written % args.exact_count_every == 0:
|
||||
actual = sum(len(tokenizer(t, add_special_tokens=False).input_ids)
|
||||
for t in [text])
|
||||
estimate = est_tok
|
||||
if estimate > 0:
|
||||
new_ratio = actual / estimate
|
||||
calibration_ratio = 0.9 * calibration_ratio + 0.1 * new_ratio
|
||||
print(f'[calibration] est_ratio updated → {calibration_ratio:.3f} '
|
||||
f'(this sample: est={estimate}, actual={actual})')
|
||||
|
||||
if est_tokens_written >= args.target_tokens:
|
||||
print(f'\n[done] target reached: {est_tokens_written/1e9:.2f}B est tokens')
|
||||
break
|
||||
|
||||
elapsed = time.time() - t0
|
||||
print(f'\n{"=" * 60}')
|
||||
print(f'Total: {n_written:,} samples, ~{est_tokens_written/1e9:.2f}B est tokens')
|
||||
print(f'Elapsed: {elapsed/60:.1f} min ({elapsed/3600:.1f}h)')
|
||||
print(f'Skipped (too short): {skipped_short:,}')
|
||||
print(f'Output: {output_path} ({os.path.getsize(output_path)/1e9:.2f} GB)')
|
||||
print(f'{"=" * 60}')
|
||||
if tokenizer:
|
||||
print(f'⚠️ 估算 token 数有 ~{abs(1-calibration_ratio)*100:.1f}% 误差。')
|
||||
print(f' 精确数：跑 eval_perplexity.py 时会有 actual count。')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,147 @@
|
|||
"""
|
||||
Qualitative eval of a MiniMind pretrain (or post-pretrain) checkpoint by sampling
|
||||
text continuations from a fixed prompt set spanning multiple domains.
|
||||
|
||||
For pretrain (base) models — does NOT use chat template. Use eval_llm.py for SFT models.
|
||||
|
||||
Usage:
|
||||
# Default: 64M pretrain ckpt
|
||||
python eval_completion.py
|
||||
|
||||
# 1B ckpt
|
||||
python eval_completion.py --hidden_size 2048 --num_hidden_layers 20 --weight pretrain
|
||||
|
||||
# Custom prompt
|
||||
python eval_completion.py --prompt "今天天气不错，我打算"
|
||||
|
||||
# Save outputs to markdown for design doc
|
||||
python eval_completion.py --output_md eval_results.md
|
||||
|
||||
What to look for:
|
||||
✓ Grammatically coherent Chinese, no repetition loops
|
||||
✓ Stays roughly on topic of the prefix
|
||||
✓ Variety across runs (different temperature should give different outputs)
|
||||
✗ Repetitive: "的的的的" / 同一句重复 → 训练或解码有问题
|
||||
✗ 全部生成都一模一样 → 模型坍缩 / temperature 太低
|
||||
✗ 乱码 / 非中文片段（如果是纯中文 prefix）→ tokenizer 问题
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "scripts"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import warnings
|
||||
import time
|
||||
import torch
|
||||
from transformers import AutoTokenizer, TextStreamer
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
DEFAULT_PROMPTS = [
|
||||
("叙事", "在很久很久以前，有一个小村庄，"),
|
||||
("日常", "今天天气不错，我打算"),
|
||||
("技术", "Python 是一种编程语言，它的特点是"),
|
||||
("事实", "中华人民共和国成立于"),
|
||||
("代码", "def fibonacci(n):\n if n < 2:"),
|
||||
("科学", "牛顿第一定律说的是"),
|
||||
("文学", "李白是唐代著名的"),
|
||||
("数学", "二的十次方等于"),
|
||||
("常识", "水的沸点在标准大气压下是"),
|
||||
("地理", "中国的首都是"),
|
||||
]
|
||||
|
||||
|
||||
def load_model(args):
|
||||
cfg = MiniMindConfig(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
use_moe=bool(args.use_moe),
|
||||
)
|
||||
model = MiniMindForCausalLM(cfg)
|
||||
moe_suffix = '_moe' if args.use_moe else ''
|
||||
ckp = f'{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
|
||||
state_dict = torch.load(ckp, map_location='cpu', weights_only=True)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
model = model.half().eval().to(args.device)
|
||||
return model, cfg
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="MiniMind Pretrain Completion Eval")
|
||||
parser.add_argument('--hidden_size', type=int, default=768)
|
||||
parser.add_argument('--num_hidden_layers', type=int, default=8)
|
||||
parser.add_argument('--use_moe', type=int, default=0, choices=[0, 1])
|
||||
parser.add_argument('--save_dir', type=str, default='../out')
|
||||
parser.add_argument('--weight', type=str, default='pretrain')
|
||||
parser.add_argument('--tokenizer_path', type=str, default='../model')
|
||||
parser.add_argument('--prompt', type=str, default=None,
|
||||
help='单条自定义 prompt；不指定则跑 DEFAULT_PROMPTS 全集')
|
||||
parser.add_argument('--max_new_tokens', type=int, default=80)
|
||||
parser.add_argument('--temperature', type=float, default=0.7)
|
||||
parser.add_argument('--top_p', type=float, default=0.9)
|
||||
parser.add_argument('--top_k', type=int, default=50)
|
||||
parser.add_argument('--repetition_penalty', type=float, default=1.1)
|
||||
parser.add_argument('--num_samples_per_prompt', type=int, default=1,
|
||||
help='每个 prompt 采样几次（验证多样性）')
|
||||
parser.add_argument('--output_md', type=str, default=None,
|
||||
help='输出 markdown 报告路径（可选）')
|
||||
parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
|
||||
parser.add_argument('--seed', type=int, default=42)
|
||||
args = parser.parse_args()
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
|
||||
print(f'[load] ckpt={args.save_dir}/{args.weight}_{args.hidden_size}.pth')
|
||||
model, cfg = load_model(args)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
||||
|
||||
prompts = [('custom', args.prompt)] if args.prompt else DEFAULT_PROMPTS
|
||||
|
||||
md_lines = []
|
||||
md_lines.append(f'# Completion Eval: `{args.weight}_{args.hidden_size}.pth`')
|
||||
md_lines.append(f'temperature={args.temperature}, top_p={args.top_p}, top_k={args.top_k}, '
|
||||
f'rep_penalty={args.repetition_penalty}, max_new={args.max_new_tokens}\n')
|
||||
|
||||
for label, prompt in prompts:
|
||||
print(f'\n{"=" * 60}')
|
||||
print(f'[{label}] Prompt: {prompt!r}')
|
||||
md_lines.append(f'## [{label}] `{prompt}`\n')
|
||||
for sample_i in range(args.num_samples_per_prompt):
|
||||
print(f'\n--- Sample {sample_i + 1} ---')
|
||||
# 必须加 BOS：训练时所有样本以 [BOS] 开头，eval 时不加 BOS 模型会产生 degenerate
|
||||
# 输出（"的的的"/"是是是" 重复），因为它从没见过"无 BOS 开头"的序列
|
||||
body = tokenizer(prompt, add_special_tokens=False).input_ids
|
||||
ids = torch.tensor([[tokenizer.bos_token_id] + body], device=args.device)
|
||||
t0 = time.time()
|
||||
output_ids = model.generate(
|
||||
ids,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
top_k=args.top_k,
|
||||
repetition_penalty=args.repetition_penalty,
|
||||
do_sample=True,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
)
|
||||
elapsed = time.time() - t0
|
||||
text = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True)
|
||||
continuation = text[len(prompt):]
|
||||
print(continuation)
|
||||
new_tokens = output_ids.shape[1] - ids.shape[1]
|
||||
print(f' [{new_tokens} new tokens in {elapsed:.2f}s = {new_tokens/elapsed:.1f} tok/s]')
|
||||
md_lines.append(f'**Sample {sample_i + 1}**: {continuation.strip()}\n')
|
||||
md_lines.append('')
|
||||
|
||||
if args.output_md:
|
||||
with open(args.output_md, 'w') as f:
|
||||
f.write('\n'.join(md_lines))
|
||||
print(f'\n[saved] {args.output_md}')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,147 @@
|
|||
"""
|
||||
Compute perplexity of a MiniMind checkpoint on a holdout text dataset.
|
||||
|
||||
Usage:
|
||||
# Default: 64M pretrain ckpt, hash-sampled 1% of pretrain_t2t.jsonl
|
||||
python eval_perplexity.py
|
||||
|
||||
# 1B ckpt with explicit holdout
|
||||
python eval_perplexity.py --hidden_size 2048 --num_hidden_layers 20 \\
|
||||
--weight pretrain --holdout_file ../dataset/holdout.jsonl
|
||||
|
||||
# Compare two ckpts (e.g. before/after QAT)
|
||||
python eval_perplexity.py --weight pretrain
|
||||
python eval_perplexity.py --weight qat
|
||||
|
||||
Reference PPL ranges (rough):
|
||||
64M from-scratch on Chinese pretrain holdout: 8-15
|
||||
200M: 6-10
|
||||
1B: 5-8
|
||||
Anything > 30 → training is broken (data leak, lr too high, etc.)
|
||||
Anything < 5 → suspect data leak (test set seen during training)
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "scripts"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import warnings
|
||||
import time
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def load_model(args):
|
||||
cfg = MiniMindConfig(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
use_moe=bool(args.use_moe),
|
||||
)
|
||||
model = MiniMindForCausalLM(cfg)
|
||||
moe_suffix = '_moe' if args.use_moe else ''
|
||||
ckp = f'{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
|
||||
state_dict = torch.load(ckp, map_location='cpu', weights_only=True)
|
||||
# strict=False — tied embeddings have duplicate keys; load handles via shared storage
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
model = model.half().eval().to(args.device)
|
||||
return model, cfg
|
||||
|
||||
|
||||
def iter_holdout(args):
|
||||
"""Yield text strings from holdout. Either a separate file or hash-sampled from pretrain."""
|
||||
if args.holdout_file:
|
||||
with open(args.holdout_file) as f:
|
||||
for line in f:
|
||||
d = json.loads(line)
|
||||
yield d.get('text') or d.get('content') or ''
|
||||
return
|
||||
# Hash-based holdout: take 1/sampling_mod of training file (deterministic, repeatable)
|
||||
# NOTE: only meaningful as proxy — model HAS seen these during training.
|
||||
# True holdout requires reserving samples BEFORE training (Week 1 trainer fix).
|
||||
with open(args.train_file) as f:
|
||||
for line in f:
|
||||
if hash(line) % args.sampling_mod == 0:
|
||||
d = json.loads(line)
|
||||
yield d.get('text') or d.get('content') or ''
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="MiniMind Perplexity Eval")
|
||||
parser.add_argument('--hidden_size', type=int, default=768)
|
||||
parser.add_argument('--num_hidden_layers', type=int, default=8)
|
||||
parser.add_argument('--use_moe', type=int, default=0, choices=[0, 1])
|
||||
parser.add_argument('--save_dir', type=str, default='../out')
|
||||
parser.add_argument('--weight', type=str, default='pretrain',
|
||||
help='ckpt 前缀，如 pretrain / full_sft / qat')
|
||||
parser.add_argument('--tokenizer_path', type=str, default='../model')
|
||||
parser.add_argument('--holdout_file', type=str, default=None,
|
||||
help='独立 holdout jsonl 文件；若为 None 则从 train_file 按 hash 抽样')
|
||||
parser.add_argument('--train_file', type=str, default='../dataset/pretrain_t2t.jsonl',
|
||||
help='没 holdout 时从这里 hash 抽样作 proxy holdout')
|
||||
parser.add_argument('--sampling_mod', type=int, default=1000,
|
||||
help='hash 抽样模数：1000=取 0.1%, 100=取 1%')
|
||||
parser.add_argument('--max_samples', type=int, default=500)
|
||||
parser.add_argument('--max_seq_len', type=int, default=512)
|
||||
parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f'[load] ckpt={args.save_dir}/{args.weight}_{args.hidden_size}.pth')
|
||||
model, cfg = load_model(args)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
||||
pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
|
||||
|
||||
total_loss, total_tokens, n_samples = 0.0, 0, 0
|
||||
t0 = time.time()
|
||||
# ⚠️ 必须加 BOS/EOS：minimind PretrainDataset 训练时格式是 [BOS] + tokens + [EOS] + pad。
|
||||
# eval 时不加 BOS，第一个 token 的 loss 会异常高（模型从来没见过"无 BOS 开头"），
|
||||
# 整体 PPL 会虚高 ~4x（实测从 4.7 → 20.2）。
|
||||
bos_id = tokenizer.bos_token_id
|
||||
eos_id = tokenizer.eos_token_id
|
||||
for text in iter_holdout(args):
|
||||
if not text or len(text) < 10:
|
||||
continue
|
||||
body = tokenizer(text, max_length=args.max_seq_len - 2,
|
||||
truncation=True, add_special_tokens=False).input_ids
|
||||
tokens = [bos_id] + body + [eos_id]
|
||||
if len(tokens) < 2:
|
||||
continue
|
||||
ids = torch.tensor([tokens], device=args.device)
|
||||
labels = ids.clone()
|
||||
labels[labels == pad_id] = -100
|
||||
out = model(input_ids=ids, labels=labels)
|
||||
n_tokens = (labels != -100).sum().item() - 1 # 减 1：last token 不算 loss
|
||||
if n_tokens <= 0:
|
||||
continue
|
||||
total_loss += out.loss.item() * n_tokens
|
||||
total_tokens += n_tokens
|
||||
n_samples += 1
|
||||
if n_samples >= args.max_samples:
|
||||
break
|
||||
|
||||
avg_loss = total_loss / max(total_tokens, 1)
|
||||
ppl = math.exp(avg_loss)
|
||||
elapsed = time.time() - t0
|
||||
print(f'\n{"=" * 60}')
|
||||
print(f'Samples evaluated: {n_samples}')
|
||||
print(f'Total tokens: {total_tokens:,}')
|
||||
print(f'Avg loss: {avg_loss:.4f}')
|
||||
print(f'Perplexity (PPL): {ppl:.2f}')
|
||||
print(f'Elapsed: {elapsed:.1f}s ({total_tokens/elapsed:.0f} tokens/s)')
|
||||
print(f'{"=" * 60}')
|
||||
|
||||
if args.holdout_file is None:
|
||||
print(f'⚠️ Used hash-sampled "holdout" from train_file (model HAS seen these).')
|
||||
print(f' For real eval, reserve a holdout BEFORE training.')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -16,17 +16,23 @@ from torch.nn.parallel import DistributedDataParallel
|
|||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (
|
||||
get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode,
|
||||
setup_seed, init_model, SkipBatchSampler,
|
||||
detect_gpu_peak_tflops_bf16, compute_model_flops_per_token,
|
||||
)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
||||
start_time = time.time()
|
||||
last_log_time, last_log_step = start_time, start_step
|
||||
last_step = start_step
|
||||
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
|
||||
input_ids = input_ids.to(args.device)
|
||||
labels = labels.to(args.device)
|
||||
# non_blocking 仅在 --enable_optimizations=1 + pin_memory=True 时有效
|
||||
input_ids = input_ids.to(args.device, non_blocking=args.enable_optimizations == 1)
|
||||
labels = labels.to(args.device, non_blocking=args.enable_optimizations == 1)
|
||||
last_step = step
|
||||
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
|
||||
for param_group in optimizer.param_groups:
|
||||
|
|
@ -49,14 +55,21 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
|||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0 or step == iters:
|
||||
spend_time = time.time() - start_time
|
||||
now = time.time()
|
||||
spend_time = now - start_time
|
||||
window_dt = max(now - last_log_time, 1e-6)
|
||||
window_steps = step - last_log_step
|
||||
window_its = window_steps / window_dt
|
||||
tokens_window = window_steps * args.batch_size * args.max_seq_len * world_size
|
||||
mfu = (flops_per_token * tokens_window / window_dt) / (gpu_peak_tflops * 1e12) * 100
|
||||
last_log_time, last_log_step = now, step
|
||||
current_loss = loss.item() * args.accumulation_steps
|
||||
current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
|
||||
current_logits_loss = current_loss - current_aux_loss
|
||||
current_lr = optimizer.param_groups[-1]['lr']
|
||||
eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
|
||||
Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
|
||||
if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
|
||||
Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, MFU: {mfu:.1f}%, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
|
||||
if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min, "it/s": window_its, "MFU": mfu})
|
||||
|
||||
if (step % args.save_interval == 0 or step == iters) and is_main_process():
|
||||
model.eval()
|
||||
|
|
@ -104,7 +117,13 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
|
||||
parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT", help="wandb项目名")
|
||||
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
|
||||
parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1],
|
||||
help="是否使用 torch.compile 加速（0=否，1=是，A100 上推荐 1）")
|
||||
parser.add_argument("--enable_optimizations", default=1, type=int, choices=[0, 1],
|
||||
help="开关：A/B 对比优化效果。1=启用 fused AdamW + non_blocking + persistent_workers（默认）；"
|
||||
"0=baseline（关掉这些，仅保留 log 测量）。compile 由 --use_compile 独立控制。")
|
||||
parser.add_argument("--gpu_peak_tflops", type=float, default=0,
|
||||
help="GPU bf16 理论峰值 TFLOPS（用于 MFU 计算）；默认 0 表示按 GPU 名自动检测")
|
||||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
|
|
@ -136,7 +155,19 @@ if __name__ == "__main__":
|
|||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
# fused=True 把所有 param 的 Adam 更新融成一个 CUDA kernel
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate,
|
||||
fused=(args.enable_optimizations == 1))
|
||||
|
||||
# ========== 5b. MFU 测量准备（始终启用，无 perf 影响）==========
|
||||
flops_per_token = compute_model_flops_per_token(model, lm_config)
|
||||
gpu_peak_tflops = detect_gpu_peak_tflops_bf16(override=args.gpu_peak_tflops)
|
||||
world_size = dist.get_world_size() if dist.is_initialized() else 1
|
||||
Logger(f'[MFU] params_flops/token={flops_per_token/1e9:.3f} GFLOPS, '
|
||||
f'GPU peak={gpu_peak_tflops:.0f} TFLOPS bf16, world_size={world_size}')
|
||||
Logger(f'[OPT] enable_optimizations={args.enable_optimizations} '
|
||||
f'(non_blocking={args.enable_optimizations}, fused_adam={args.enable_optimizations}, '
|
||||
f'persistent_workers={args.enable_optimizations}), use_compile={args.use_compile}')
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
start_epoch, start_step = 0, 0
|
||||
|
|
@ -149,8 +180,10 @@ if __name__ == "__main__":
|
|||
|
||||
# ========== 7. 编译和分布式包装 ==========
|
||||
if args.use_compile == 1:
|
||||
model = torch.compile(model)
|
||||
Logger('torch.compile enabled')
|
||||
# enable_optimizations=1 时用 max-autotune（找更优 kernel），0 时用 default mode
|
||||
compile_mode = "max-autotune" if args.enable_optimizations == 1 else "default"
|
||||
model = torch.compile(model, mode=compile_mode)
|
||||
Logger(f'torch.compile enabled (mode={compile_mode})')
|
||||
if dist.is_initialized():
|
||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||
|
||||
|
|
@ -160,7 +193,9 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); 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)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=True,
|
||||
persistent_workers=(args.enable_optimizations == 1 and args.num_workers > 0))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
|
||||
|
|
|
|||
|
|
@ -16,17 +16,23 @@ from torch.nn.parallel import DistributedDataParallel
|
|||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import PretrainDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (
|
||||
get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode,
|
||||
setup_seed, init_model, SkipBatchSampler,
|
||||
detect_gpu_peak_tflops_bf16, compute_model_flops_per_token,
|
||||
)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
||||
start_time = time.time()
|
||||
last_log_time, last_log_step = start_time, start_step
|
||||
last_step = start_step
|
||||
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
|
||||
input_ids = input_ids.to(args.device)
|
||||
labels = labels.to(args.device)
|
||||
# non_blocking=True overlaps host->device copy with compute (pin_memory=True 已在 DataLoader)
|
||||
input_ids = input_ids.to(args.device, non_blocking=True)
|
||||
labels = labels.to(args.device, non_blocking=True)
|
||||
last_step = step
|
||||
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
|
||||
for param_group in optimizer.param_groups:
|
||||
|
|
@ -49,13 +55,22 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
|||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0 or step == iters:
|
||||
spend_time = time.time() - start_time
|
||||
now = time.time()
|
||||
spend_time = now - start_time
|
||||
window_dt = max(now - last_log_time, 1e-6)
|
||||
window_steps = step - last_log_step
|
||||
window_its = window_steps / window_dt
|
||||
# MFU = 实际 FLOPs/s / GPU 峰值；按 window 内的 token 吞吐算（避开 warmup / ckpt save 污染）
|
||||
# tokens 用 per-rank batch × seq_len × world_size × window_steps（DDP 每卡同时处理 batch）
|
||||
tokens_window = window_steps * args.batch_size * args.max_seq_len * world_size
|
||||
mfu = (flops_per_token * tokens_window / window_dt) / (gpu_peak_tflops * 1e12) * 100
|
||||
last_log_time, last_log_step = now, step
|
||||
current_loss = loss.item() * args.accumulation_steps
|
||||
current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
|
||||
current_logits_loss = current_loss - current_aux_loss
|
||||
current_lr = optimizer.param_groups[-1]['lr']
|
||||
eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
|
||||
Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
|
||||
Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, MFU: {mfu:.1f}%, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
|
||||
if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
|
||||
|
||||
if (step % args.save_interval == 0 or step == iters) and is_main_process():
|
||||
|
|
@ -103,7 +118,9 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
|
||||
parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名")
|
||||
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
|
||||
parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是，A100 上推荐默认 1）")
|
||||
parser.add_argument("--gpu_peak_tflops", type=float, default=0,
|
||||
help="GPU bf16 理论峰值 TFLOPS，用于 MFU 计算；默认 0 表示按 GPU 名自动检测（A100=312）")
|
||||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
|
|
@ -135,7 +152,15 @@ if __name__ == "__main__":
|
|||
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
# fused=True 把所有 param 的 Adam 更新融成一个 CUDA kernel（A100/3090 + CUDA tensor 支持）
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=True)
|
||||
|
||||
# ========== 5b. MFU 测量准备 ==========
|
||||
flops_per_token = compute_model_flops_per_token(model, lm_config)
|
||||
gpu_peak_tflops = detect_gpu_peak_tflops_bf16(override=args.gpu_peak_tflops)
|
||||
world_size = dist.get_world_size() if dist.is_initialized() else 1
|
||||
Logger(f'[MFU] params_flops/token={flops_per_token/1e9:.3f} GFLOPS, '
|
||||
f'GPU peak={gpu_peak_tflops:.0f} TFLOPS bf16, world_size={world_size}')
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
start_epoch, start_step = 0, 0
|
||||
|
|
@ -148,8 +173,10 @@ if __name__ == "__main__":
|
|||
|
||||
# ========== 7. 编译和分布式包装 ==========
|
||||
if args.use_compile == 1:
|
||||
model = torch.compile(model)
|
||||
Logger('torch.compile enabled')
|
||||
# max-autotune 在 Triton 层为每个 fused kernel 试多套 tile/block 配置，挑最快的。
|
||||
# 编译时间多 ~30-60s，稳态可比 default 快 5-15%。
|
||||
model = torch.compile(model, mode="max-autotune")
|
||||
Logger('torch.compile enabled (mode=max-autotune)')
|
||||
if dist.is_initialized():
|
||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||
|
||||
|
|
@ -159,7 +186,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); 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)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=True, persistent_workers=args.num_workers > 0)
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,221 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
|
||||
import argparse
|
||||
import time
|
||||
import warnings
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from model.quant import prepare_qat, DEFAULT_SKIP_PATTERNS
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from trainer.trainer_utils import (
|
||||
get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model, SkipBatchSampler,
|
||||
)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
||||
start_time = time.time()
|
||||
last_step = start_step
|
||||
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
|
||||
input_ids = input_ids.to(args.device)
|
||||
labels = labels.to(args.device)
|
||||
last_step = step
|
||||
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
with autocast_ctx:
|
||||
res = model(input_ids, labels=labels)
|
||||
loss = res.loss + res.aux_loss
|
||||
loss = loss / args.accumulation_steps
|
||||
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
if step % args.accumulation_steps == 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0 or step == iters:
|
||||
spend_time = time.time() - start_time
|
||||
current_loss = loss.item() * args.accumulation_steps
|
||||
current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
|
||||
current_logits_loss = current_loss - current_aux_loss
|
||||
current_lr = optimizer.param_groups[-1]['lr']
|
||||
eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
|
||||
Logger(
|
||||
f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
|
||||
f'loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, '
|
||||
f'aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, '
|
||||
f'epoch_time: {eta_min:.1f}min'
|
||||
)
|
||||
if wandb:
|
||||
wandb.log({
|
||||
"loss": current_loss, "logits_loss": current_logits_loss,
|
||||
"aux_loss": current_aux_loss, "learning_rate": current_lr,
|
||||
"epoch_time": eta_min,
|
||||
})
|
||||
|
||||
if (step % args.save_interval == 0 or step == iters) and is_main_process():
|
||||
model.eval()
|
||||
moe_suffix = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
|
||||
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
|
||||
raw_model = getattr(raw_model, '_orig_mod', raw_model)
|
||||
state_dict = raw_model.state_dict()
|
||||
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
|
||||
lm_checkpoint(
|
||||
lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler,
|
||||
)
|
||||
model.train()
|
||||
del state_dict
|
||||
|
||||
del input_ids, labels, res, loss
|
||||
|
||||
if last_step > start_step and last_step % args.accumulation_steps != 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind QAT (W8A8 fake-quant fine-tune)")
|
||||
parser.add_argument("--save_dir", type=str, default="../out")
|
||||
parser.add_argument('--save_weight', default='qat', type=str)
|
||||
parser.add_argument("--epochs", type=int, default=1, help="QAT 通常 1~2 epoch 即可")
|
||||
parser.add_argument("--batch_size", type=int, default=16)
|
||||
parser.add_argument("--learning_rate", type=float, default=2e-6,
|
||||
help="QAT 用比 SFT 更小的 lr（默认 1/5），防止打乱已收敛的权重分布")
|
||||
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("--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('--max_seq_len', default=768, type=int)
|
||||
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl")
|
||||
parser.add_argument('--from_weight', default='full_sft', type=str,
|
||||
help="QAT 起点权重前缀（默认基于 full_sft 微调）")
|
||||
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1])
|
||||
parser.add_argument("--use_wandb", action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-QAT")
|
||||
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1])
|
||||
|
||||
# ---- QAT-specific ----
|
||||
parser.add_argument("--w_bits", type=int, default=8, help="weight 量化位宽")
|
||||
parser.add_argument("--a_bits", type=int, default=8, help="activation 量化位宽")
|
||||
parser.add_argument("--quant_observer_steps", type=int, default=200,
|
||||
help="activation observer 预热步数（仅采集 EMA，不做 fake-quant）")
|
||||
parser.add_argument("--quant_skip", type=str, default=",".join(DEFAULT_SKIP_PATTERNS),
|
||||
help="名称中包含这些子串的 Linear 跳过量化（逗号分隔）"
|
||||
"；默认跳过 lm_head 和 MoE 路由 mlp.gate")
|
||||
args = parser.parse_args()
|
||||
|
||||
# 1. dist + seed
|
||||
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. config + checkpoint probe
|
||||
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=bool(args.use_moe),
|
||||
)
|
||||
ckp_data = lm_checkpoint(
|
||||
lm_config, weight=args.save_weight, save_dir='../checkpoints',
|
||||
) if args.from_resume == 1 else None
|
||||
|
||||
# 3. autocast
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||
|
||||
# 4. wandb
|
||||
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
|
||||
resume = 'must' if wandb_id else None
|
||||
wandb_run_name = (
|
||||
f"MiniMind-QAT-W{args.w_bits}A{args.a_bits}-Epoch-{args.epochs}-"
|
||||
f"BS-{args.batch_size}-LR-{args.learning_rate}"
|
||||
)
|
||||
wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
|
||||
|
||||
# 5. model: load fp weights -> wrap with QATLinear -> dataset/optimizer
|
||||
if args.from_weight == 'none':
|
||||
raise SystemExit("--from_weight=none is not supported for QAT; load a pretrained/SFT ckpt.")
|
||||
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
|
||||
|
||||
skip_patterns = tuple(s.strip() for s in args.quant_skip.split(',') if s.strip())
|
||||
replaced = prepare_qat(
|
||||
model, w_bits=args.w_bits, a_bits=args.a_bits,
|
||||
skip_patterns=skip_patterns, a_observer_steps=args.quant_observer_steps,
|
||||
)
|
||||
Logger(f'QAT: wrapped {replaced} Linear modules '
|
||||
f'(W{args.w_bits}A{args.a_bits}, skip={skip_patterns}, '
|
||||
f'observer_warmup={args.quant_observer_steps} steps)')
|
||||
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# 6. resume
|
||||
start_epoch, start_step = 0, 0
|
||||
if ckp_data:
|
||||
model.load_state_dict(ckp_data['model']) # state dict already contains QAT buffers
|
||||
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)
|
||||
|
||||
# 7. compile + DDP
|
||||
if args.use_compile == 1:
|
||||
model = torch.compile(model)
|
||||
Logger('torch.compile enabled')
|
||||
if dist.is_initialized():
|
||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||
|
||||
# 8. train
|
||||
for epoch in range(start_epoch, args.epochs):
|
||||
train_sampler and train_sampler.set_epoch(epoch)
|
||||
setup_seed(42 + epoch)
|
||||
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(epoch, loader, len(loader) + skip, start_step, wandb)
|
||||
else:
|
||||
train_epoch(epoch, loader, len(loader), 0, wandb)
|
||||
|
||||
# 9. cleanup
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
|
|
@ -41,6 +41,53 @@ def get_lr(current_step, total_steps, lr):
|
|||
return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
|
||||
|
||||
|
||||
# A100/H100/3090 等常见卡的 bf16 理论峰值 TFLOPS。MoE 等价 FLOPs 是 active params 不是 total，
|
||||
# active params 算出来的 MFU 才有意义；MoE 的 active 比例由 num_experts_per_tok / num_experts 决定。
|
||||
_GPU_PEAK_BF16_TFLOPS = {
|
||||
'A100': 312, 'A100-SXM4-80GB': 312, 'A100-SXM4-40GB': 312, 'A100-PCIE-80GB': 312, 'A100-PCIE-40GB': 312,
|
||||
'H100': 989, 'H100-SXM5-80GB': 989, 'H100-PCIE-80GB': 756, 'H100 80GB HBM3': 989,
|
||||
'H200': 989,
|
||||
'RTX 3090': 71, 'GeForce RTX 3090': 71, 'RTX 3090 Ti': 80,
|
||||
'RTX 4090': 165, 'GeForce RTX 4090': 165,
|
||||
'V100': 125, 'V100-SXM2-32GB': 125,
|
||||
'L40': 181, 'L40S': 362,
|
||||
}
|
||||
|
||||
|
||||
def detect_gpu_peak_tflops_bf16(device_index: int = 0, override: float = None) -> float:
|
||||
"""Look up the bf16 TensorCore peak for the current GPU; fall back to A100 (312)."""
|
||||
if override is not None and override > 0:
|
||||
return float(override)
|
||||
if not torch.cuda.is_available():
|
||||
return 312.0
|
||||
name = torch.cuda.get_device_name(device_index)
|
||||
for key, peak in _GPU_PEAK_BF16_TFLOPS.items():
|
||||
if key in name:
|
||||
return float(peak)
|
||||
Logger(f'[MFU] Unknown GPU "{name}", falling back to A100 (312 TFLOPS). Pass --gpu_peak_tflops to override.')
|
||||
return 312.0
|
||||
|
||||
|
||||
def compute_model_flops_per_token(model, config=None) -> float:
|
||||
"""Compute the "model FLOPs" per training token (forward + backward).
|
||||
|
||||
Standard approximation for dense transformers: 6 × P, where P is the total
|
||||
trainable params count. For MoE we use active params = base + experts_per_tok × per_expert.
|
||||
"""
|
||||
raw = model.module if isinstance(model, DistributedDataParallel) else model
|
||||
raw = getattr(raw, '_orig_mod', raw)
|
||||
total = sum(p.numel() for p in raw.parameters())
|
||||
if config is not None and getattr(config, 'use_moe', False):
|
||||
n_routed = getattr(config, 'num_experts', 0)
|
||||
n_active = getattr(config, 'num_experts_per_tok', 1)
|
||||
if n_routed > 0:
|
||||
expert = sum(p.numel() for n, p in raw.named_parameters() if 'mlp.experts.0.' in n)
|
||||
base = total - expert * n_routed
|
||||
active = base + expert * n_active
|
||||
return 6.0 * active
|
||||
return 6.0 * total
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
if int(os.environ.get("RANK", -1)) == -1:
|
||||
return 0 # 非DDP模式
|
||||
|
|
@ -51,14 +98,18 @@ def init_distributed_mode():
|
|||
return local_rank
|
||||
|
||||
|
||||
def setup_seed(seed: int):
|
||||
def setup_seed(seed: int, deterministic: bool = False):
|
||||
# cudnn 决定性 + 禁 benchmark 是预训练吞吐杀手（A100 上 ~3-5x 慢）。
|
||||
# 默认走 benchmark=True 让 cudnn 选最快 kernel；要严格复现传 deterministic=True。
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = deterministic
|
||||
torch.backends.cudnn.benchmark = not deterministic
|
||||
# A100 TF32 矩阵乘加速（autocast bf16 路径之外的 fp32 matmul 才生效，但无副作用）
|
||||
torch.set_float32_matmul_precision('high')
|
||||
|
||||
def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
|
|
|||
Loading…
Reference in New Issue