chore: 修改文件夹结构
This commit is contained in:
@@ -0,0 +1,12 @@
|
||||
from pathlib import Path
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
|
||||
|
||||
if __name__ == "__main__":
|
||||
snapshot_download(
|
||||
repo_id="ViperEk/AstrAI",
|
||||
local_dir=PARAMETER_ROOT,
|
||||
force_download=True,
|
||||
)
|
||||
@@ -0,0 +1,35 @@
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from astrai.config.param_config import ModelParameter
|
||||
from astrai.inference.core import disable_random_init
|
||||
from astrai.inference.generator import GeneratorFactory, GenerationRequest
|
||||
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
|
||||
|
||||
|
||||
def generate_text():
|
||||
|
||||
with disable_random_init():
|
||||
param = ModelParameter.load(PARAMETER_ROOT)
|
||||
param.to(device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
query = input(">> ")
|
||||
|
||||
request = GenerationRequest(
|
||||
query=query,
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
top_k=50,
|
||||
max_len=param.config.max_len,
|
||||
history=None,
|
||||
system_prompt=None,
|
||||
)
|
||||
generator = GeneratorFactory.create(param, request)
|
||||
response = generator.generate(request)
|
||||
|
||||
print(response)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
generate_text()
|
||||
@@ -0,0 +1,42 @@
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from astrai.config.param_config import ModelParameter
|
||||
from astrai.inference.core import disable_random_init
|
||||
from astrai.inference.generator import GeneratorFactory, GenerationRequest
|
||||
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
|
||||
|
||||
|
||||
def batch_generate():
|
||||
|
||||
with disable_random_init():
|
||||
param = ModelParameter.load(PARAMETER_ROOT)
|
||||
param.to(device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
inputs = [
|
||||
"你好",
|
||||
"请问什么是人工智能",
|
||||
"今天天气如何",
|
||||
"我感到焦虑， 请问我应该怎么办",
|
||||
"请问什么是显卡",
|
||||
]
|
||||
|
||||
request = GenerationRequest(
|
||||
query=inputs,
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
top_k=50,
|
||||
max_len=param.config.max_len,
|
||||
history=None,
|
||||
system_prompt=None,
|
||||
)
|
||||
generator = GeneratorFactory.create(param, request)
|
||||
responses = generator.generate(request)
|
||||
|
||||
for q, r in zip(inputs, responses):
|
||||
print((q, r))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
batch_generate()
|
||||
@@ -0,0 +1,47 @@
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from astrai.config.param_config import ModelParameter
|
||||
from astrai.inference.core import disable_random_init
|
||||
from astrai.inference.generator import GeneratorFactory, GenerationRequest
|
||||
|
||||
PROJECT_ROOT = Path(__file__).parent.parent
|
||||
PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
|
||||
|
||||
|
||||
def chat():
|
||||
|
||||
with disable_random_init():
|
||||
param = ModelParameter.load(PARAMETER_ROOT)
|
||||
param.to(device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
history = []
|
||||
while True:
|
||||
query = input(">> ")
|
||||
if query == "!exit":
|
||||
break
|
||||
|
||||
request = GenerationRequest(
|
||||
query=query,
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
top_k=50,
|
||||
max_len=param.config.max_len,
|
||||
history=history,
|
||||
system_prompt=None,
|
||||
)
|
||||
generator = GeneratorFactory.create(param, request)
|
||||
|
||||
response_size = 0
|
||||
full_response = ""
|
||||
for response in generator.generate(request):
|
||||
# response is the cumulative response up to current token
|
||||
print(response[response_size:], end="", flush=True)
|
||||
response_size = len(response)
|
||||
full_response = response
|
||||
|
||||
# After generation, update history
|
||||
history.append((query, full_response.strip()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
chat()
|
||||
@@ -0,0 +1,221 @@
|
||||
import torch
|
||||
from typing import Dict, Any
|
||||
from dataclasses import dataclass
|
||||
from astrai.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.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,
|
||||
)
|
||||
|
||||
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,
|
||||
dim=1536,
|
||||
n_heads=24,
|
||||
n_kv_heads=4,
|
||||
dim_ffn=6912,
|
||||
max_len=2048,
|
||||
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
|
||||
)
|
||||
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,102 @@
|
||||
import torch
|
||||
import json
|
||||
import argparse
|
||||
|
||||
from astrai.config.param_config import ModelParameter
|
||||
from astrai.inference.generator import BatchGenerator, GenerationRequest
|
||||
from astrai.inference.core import disable_random_init
|
||||
|
||||
|
||||
def processor(
|
||||
model_dir: str,
|
||||
input_json_file: str,
|
||||
output_json_file: str,
|
||||
batch_size: int,
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
question_key: str,
|
||||
response_key: str,
|
||||
):
|
||||
with disable_random_init():
|
||||
param = ModelParameter.load(model_dir)
|
||||
|
||||
param.to(device="cuda", dtype=torch.bfloat16)
|
||||
generator = BatchGenerator(param)
|
||||
|
||||
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]
|
||||
|
||||
request = GenerationRequest(
|
||||
query=queries,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
top_k=top_k,
|
||||
max_len=param.config.max_len,
|
||||
history=None,
|
||||
system_prompt=None,
|
||||
)
|
||||
|
||||
responses = generator.generate(request)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
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."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
with torch.inference_mode():
|
||||
processor(**vars(args))
|
||||
@@ -0,0 +1,110 @@
|
||||
import json
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import argparse
|
||||
import tqdm
|
||||
|
||||
from torch import Tensor
|
||||
from astrai.config.param_config import ModelParameter
|
||||
from astrai.inference.core import disable_random_init
|
||||
|
||||
|
||||
def compute_perplexity(
|
||||
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))).
|
||||
"""
|
||||
|
||||
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]
|
||||
|
||||
loss = F.cross_entropy(
|
||||
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 * shifted_mask
|
||||
sentence_loss = (loss).sum(dim=1) / shifted_mask.sum(dim=1)
|
||||
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
|
||||
):
|
||||
with disable_random_init():
|
||||
param = ModelParameter.load(model_dir)
|
||||
|
||||
param.to(device="cuda", dtype=torch.bfloat16)
|
||||
model = param.model
|
||||
tokenizer = param.tokenizer
|
||||
|
||||
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]
|
||||
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:
|
||||
for item in output_data:
|
||||
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.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
with torch.inference_mode():
|
||||
process_file(**vars(args))
|
||||
@@ -0,0 +1,259 @@
|
||||
import os
|
||||
import argparse
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
|
||||
from functools import partial
|
||||
from astrai.data import DatasetLoader
|
||||
from astrai.config import ModelParameter, TrainConfig, CosineScheduleConfig
|
||||
from astrai.trainer import Trainer, SchedulerFactory
|
||||
from astrai.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("--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(
|
||||
"--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."
|
||||
)
|
||||
|
||||
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)
|
||||
ddp_model = DDP(
|
||||
model,
|
||||
device_ids=[local_rank],
|
||||
output_device=local_rank,
|
||||
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:
|
||||
return SchedulerFactory.load(optimizer, **kwargs)
|
||||
|
||||
|
||||
def prepare_checkpoint(model: nn.Module) -> dict:
|
||||
return model.module.state_dict()
|
||||
|
||||
|
||||
def train(
|
||||
train_type: str,
|
||||
param_path: str,
|
||||
data_root_path: str,
|
||||
max_lr: float,
|
||||
n_epoch: int,
|
||||
batch_size: int,
|
||||
start_epoch: int,
|
||||
start_batch: int,
|
||||
accumulation_steps: int,
|
||||
warmup_steps: int,
|
||||
ckpt_interval: int,
|
||||
ckpt_dir: str,
|
||||
dpo_beta: float,
|
||||
adamw_beta1: float,
|
||||
adamw_beta2: float,
|
||||
adamw_weight_decay: float,
|
||||
max_grad_norm: float,
|
||||
label_smoothing: float,
|
||||
random_seed: int,
|
||||
num_workers: int,
|
||||
pin_memory: bool,
|
||||
window_size: int,
|
||||
stride: int,
|
||||
nprocs: int,
|
||||
device_type: str,
|
||||
):
|
||||
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}
|
||||
|
||||
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),
|
||||
)
|
||||
|
||||
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,
|
||||
ckpt_dir=ckpt_dir,
|
||||
n_epoch=n_epoch,
|
||||
batch_size=batch_size,
|
||||
start_epoch=start_epoch,
|
||||
start_batch=start_batch,
|
||||
ckpt_interval=ckpt_interval,
|
||||
accumulation_steps=accumulation_steps,
|
||||
max_grad_norm=max_grad_norm,
|
||||
random_seed=random_seed,
|
||||
num_workers=num_workers,
|
||||
pin_memory=pin_memory,
|
||||
nprocs=nprocs,
|
||||
parallel_wrapper=ddp_wrap,
|
||||
state_dict_fn=prepare_checkpoint,
|
||||
device_type=device_type,
|
||||
extra_kwargs=strategy_kwargs,
|
||||
)
|
||||
|
||||
trainer = Trainer(train_config)
|
||||
trainer.train()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
train(**vars(args))
|
||||
Reference in New Issue
Block a user