[2025.06.17] update train code
This commit is contained in:
root
Executable
+74
@@ -0,0 +1,74 @@
|
||||
import os
|
||||
import multiprocessing as mp
|
||||
from glob import glob
|
||||
import numpy as np
|
||||
from transformers import AutoTokenizer
|
||||
from datasets import load_dataset
|
||||
from tqdm import tqdm
|
||||
|
||||
DATA_DIR = "../corpus"
|
||||
STATS_FILE = "token_stats.json"
|
||||
NUM_PROC = min(32, os.cpu_count())
|
||||
SAMPLE_SIZE = 10000 #None
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("../weights/weights_tokenizer")
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
def get_token_lengths(texts):
|
||||
return [len(tokenizer.encode(t, add_special_tokens=False)) for t in texts]
|
||||
|
||||
def process_file(file_path):
|
||||
try:
|
||||
dataset = load_dataset("parquet", data_files=file_path, split="train")
|
||||
if SAMPLE_SIZE and len(dataset) > SAMPLE_SIZE:
|
||||
dataset = dataset.shuffle().select(range(SAMPLE_SIZE))
|
||||
texts = dataset["text"]
|
||||
lengths = get_token_lengths(texts)
|
||||
return {
|
||||
"file": os.path.basename(file_path),
|
||||
"count": len(lengths),
|
||||
"mean": np.mean(lengths),
|
||||
"std": np.std(lengths),
|
||||
"min": np.min(lengths),
|
||||
"max": np.max(lengths),
|
||||
"percentiles": {
|
||||
str(p): np.percentile(lengths, p)
|
||||
for p in [50, 75, 90, 95, 99]
|
||||
},
|
||||
"all_lengths": lengths
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"处理 {file_path} 出错: {e}")
|
||||
return None
|
||||
|
||||
if __name__ == "__main__":
|
||||
file_paths = glob(os.path.join(DATA_DIR, "*.parquet"))
|
||||
print(f"找到 {len(file_paths)} 个parquet文件")
|
||||
all_stats = []
|
||||
with mp.Pool(processes=NUM_PROC) as pool:
|
||||
results = pool.imap(process_file, file_paths)
|
||||
for stat in tqdm(results, total=len(file_paths)):
|
||||
if stat is not None:
|
||||
all_stats.append(stat)
|
||||
|
||||
all_lengths = []
|
||||
for stat in all_stats:
|
||||
all_lengths.extend(stat["all_lengths"])
|
||||
stat.pop("all_lengths")
|
||||
|
||||
global_stats = {
|
||||
"total_samples": len(all_lengths),
|
||||
"global_mean": np.mean(all_lengths),
|
||||
"global_max": np.max(all_lengths),
|
||||
"global_95_percentile": np.percentile(all_lengths, 95),
|
||||
"recommended_max_length": int(np.percentile(all_lengths, 95) * 1.05)
|
||||
}
|
||||
|
||||
print(f"\n{'='*50}")
|
||||
print("全局统计结果:")
|
||||
print(f"总样本数: {global_stats['total_samples']:,}")
|
||||
print(f"平均token长度: {global_stats['global_mean']:.1f}")
|
||||
print(f"最大token长度: {global_stats['global_max']}")
|
||||
print(f"95百分位长度: {global_stats['global_95_percentile']}")
|
||||
print(f"建议最大长度: {global_stats['recommended_max_length']}")
|
||||
@@ -3,43 +3,33 @@ import multiprocessing as mp
|
||||
from glob import glob
|
||||
import numpy as np
|
||||
from transformers import AutoTokenizer
|
||||
from datasets import load_dataset, Dataset
|
||||
from datasets import load_dataset
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import pandas as pd
|
||||
from tqdm import tqdm
|
||||
import gc
|
||||
import psutil
|
||||
import time
|
||||
|
||||
# 配置参数
|
||||
MODEL_PATH = "models/qwen2" # 本地模型路径
|
||||
DATA_DIR = "datas" # 原始数据目录
|
||||
OUTPUT_DIR = "tokenized_data" # tokenized数据输出目录
|
||||
MAX_SEQ_LENGTH = 1024 # 最大序列长度
|
||||
NUM_PROC = min(32, psutil.cpu_count(logical=False)) # 物理核心数
|
||||
CHUNK_SIZE = 50000 # 每个处理块的大小
|
||||
MAX_MEMORY_PERCENT = 70 # 最大内存使用百分比
|
||||
|
||||
# 创建输出目录
|
||||
DATA_DIR = "../corpus"
|
||||
OUTPUT_DIR = "../corpus/tokens"
|
||||
MAX_SEQ_LENGTH = 1024
|
||||
NUM_PROC = min(32, psutil.cpu_count(logical=False))
|
||||
CHUNK_SIZE = 50000
|
||||
MAX_MEMORY_PERCENT = 70
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
# 加载分词器
|
||||
tokenizer = AutoTokenizer.from_pretrained("/mnt/d/VMshare/LLM-from-scratch/weights/weights_tokenizer")
|
||||
tokenizer = AutoTokenizer.from_pretrained("../weights/weights_tokenizer")
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
# 获取所有parquet文件
|
||||
file_paths = ["/mnt/d/VMshare/LLM-from-scratch/corpus/parquet/train-00001-of-00399.parquet"]
|
||||
file_paths = glob(os.path.join(DATA_DIR, "*.parquet"))
|
||||
print(f"找到 {len(file_paths)} 个parquet文件")
|
||||
|
||||
def memory_safe():
|
||||
"""检查内存是否安全"""
|
||||
mem = psutil.virtual_memory()
|
||||
return mem.percent < MAX_MEMORY_PERCENT
|
||||
|
||||
def process_chunk(texts, tokenizer):
|
||||
"""处理文本块并返回tokenized结果"""
|
||||
tokenized = tokenizer(
|
||||
texts,
|
||||
truncation=True,
|
||||
@@ -50,31 +40,20 @@ def process_chunk(texts, tokenizer):
|
||||
return tokenized["input_ids"], tokenized["attention_mask"]
|
||||
|
||||
def process_file(file_path):
|
||||
"""处理单个大型parquet文件"""
|
||||
try:
|
||||
base_name = os.path.basename(file_path)
|
||||
output_path = os.path.join(OUTPUT_DIR, f"tokenized_{base_name}")
|
||||
|
||||
# 如果文件已处理则跳过
|
||||
if os.path.exists(output_path):
|
||||
print(f"文件已存在: {output_path}, 跳过")
|
||||
return
|
||||
|
||||
print(f"开始处理: {base_name}")
|
||||
start_time = time.time()
|
||||
|
||||
# 1. 获取文件总行数
|
||||
dataset = load_dataset("parquet", data_files=file_path, split="train")
|
||||
total_rows = len(dataset)
|
||||
print(f"文件包含 {total_rows:,} 行")
|
||||
|
||||
# 2. 创建空DataFrame用于收集结果
|
||||
all_input_ids = []
|
||||
all_attention_mask = []
|
||||
|
||||
# 3. 分块处理
|
||||
num_chunks = (total_rows + CHUNK_SIZE - 1) // CHUNK_SIZE
|
||||
|
||||
for chunk_idx in range(num_chunks):
|
||||
if not memory_safe():
|
||||
print("内存使用过高，等待释放...")
|
||||
@@ -83,58 +62,39 @@ def process_file(file_path):
|
||||
if not memory_safe():
|
||||
print("内存仍然不足，跳过当前文件")
|
||||
return
|
||||
|
||||
start_row = chunk_idx * CHUNK_SIZE
|
||||
end_row = min((chunk_idx + 1) * CHUNK_SIZE, total_rows)
|
||||
|
||||
# 读取当前块
|
||||
chunk = dataset.select(range(start_row, end_row))
|
||||
texts = chunk["text"] # 根据实际字段名调整
|
||||
|
||||
# 处理当前块
|
||||
texts = chunk["text"]
|
||||
input_ids, attention_mask = process_chunk(texts, tokenizer)
|
||||
|
||||
# 收集结果
|
||||
all_input_ids.append(input_ids)
|
||||
all_attention_mask.append(attention_mask)
|
||||
|
||||
# 每处理5个块显示进度
|
||||
if (chunk_idx + 1) % 5 == 0 or (chunk_idx + 1) == num_chunks:
|
||||
elapsed = time.time() - start_time
|
||||
rows_done = min((chunk_idx + 1) * CHUNK_SIZE, total_rows)
|
||||
percent = rows_done / total_rows * 100
|
||||
print(f"进度: {rows_done:,}/{total_rows:,} 行 ({percent:.1f}%) "
|
||||
f"耗时: {elapsed:.1f}秒")
|
||||
|
||||
# 4. 合并所有块
|
||||
all_input_ids = np.vstack(all_input_ids)
|
||||
all_attention_mask = np.vstack(all_attention_mask)
|
||||
|
||||
# 5. 创建PyArrow表并保存
|
||||
table = pa.Table.from_arrays([
|
||||
pa.array(all_input_ids.reshape(-1, MAX_SEQ_LENGTH).tolist()),
|
||||
pa.array(all_attention_mask.reshape(-1, MAX_SEQ_LENGTH).tolist())
|
||||
], names=["input_ids", "attention_mask"])
|
||||
|
||||
pq.write_table(table, output_path)
|
||||
|
||||
# 6. 清理内存
|
||||
del dataset, all_input_ids, all_attention_mask, table
|
||||
gc.collect()
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"完成处理: {base_name} -> {output_path} "
|
||||
f"总耗时: {elapsed:.1f}秒, 平均速度: {total_rows/elapsed:.1f}行/秒")
|
||||
|
||||
except Exception as e:
|
||||
print(f"处理文件 {file_path} 时出错: {str(e)}")
|
||||
# 清理可能的残留资源
|
||||
if 'dataset' in locals():
|
||||
del dataset
|
||||
gc.collect()
|
||||
|
||||
def process_file_wrapper(args):
|
||||
"""包装函数用于处理内存监控"""
|
||||
file_path, idx, total = args
|
||||
print(f"\n{'='*50}\n处理文件 [{idx+1}/{total}]: {os.path.basename(file_path)}\n{'='*50}")
|
||||
process_file(file_path)
|
||||
@@ -142,28 +102,7 @@ def process_file_wrapper(args):
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(f"开始预处理 {len(file_paths)} 个文件，使用 {NUM_PROC} 个进程")
|
||||
print(f"每个文件约18万行，总数据量约 {len(file_paths)*180000/1000000:.1f} 百万行")
|
||||
|
||||
# 准备参数
|
||||
file_args = [(fp, idx, len(file_paths)) for idx, fp in enumerate(file_paths)]
|
||||
|
||||
# 使用进程池并行处理
|
||||
with mp.Pool(processes=NUM_PROC) as pool:
|
||||
results = list(tqdm(pool.imap(process_file_wrapper, file_args), total=len(file_paths)))
|
||||
|
||||
print("所有文件处理完成!")
|
||||
|
||||
# 生成元数据文件
|
||||
metadata = {
|
||||
"model": MODEL_PATH,
|
||||
"tokenized_date": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"max_seq_length": MAX_SEQ_LENGTH,
|
||||
"num_files": len(file_paths),
|
||||
"output_dir": OUTPUT_DIR
|
||||
}
|
||||
|
||||
with open(os.path.join(OUTPUT_DIR, "metadata.json"), "w") as f:
|
||||
import json
|
||||
json.dump(metadata, f, indent=2)
|
||||
|
||||
print(f"元数据已保存到 {OUTPUT_DIR}/metadata.json")
|
||||
@@ -1,54 +0,0 @@
|
||||
import os
|
||||
import torch
|
||||
from transformers import (
|
||||
AutoTokenizer,
|
||||
AutoModelForCausalLM,
|
||||
TrainingArguments,
|
||||
Trainer,
|
||||
DataCollatorForLanguageModeling
|
||||
)
|
||||
from datasets import Dataset, load_from_disk
|
||||
import deepspeed
|
||||
import numpy as np
|
||||
from glob import glob
|
||||
import logging
|
||||
from torch.utils.data import IterableDataset, DataLoader
|
||||
import pyarrow.parquet as pq
|
||||
# 配置参数
|
||||
MODEL_PATH = "models/qwen2" # 本地模型路径
|
||||
TOKENIZED_DATA_DIR = "corpus" # tokenized数据目录
|
||||
OUTPUT_DIR = "./output" # 训练输出目录
|
||||
DEEPSPEED_CONFIG = "ds_config.json" # DeepSpeed配置文件
|
||||
NUM_GPUS = 4 # GPU数量
|
||||
PER_DEVICE_BATCH_SIZE = 4 # 每张GPU的batch size
|
||||
GRADIENT_ACCUM_STEPS = 8 # 梯度累积步数
|
||||
MAX_TRAIN_STEPS = 100000 # 最大训练步数
|
||||
SAVE_STEPS = 2000 # 保存间隔步数
|
||||
LOGGING_STEPS = 50 # 日志记录间隔
|
||||
LEARNING_RATE = 5e-5 # 学习率
|
||||
# 设置日志
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
# ===== 1. 准备模型和分词器 =====
|
||||
logger.info("加载模型和分词器...")
|
||||
from transformers import AutoTokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained("/mnt/d/VMshare/LLM-from-scratch/weights/weights_tokenizer")
|
||||
file_path="tokenized_data/tokenized_train-00001-of-00399.parquet"
|
||||
table = pq.read_table(file_path)
|
||||
input_ids = table["input_ids"].to_numpy()
|
||||
attention_mask = table["attention_mask"].to_numpy()
|
||||
print(len(input_ids))
|
||||
for input_id in input_ids:
|
||||
print(input_id)
|
||||
print(len(input_id))
|
||||
res=tokenizer.decode(input_id)
|
||||
# print(res)
|
||||
break
|
||||
|
||||
file_path2="/mnt/d/VMshare/LLM-from-scratch/corpus/parquet/train-00001-of-00399.parquet"
|
||||
table2 = pq.read_table(file_path2)
|
||||
text = table2["text"]
|
||||
print(len(text))
|
||||
for value in text:
|
||||
print(value)
|
||||
break
|
||||
Reference in New Issue
Block a user