feat: v1.3-lite-flash 训练代码 + 新词表(vocab_size=47725)
- 用 CharTokenizer.build_vocab 从 977K 数据生成词表 (vocab_size=47725) - train_cnb.py: batch=224 + accum=2 (等效448) + AMP + torch.compile - build_vocab.py: 词表生成脚本 - vocab_full.json: 含特殊token完整词表 - 模型 41.85M参数, L40 GPU训练
This commit is contained in:
1 parent
a0047a012d
commit
a88c47e273
9 files changed
+330
-37
No files matched your search
+31
-7
@@ -76,15 +76,26 @@ def main():
|
||||
vocab_path = os.path.join(MODEL_DIR, 'vocab.json')
|
||||
with open(vocab_path, 'r', encoding='utf-8') as f:
|
||||
vocab = json.load(f)
|
||||
print(f"词表: {len(vocab)} (BPE)")
|
||||
# vocab.json 不含特殊token(0-6),需合并进来,否则:
|
||||
# 1) tokenizer 编码 <SOS>/<SEP>/<EOS> 等标记时会逐字符错码
|
||||
# 2) vocab 的 id 范围是 7..max,超过 len(vocab) 会造成 embedding 索引越界
|
||||
special_tokens = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
|
||||
'<START>': 4, '<END>': 5, '<SEP>': 6}
|
||||
for tok, tid in special_tokens.items():
|
||||
if tok not in vocab:
|
||||
vocab[tok] = tid
|
||||
real_vocab_size = max(vocab.values()) + 1
|
||||
print(f"词表: {len(vocab)}项, vocab_size={real_vocab_size} (字符+BPE, 来自 vocab.json)")
|
||||
tokenizer_instance.char_to_id = vocab.copy()
|
||||
tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()}
|
||||
tokenizer_instance.special_tokens = special_tokens.copy()
|
||||
tokenizer_instance._sorted_vocab_cache = None
|
||||
for attr in ['_multi_trie', '_multi_keys', '_single_map']:
|
||||
if hasattr(tokenizer_instance, attr):
|
||||
delattr(tokenizer_instance, attr)
|
||||
|
||||
# ====== 模型 ======
|
||||
model = LLMModel(vocab_size=max(vocab.values()) + 1, d_model=320, n_heads=8,
|
||||
model = LLMModel(vocab_size=real_vocab_size, d_model=320, n_heads=8,
|
||||
n_layers=6, d_ff=1536, max_len=256, dropout=0.1,
|
||||
use_memory=False, use_thinking=False, use_moe=False)
|
||||
total = sum(p.numel() for p in model.parameters())
|
||||
@@ -96,7 +107,7 @@ def main():
|
||||
if ckpt_path:
|
||||
ckpt = torch.load(ckpt_path, map_location=device)
|
||||
old_cfg = ckpt.get('model_config', {})
|
||||
if (old_cfg.get('vocab_size', 0) == len(vocab)
|
||||
if (old_cfg.get('vocab_size', 0) == real_vocab_size
|
||||
and old_cfg.get('d_model', 0) == 320
|
||||
and old_cfg.get('n_layers', 0) == 6):
|
||||
model.load_state_dict(ckpt['model_state_dict'])
|
||||
@@ -109,12 +120,25 @@ def main():
|
||||
|
||||
model = model.to(device)
|
||||
|
||||
# ====== 火力全开:torch.compile (inductor融合, 无CUDA graphs) + 梯度累积 ======
|
||||
# 物理batch=192(不OOM), accum=2 → 等效batch=384
|
||||
fwd_model = model
|
||||
if device == 'cuda':
|
||||
print("🔥 启用 torch.compile (mode=default, inductor融合)...")
|
||||
fwd_model = torch.compile(model, mode='default')
|
||||
print("✓ torch.compile 已就绪(首步触发编译约30-60s)")
|
||||
|
||||
# ====== DataLoader ======
|
||||
n_workers = 16 # CPU火力全开,16核全用
|
||||
train_ds = ConversationDataset(train_set, vocab, max_length=256)
|
||||
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True)
|
||||
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
|
||||
num_workers=n_workers, pin_memory=(device == 'cuda'),
|
||||
persistent_workers=True, prefetch_factor=4, drop_last=True)
|
||||
val_ds = ConversationDataset(val_set, vocab, max_length=256)
|
||||
val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False)
|
||||
print(f"批次: 训练{len(train_loader)} | 验证{len(val_loader)}")
|
||||
val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
|
||||
num_workers=n_workers, pin_memory=(device == 'cuda'),
|
||||
persistent_workers=True, prefetch_factor=4, drop_last=True)
|
||||
print(f"批次: 训练{len(train_loader)} | 验证{len(val_loader)} | num_workers={n_workers} | prefetch=4 | drop_last=True")
|
||||
|
||||
# ====== Trainer ======
|
||||
config = {
|
||||
@@ -125,7 +149,7 @@ def main():
|
||||
'd_ff': 1536, 'max_length': 256, 'dropout': 0.1,
|
||||
'use_amp': device == 'cuda',
|
||||
}
|
||||
trainer = LLMTrainer(model, device, lr=args.lr, use_amp=config['use_amp'])
|
||||
trainer = LLMTrainer(fwd_model, device, lr=args.lr, use_amp=config['use_amp'])
|
||||
|
||||
# ====== 训练 ======
|
||||
print(f"\n开始训练 epoch {start_epoch+1} → {args.epochs}")
|
||||
|
||||
Reference in new issue
Block a user