refactor(tools): 将工具脚本移动到tools目录下
This commit is contained in:
@@ -0,0 +1,198 @@
|
||||
import torch
|
||||
from typing import Dict, Any
|
||||
from dataclasses import dataclass
|
||||
from khaosz.model.transformer import ModelConfig, Transformer
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
total_tokens: int
|
||||
total_time: float
|
||||
tokens_per_second: float
|
||||
metadata: Dict[str, Any]
|
||||
|
||||
|
||||
class GenerationBenchmark:
|
||||
def __init__(
|
||||
self,
|
||||
config: ModelConfig,
|
||||
device: str = "cuda",
|
||||
dtype: torch.dtype = torch.float16
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
self.model = Transformer(config).to(device=device, dtype=dtype)
|
||||
self.model.eval()
|
||||
|
||||
def _initialize_kv_cache(self, batch_size: int) -> list:
|
||||
"""初始化KV缓存"""
|
||||
config = self.config
|
||||
shape = (batch_size, config.n_layer, config.m_len, config.n_kvhead, config.n_dim // config.n_head)
|
||||
k_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
|
||||
v_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
|
||||
return (k_cache, v_cache)
|
||||
|
||||
def _prepare_inputs(self, batch_size: int, prompt_length: int, total_length: int):
|
||||
prompt_ids = torch.randint(
|
||||
low=0,
|
||||
high=self.config.vocab_size,
|
||||
size=(batch_size, prompt_length),
|
||||
device=self.device,
|
||||
dtype=torch.long
|
||||
)
|
||||
|
||||
gen_ids = torch.randint(
|
||||
low=0,
|
||||
high=self.config.vocab_size,
|
||||
size=(batch_size, total_length - prompt_length),
|
||||
device=self.device,
|
||||
dtype=torch.long
|
||||
)
|
||||
|
||||
return prompt_ids, gen_ids
|
||||
|
||||
@torch.inference_mode()
|
||||
def run_prefill_benchmark(
|
||||
self,
|
||||
batch_size: int = 1,
|
||||
prompt_length: int = 512,
|
||||
num_trials: int = 10,
|
||||
) -> BenchmarkResult:
|
||||
|
||||
for _ in range(3):
|
||||
prompt_ids, _ = self._prepare_inputs(batch_size, prompt_length, prompt_length)
|
||||
_ = self.model(prompt_ids)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
total_time = 0.0
|
||||
total_tokens = batch_size * prompt_length * num_trials
|
||||
|
||||
for trial in range(num_trials):
|
||||
prompt_ids, _ = self._prepare_inputs(batch_size, prompt_length, prompt_length)
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
start_event.record()
|
||||
_ = self.model(prompt_ids)
|
||||
end_event.record()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
trial_time = start_event.elapsed_time(end_event) / 1000
|
||||
total_time += trial_time
|
||||
|
||||
print(f"Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
|
||||
f"({prompt_length / trial_time:.1f} tokens/s)")
|
||||
|
||||
return BenchmarkResult(
|
||||
total_tokens=total_tokens,
|
||||
total_time=total_time,
|
||||
tokens_per_second=total_tokens / total_time,
|
||||
metadata={
|
||||
"benchmark_type": "prefill",
|
||||
"batch_size": batch_size,
|
||||
"prompt_length": prompt_length,
|
||||
"dtype": self.dtype,
|
||||
"device": self.device,
|
||||
}
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def run_decoding_benchmark(
|
||||
self,
|
||||
batch_size: int = 1,
|
||||
prompt_length: int = 512,
|
||||
gen_length: int = 128,
|
||||
num_trials: int = 5,
|
||||
) -> BenchmarkResult:
|
||||
|
||||
total_time = 0.0
|
||||
total_tokens = batch_size * gen_length * num_trials
|
||||
|
||||
for trial in range(num_trials):
|
||||
|
||||
prompt_ids, gen_ids = self._prepare_inputs(batch_size, prompt_length, prompt_length + gen_length)
|
||||
kv_cache = self._initialize_kv_cache(batch_size)
|
||||
_ = self.model(prompt_ids, persistent_key_values=kv_cache, start_pos=0)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
start_event.record()
|
||||
|
||||
current_pos = prompt_length
|
||||
for i in range(gen_length):
|
||||
input_token = gen_ids[:, i:i+1]
|
||||
_ = self.model(input_token, persistent_key_values=kv_cache, start_pos=current_pos)
|
||||
current_pos += 1
|
||||
|
||||
end_event.record()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
trial_time = start_event.elapsed_time(end_event) / 1000
|
||||
total_time += trial_time
|
||||
|
||||
print(f"Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
|
||||
f"({gen_length / trial_time:.1f} tokens/s)")
|
||||
|
||||
|
||||
return BenchmarkResult(
|
||||
total_tokens=total_tokens,
|
||||
total_time=total_time,
|
||||
tokens_per_second=total_tokens / total_time,
|
||||
metadata={
|
||||
"benchmark_type": "decoding",
|
||||
"batch_size": batch_size,
|
||||
"prompt_length": prompt_length,
|
||||
"gen_length": gen_length,
|
||||
"dtype": self.dtype,
|
||||
"device": self.device,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def print_benchmark_result(result: BenchmarkResult):
|
||||
"""打印基准测试结果"""
|
||||
benchmark_type = result.metadata["benchmark_type"]
|
||||
|
||||
print(f"\n{' ' + benchmark_type.upper().replace('_', ' ') + ' Benchmark ':-^80}")
|
||||
print(f"Total Tokens Processed: {result.total_tokens:,}")
|
||||
print(f"Time Consumed: {result.total_time:.3f}s")
|
||||
print(f"Throughput: {result.tokens_per_second:,.1f} tokens/s")
|
||||
|
||||
if benchmark_type == "prefill":
|
||||
print(f"Batch Size: {result.metadata['batch_size']} | Prompt Length: {result.metadata['prompt_length']}")
|
||||
elif benchmark_type == "decoding":
|
||||
print(f"Batch Size: {result.metadata['batch_size']} | Gen Length: {result.metadata['gen_length']}")
|
||||
|
||||
print(f"Device: {result.metadata['device']} | Dtype: {result.metadata['dtype']}")
|
||||
print("-" * 80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
config = ModelConfig(
|
||||
vocab_size=10000,
|
||||
n_dim=1536,
|
||||
n_head=24,
|
||||
n_kvhead=4,
|
||||
d_ffn=6912,
|
||||
m_len=2048,
|
||||
n_layer=24,
|
||||
norm_eps=1e-5,
|
||||
)
|
||||
|
||||
benchmark = GenerationBenchmark(config)
|
||||
|
||||
print("=" * 80)
|
||||
print("Running Transformer Generation Benchmark")
|
||||
print("=" * 80)
|
||||
|
||||
prefill_result = benchmark.run_prefill_benchmark(batch_size=4, prompt_length=512, num_trials=5)
|
||||
print_benchmark_result(prefill_result)
|
||||
|
||||
gen_result = benchmark.run_decoding_benchmark(batch_size=4, prompt_length=512, gen_length=128, num_trials=5)
|
||||
print_benchmark_result(gen_result)
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import torch
|
||||
import json
|
||||
import torch
|
||||
import argparse
|
||||
|
||||
from khaosz import Khaosz
|
||||
from typing import List
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def batch_generate(
|
||||
model: Khaosz,
|
||||
queries: List[str],
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
batch_size: int,
|
||||
) -> List:
|
||||
assert batch_size > 0
|
||||
sorted_queries = sorted(queries, key=lambda x: len(x), reverse=True)
|
||||
original_indices = {query: idx for idx, query in enumerate(queries)}
|
||||
|
||||
responses = [None] * len(queries)
|
||||
total_batches = (len(sorted_queries) + batch_size - 1) // batch_size
|
||||
|
||||
for i in tqdm(range(0, total_batches * batch_size, batch_size), desc="Generating responses"):
|
||||
batch_queries = sorted_queries[i: min(i + batch_size, len(queries))]
|
||||
if not isinstance(batch_queries, list):
|
||||
batch_queries = [batch_queries]
|
||||
|
||||
batch_responses = model.batch_generate(
|
||||
queries=batch_queries,
|
||||
temperature=temperature,
|
||||
top_k=top_k,
|
||||
top_p=top_p
|
||||
)
|
||||
|
||||
for batch_query, batch_response in zip(batch_queries, batch_responses):
|
||||
print(f"Q: {batch_query[:50]} \nR: {batch_response[:50]})")
|
||||
|
||||
for query, response in zip(batch_queries, batch_responses):
|
||||
original_idx = original_indices[query]
|
||||
responses[original_idx] = response
|
||||
|
||||
return responses
|
||||
|
||||
|
||||
def processor(
|
||||
model: Khaosz,
|
||||
input_json_file: str,
|
||||
output_json_file: str,
|
||||
batch_size: int,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
top_k: int,
|
||||
question_key: str="question",
|
||||
):
|
||||
with open(input_json_file, "r", encoding='utf-8') as f:
|
||||
input_dict = [json.loads(line) for line in f]
|
||||
queries = [item[question_key] for item in input_dict]
|
||||
|
||||
output_dict = batch_generate(
|
||||
model=model,
|
||||
queries=queries,
|
||||
temperature=temperature,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
batch_size=batch_size
|
||||
)
|
||||
|
||||
with open(output_json_file, "w", encoding='utf-8') as f:
|
||||
json.dump(output_dict, f, indent=4, ensure_ascii=False)
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run generate with a Khaosz model.")
|
||||
|
||||
parser.add_argument("--model_dir", type=str, required=True, help="Path to the model directory.")
|
||||
parser.add_argument("--input_json_file", type=str, required=True, help="Path to the input JSONL file.")
|
||||
parser.add_argument("--output_json_file", type=str, required=True, help="Path to the output JSONL file.")
|
||||
parser.add_argument("--question_key", type=str, default="question", help="Key for the question in the input JSON.")
|
||||
parser.add_argument("--temperature", type=float, default=0.60, help="Temperature for generating responses.")
|
||||
parser.add_argument("--top_p", type=float, default=0.95, help="Top-p value for generating responses.")
|
||||
parser.add_argument("--top_k", type=int, default=30, help="Top-k value for generating responses.")
|
||||
parser.add_argument("--batch_size", type=int, default=1, help="Batch size for generating responses.")
|
||||
|
||||
args = parser.parse_args()
|
||||
model = Khaosz(args.model_dir).to(device='cuda', dtype=torch.bfloat16)
|
||||
|
||||
processor(
|
||||
model,
|
||||
input_json_file=args.input_json_file,
|
||||
output_json_file=args.output_json_file,
|
||||
question_key=args.question_key,
|
||||
batch_size=args.batch_size,
|
||||
temperature=args.temperature,
|
||||
top_k=args.top_k,
|
||||
top_p=args.top_p
|
||||
)
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
import os
|
||||
import argparse
|
||||
import torch
|
||||
|
||||
from torch.optim import AdamW
|
||||
from khaosz.config import ParameterLoader, Checkpoint, TrainConfig, CosineScheduleConfig
|
||||
from khaosz.trainer import Trainer, StrategyFactory
|
||||
from khaosz.data import DatasetLoader
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def get_files(root_path: str) -> list[str]:
|
||||
paths = []
|
||||
for root, _, files in os.walk(root_path):
|
||||
paths.extend([os.path.join(root, file) for file in files])
|
||||
|
||||
return paths
|
||||
|
||||
def train(
|
||||
train_type: str,
|
||||
param_path: str,
|
||||
data_root_path: str,
|
||||
max_lr: int,
|
||||
n_epoch: int,
|
||||
batch_size: int,
|
||||
start_epoch: int,
|
||||
start_batch: int,
|
||||
accumulation_steps: int,
|
||||
warmup_steps: int,
|
||||
checkpoint_interval: int,
|
||||
checkpoint_dir: str,
|
||||
dpo_beta: float,
|
||||
adamw_betas: tuple,
|
||||
adamw_weight_decay: float,
|
||||
max_grad_norm: float,
|
||||
embdeding_lr_rate: int,
|
||||
random_seed: int,
|
||||
window_size: int,
|
||||
stride: int,
|
||||
resume_from_checkpoint: bool
|
||||
):
|
||||
assert train_type in ["seq", "sft", "dpo"]
|
||||
assert os.path.exists(param_path)
|
||||
|
||||
parameter = ParameterLoader.load(param_path)
|
||||
checkpoint = None
|
||||
|
||||
if isinstance(parameter, Checkpoint) and resume_from_checkpoint:
|
||||
checkpoint = parameter
|
||||
|
||||
if window_size is None:
|
||||
window_size = parameter.config.m_len
|
||||
|
||||
model = parameter.model
|
||||
device = torch.device("cuda")
|
||||
model = model.to(device=device, dtype=torch.bfloat16)
|
||||
cache_files = get_files(data_root_path)
|
||||
|
||||
kwargs = {
|
||||
"dpo_beta": dpo_beta,
|
||||
"bos_token_id": parameter.tokenizer.bos_id,
|
||||
"eos_token_id": parameter.tokenizer.eos_id,
|
||||
"pad_token_id": parameter.tokenizer.pad_id,
|
||||
}
|
||||
|
||||
strategy = StrategyFactory.load(
|
||||
model,
|
||||
train_type,
|
||||
device,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
dataset = DatasetLoader.load(
|
||||
train_type=train_type,
|
||||
load_path=cache_files,
|
||||
window_size=window_size,
|
||||
stride=stride,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
param_groups = [
|
||||
{"params": [p for n, p in model.named_parameters() if "embedding" in n], "lr": max_lr * embdeding_lr_rate},
|
||||
{"params": [p for n, p in model.named_parameters() if "embedding" not in n], "lr": max_lr}
|
||||
]
|
||||
|
||||
optim = AdamW(
|
||||
param_groups,
|
||||
betas=adamw_betas,
|
||||
weight_decay=adamw_weight_decay
|
||||
)
|
||||
|
||||
train_config = TrainConfig(
|
||||
strategy=strategy,
|
||||
dataset=dataset,
|
||||
optimizer=optim,
|
||||
checkpoint_dir=checkpoint_dir,
|
||||
n_epoch=n_epoch,
|
||||
batch_size=batch_size,
|
||||
start_epoch=start_epoch,
|
||||
start_batch=start_batch,
|
||||
checkpoint_interval=checkpoint_interval,
|
||||
accumulation_steps=accumulation_steps,
|
||||
max_grad_norm=max_grad_norm,
|
||||
random_seed=random_seed,
|
||||
num_workers=4,
|
||||
pin_memory=True
|
||||
)
|
||||
|
||||
schedule_config = CosineScheduleConfig(
|
||||
warmup_steps=warmup_steps,
|
||||
total_steps=len(dataset) * n_epoch // batch_size,
|
||||
)
|
||||
|
||||
trainer = Trainer(
|
||||
parameter=parameter,
|
||||
train_config=train_config,
|
||||
schedule_config=schedule_config,
|
||||
)
|
||||
trainer.train(checkpoint)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Train the Transformer model.")
|
||||
# train args
|
||||
parser.add_argument("--train_type",choices=["seq", "sft", "dpo"], help="Train type.")
|
||||
parser.add_argument("--data_root_path", type=str, required=True, help="Path to the root directory of the dataset.")
|
||||
parser.add_argument("--param_path", type=str, required=True, help="Path to the model parameters or resume checkpoint.")
|
||||
parser.add_argument("--n_epoch", type=int, default=1, help="Number of epochs to train.")
|
||||
parser.add_argument("--batch_size", type=int, default=1, help="Batch size for training.")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="Number of iterations between each optimizer step.")
|
||||
parser.add_argument("--warmup_steps", type=int, default=1000, help="Number of iters between warnings.")
|
||||
parser.add_argument("--max_lr", type=float, default=3e-4, help="Max learning rate for training.")
|
||||
parser.add_argument("--checkpoint_interval", type=int, default=5000, help="Number of iters between checkpoints.")
|
||||
parser.add_argument("--checkpoint_dir", type=str, default="checkpoint", help="Directory to save checkpoints.")
|
||||
parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping.")
|
||||
parser.add_argument("--adamw_betas", type=tuple, default=(0.9, 0.95), help="Beta values for AdamW optimizer.")
|
||||
parser.add_argument("--adamw_weight_decay", type=float, default=0.01, help="Weight decay for AdamW optimizer.")
|
||||
parser.add_argument("--embdeding_lr_rate", type=float, default=1.0, help="The rate between the embedding layers lr rate and the max lr rate.")
|
||||
parser.add_argument("--random_seed", type=int, default=3407, help="Random seed for reproducibility.")
|
||||
|
||||
# other configs
|
||||
parser.add_argument("--window_size", type=int, default=None, help="the max length of the input sequence.")
|
||||
parser.add_argument("--stride", type=int, default=None, help="the step size of the input sequence.")
|
||||
parser.add_argument("--start_epoch", type=int, default=0, help="Start epoch for training.")
|
||||
parser.add_argument("--start_batch", type=int, default=0, help="Start batch for training.")
|
||||
parser.add_argument("--resume_from_checkpoint", type=bool, default=False, help="train from checkpoint or not.")
|
||||
parser.add_argument("--dpo_beta", type=float, default=0.1, help="DPO beta value.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
train(**vars(args))
|
||||
Reference in New Issue
Block a user