ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

style: 使用ruff 工具优化代码风格

Browse Source
This commit is contained in:
ViperEkura
2026-03-30 23:32:28 +08:00
parent 345fd2f091
commit 426af2d75f
52 changed files with 1836 additions and 1493 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/benchmark.py
+74 -51
View File
@@ -17,41 +17,47 @@ class GenerationBenchmark:
self,
config: ModelConfig,
device: str = "cuda",
dtype: torch.dtype = torch.float16
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.max_len, config.n_layers, config.n_kv_heads, config.dim // config.n_heads)
shape = (
batch_size,
config.max_len,
config.n_layers,
config.n_kv_heads,
config.dim // config.n_heads,
)
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
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
dtype=torch.long,
)
return prompt_ids, gen_ids
@torch.inference_mode()
def run_prefill_benchmark(
self,
@@ -59,32 +65,38 @@ class GenerationBenchmark:
prompt_length: int = 512,
num_trials: int = 10,
) -> BenchmarkResult:
for _ in range(3):
prompt_ids, _ = self._prepare_inputs(batch_size, prompt_length, prompt_length)
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)
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)")
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,
@@ -95,9 +107,9 @@ class GenerationBenchmark:
"prompt_length": prompt_length,
"dtype": self.dtype,
"device": self.device,
}
},
)
@torch.inference_mode()
def run_decoding_benchmark(
self,
@@ -106,39 +118,43 @@ class GenerationBenchmark:
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)
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)
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)")
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,
@@ -150,24 +166,28 @@ class GenerationBenchmark:
"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']}")
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"Batch Size: {result.metadata['batch_size']} | Gen Length: {result.metadata['gen_length']}"
)
print(f"Device: {result.metadata['device']} | Dtype: {result.metadata['dtype']}")
print("-" * 80)
@@ -183,16 +203,19 @@ if __name__ == "__main__":
n_layers=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)
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)
gen_result = benchmark.run_decoding_benchmark(
batch_size=4, prompt_length=512, gen_length=128, num_trials=5
)
print_benchmark_result(gen_result)
tools/generate.py
+52 -16
View File
@@ -21,10 +21,10 @@ def processor(
with disable_random_init():
param = ModelParameter.load(model_dir)
param.to(device='cuda', dtype=torch.bfloat16)
param.to(device="cuda", dtype=torch.bfloat16)
generator = BatchGenerator(param)
with open(input_json_file, "r", encoding='utf-8') as f:
with open(input_json_file, "r", encoding="utf-8") as f:
input_data = [json.loads(line) for line in f]
queries = [item[question_key] for item in input_data]
@@ -41,26 +41,62 @@ def processor(
responses = generator.generate(request)
with open(output_json_file, "w", encoding='utf-8') as f:
with open(output_json_file, "w", encoding="utf-8") as f:
for query, response in zip(queries, responses):
output_item = {question_key: query, response_key: response}
f.write(json.dumps(output_item, ensure_ascii=False) + '\n')
f.write(json.dumps(output_item, ensure_ascii=False) + "\n")
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("--response_key", type=str, default="response", help="Key for the response in the output JSON.")
parser.add_argument("--temperature", type=float, default=0.60, help="Temperature for generating responses.")
parser.add_argument("--top_k", type=int, default=30, help="Top-k value for generating responses.")
parser.add_argument("--top_p", type=float, default=0.95, help="Top-p value for generating responses.")
parser.add_argument("--batch_size", type=int, default=1, help="Batch size for generating responses.")
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(
"--response_key",
type=str,
default="response",
help="Key for the response in the output JSON.",
)
parser.add_argument(
"--temperature",
type=float,
default=0.60,
help="Temperature for generating responses.",
)
parser.add_argument(
"--top_k", type=int, default=30, help="Top-k value for generating responses."
)
parser.add_argument(
"--top_p",
type=float,
default=0.95,
help="Top-p value for generating responses.",
)
parser.add_argument(
"--batch_size", type=int, default=1, help="Batch size for generating responses."
)
args = parser.parse_args()
with torch.inference_mode():
processor(**vars(args))
processor(**vars(args))
tools/perplexity.py
+49 -39
View File
@@ -11,89 +11,99 @@ from khaosz.inference.core import disable_random_init
def compute_perplexity(
model: nn.Module,
input_ids: Tensor,
input_mask: Tensor,
) -> Tensor:
model: nn.Module,
input_ids: Tensor,
input_mask: Tensor,
) -> Tensor:
"""
Compute the perplexity of a batch of input sequences,
where PPL = exp(-(1/N) * sum(log P(w_i | w_<i))).
Compute the perplexity of a batch of input sequences,
where PPL = exp(-(1/N) * sum(log P(w_i | w_<i))).
"""
output = model(input_ids, input_mask)
logits = output["logits"]
shifted_logits = logits[:, :-1, :] # [batch_size, seq_len-1, vocab_size]
shifted_input_ids = input_ids[:, 1:] # [batch_size, seq_len-1]
shifted_mask = input_mask[:, 1:] # [batch_size, seq_len-1]
shifted_logits = logits[:, :-1, :] # [batch_size, seq_len-1, vocab_size]
shifted_input_ids = input_ids[:, 1:] # [batch_size, seq_len-1]
shifted_mask = input_mask[:, 1:] # [batch_size, seq_len-1]
loss = F.cross_entropy(
shifted_logits.flatten(0, 1),
shifted_input_ids.flatten(0, 1),
reduction='none'
shifted_logits.flatten(0, 1), shifted_input_ids.flatten(0, 1), reduction="none"
)
loss = loss.view(shifted_input_ids.shape) # [batch_size, seq_len-1]
loss = loss.view(shifted_input_ids.shape) # [batch_size, seq_len-1]
loss = loss * shifted_mask
sentence_loss = (loss).sum(dim=1) / shifted_mask.sum(dim=1)
perplexity = torch.exp(sentence_loss) # [batch_size]
perplexity = torch.exp(sentence_loss) # [batch_size]
return perplexity
def process_file(
model_dir: str,
input_file: str,
output_file: str,
batch_size: int,
text_key: str
model_dir: str, input_file: str, output_file: str, batch_size: int, text_key: str
):
with disable_random_init():
param = ModelParameter.load(model_dir)
param.to(device='cuda', dtype=torch.bfloat16)
param.to(device="cuda", dtype=torch.bfloat16)
model = param.model
tokenizer = param.tokenizer
with open(input_file, "r", encoding='utf-8') as f:
with open(input_file, "r", encoding="utf-8") as f:
input_data = [json.loads(line) for line in f]
texts = [item[text_key] for item in input_data]
encoded_texts = [tokenizer.encode(text) for text in texts]
output_data = []
for i in tqdm(range(0, len(encoded_texts), batch_size), desc="Computing perplexity"):
batch_encoded = encoded_texts[i:i + batch_size]
batch_texts = texts[i:i + batch_size]
for i in tqdm(
range(0, len(encoded_texts), batch_size), desc="Computing perplexity"
):
batch_encoded = encoded_texts[i : i + batch_size]
batch_texts = texts[i : i + batch_size]
max_len = max(len(seq) for seq in batch_encoded)
padded_ids = []
masks = []
for seq in batch_encoded:
pad_len = max_len - len(seq)
padded_seq = [tokenizer.pad_id] * pad_len + seq
mask = [False] * pad_len + [True] * len(seq)
padded_ids.append(padded_seq)
masks.append(mask)
input_ids = torch.tensor(padded_ids, device="cuda", dtype=torch.long)
input_mask = torch.tensor(masks, device="cuda", dtype=torch.bool)
perplexity = compute_perplexity(model, input_ids, input_mask)
for text, ppl in zip(batch_texts, perplexity):
output_data.append({text_key: text, "ppl": float(ppl.item())})
with open(output_file, "w", encoding='utf-8') as f:
with open(output_file, "w", encoding="utf-8") as f:
for item in output_data:
f.write(json.dumps(item, ensure_ascii=False) + '\n')
f.write(json.dumps(item, ensure_ascii=False) + "\n")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run perplexity with a Khaosz model.")
parser.add_argument("--model_dir", type=str, required=True, help="Path to the model directory.")
parser.add_argument("--input_file", type=str, required=True, help="Path to the input file.")
parser.add_argument("--output_file", type=str, required=True, help="Path to the output file.")
parser.add_argument("--batch_size", type=int, default=4, help="Batch size for evaluation.")
parser.add_argument("--text_key", type=str, default="text", help="Key for the text field in the input data.")
parser.add_argument(
"--model_dir", type=str, required=True, help="Path to the model directory."
)
parser.add_argument(
"--input_file", type=str, required=True, help="Path to the input file."
)
parser.add_argument(
"--output_file", type=str, required=True, help="Path to the output file."
)
parser.add_argument(
"--batch_size", type=int, default=4, help="Batch size for evaluation."
)
parser.add_argument(
"--text_key",
type=str,
default="text",
help="Key for the text field in the input data.",
)
args = parser.parse_args()
with torch.inference_mode():
tools/train.py
+151 -54
View File
@@ -15,40 +15,130 @@ from khaosz.parallel import get_rank
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Train the Transformer model.")
parser.add_argument("--train_type", type=str, required=True, 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("--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping.")
parser.add_argument("--adamw_beta1", type=float, default=0.9, help="Beta values for AdamW optimizer.")
parser.add_argument("--adamw_beta2", type=float, default=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("--random_seed", type=int, default=3407, help="Random seed for reproducibility.")
parser.add_argument("--num_workers", type=int, default=4, help="Number of workers for data loading.")
parser.add_argument("--no_pin_memory", action="store_false", dest="pin_memory", help="Disable pin memory")
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(
"--train_type",
type=str,
required=True,
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(
"--max_grad_norm",
type=float,
default=1.0,
help="Max gradient norm for clipping.",
)
parser.add_argument(
"--adamw_beta1",
type=float,
default=0.9,
help="Beta values for AdamW optimizer.",
)
parser.add_argument(
"--adamw_beta2",
type=float,
default=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(
"--random_seed", type=int, default=3407, help="Random seed for reproducibility."
)
parser.add_argument(
"--num_workers", type=int, default=4, help="Number of workers for data loading."
)
parser.add_argument(
"--no_pin_memory",
action="store_false",
dest="pin_memory",
help="Disable pin memory",
)
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("--dpo_beta", type=float, default=0.1, help="DPO beta value.")
parser.add_argument("--label_smoothing", type=float, default=0.1, help="cross_entropy function label smoothing parameter")
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("--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(
"--label_smoothing",
type=float,
default=0.1,
help="cross_entropy function label smoothing parameter",
)
parser.add_argument(
"--ckpt_interval",
type=int,
default=5000,
help="Number of iters between checkpoints.",
)
parser.add_argument(
"--ckpt_dir",
type=str,
default="checkpoint",
help="Directory to save checkpoints.",
)
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("--nprocs", type=int, default=1, help="Number of GPUs to use.")
parser.add_argument("--device_type", type=str, default="cuda", help="Device type to use.")
parser.add_argument(
"--device_type", type=str, default="cuda", help="Device type to use."
)
args = parser.parse_args()
return args
def ddp_wrap(model: nn.Module):
local_rank = get_rank()
model = model.to(device=f"cuda:{local_rank}", dtype=torch.bfloat16)
@@ -56,16 +146,21 @@ def ddp_wrap(model: nn.Module):
model,
device_ids=[local_rank],
output_device=local_rank,
find_unused_parameters=False
find_unused_parameters=False,
)
return ddp_model
def create_optimizer(model: nn.Module, **kwargs) -> optim.Optimizer:
return optim.AdamW(model.parameters(), **kwargs)
def create_scheduler(optimizer: optim.Optimizer, **kwargs) -> optim.lr_scheduler.LRScheduler:
def create_scheduler(
optimizer: optim.Optimizer, **kwargs
) -> optim.lr_scheduler.LRScheduler:
return SchedulerFactory.load(optimizer, **kwargs)
def prepare_checkpoint(model: nn.Module) -> dict:
return model.module.state_dict()
@@ -81,8 +176,8 @@ def train(
start_batch: int,
accumulation_steps: int,
warmup_steps: int,
checkpoint_interval: int,
checkpoint_dir: str,
ckpt_interval: int,
ckpt_dir: str,
dpo_beta: float,
adamw_beta1: float,
adamw_beta2: float,
@@ -99,48 +194,50 @@ def train(
):
assert train_type in ["seq", "sft", "dpo"]
assert os.path.exists(param_path)
parameter = ModelParameter.load(param_path)
if window_size is None:
window_size = parameter.config.max_len
model = parameter.model
strategy_kwargs = {
"dpo_beta": dpo_beta,
"label_smoothing": label_smoothing
}
strategy_kwargs = {"dpo_beta": dpo_beta, "label_smoothing": label_smoothing}
dataset = DatasetLoader.load(
train_type=train_type,
load_path=data_root_path,
window_size=window_size,
stride=stride
)
schedule_config = CosineScheduleConfig(
warmup_steps=warmup_steps,
total_steps=len(dataset) * n_epoch // (batch_size * nprocs),
stride=stride,
)
optimizer_fn = partial(create_optimizer,
**{"lr": max_lr, "betas": (adamw_beta1, adamw_beta2), "weight_decay": adamw_weight_decay})
scheduler_fn = partial(create_scheduler,
**{"schedule_config": schedule_config})
schedule_config = CosineScheduleConfig(
warmup_steps=warmup_steps,
total_steps=len(dataset) * n_epoch // (batch_size * nprocs),
)
optimizer_fn = partial(
create_optimizer,
**{
"lr": max_lr,
"betas": (adamw_beta1, adamw_beta2),
"weight_decay": adamw_weight_decay,
},
)
scheduler_fn = partial(create_scheduler, **{"schedule_config": schedule_config})
train_config = TrainConfig(
model=model,
strategy=train_type,
dataset=dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
checkpoint_dir=checkpoint_dir,
ckpt_dir=ckpt_dir,
n_epoch=n_epoch,
batch_size=batch_size,
start_epoch=start_epoch,
start_batch=start_batch,
checkpoint_interval=checkpoint_interval,
ckpt_interval=ckpt_interval,
accumulation_steps=accumulation_steps,
max_grad_norm=max_grad_norm,
random_seed=random_seed,
@@ -152,11 +249,11 @@ def train(
device_type=device_type,
extra_kwargs=strategy_kwargs,
)
trainer = Trainer(train_config)
trainer.train()
if __name__ == "__main__":
args = parse_args()
train(**vars(args))
train(**vars(args))