add finetune scripts
This commit is contained in:
shiyu-coder
@@ -0,0 +1,131 @@
|
||||
import os
|
||||
|
||||
class Config:
|
||||
"""
|
||||
Configuration class for the entire project.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# =================================================================
|
||||
# Data & Feature Parameters
|
||||
# =================================================================
|
||||
# TODO: Update this path to your Qlib data directory.
|
||||
self.qlib_data_path = "~/.qlib/qlib_data/cn_data"
|
||||
self.instrument = 'csi300'
|
||||
|
||||
# Overall time range for data loading from Qlib.
|
||||
self.dataset_begin_time = "2011-01-01"
|
||||
self.dataset_end_time = '2025-06-05'
|
||||
|
||||
# Sliding window parameters for creating samples.
|
||||
self.lookback_window = 90 # Number of past time steps for input.
|
||||
self.predict_window = 10 # Number of future time steps for prediction.
|
||||
self.max_context = 512 # Maximum context length for the model.
|
||||
|
||||
# Features to be used from the raw data.
|
||||
self.feature_list = ['open', 'high', 'low', 'close', 'vol', 'amt']
|
||||
# Time-based features to be generated.
|
||||
self.time_feature_list = ['minute', 'hour', 'weekday', 'day', 'month']
|
||||
|
||||
# =================================================================
|
||||
# Dataset Splitting & Paths
|
||||
# =================================================================
|
||||
# Note: The validation/test set starts earlier than the training/validation set ends
|
||||
# to account for the `lookback_window`.
|
||||
self.train_time_range = ["2011-01-01", "2022-12-31"]
|
||||
self.val_time_range = ["2022-09-01", "2024-06-30"]
|
||||
self.test_time_range = ["2024-04-01", "2025-06-05"]
|
||||
self.backtest_time_range = ["2024-07-01", "2025-06-05"]
|
||||
|
||||
# TODO: Directory to save the processed, pickled datasets.
|
||||
self.dataset_path = "./data/processed_datasets"
|
||||
|
||||
# =================================================================
|
||||
# Training Hyperparameters
|
||||
# =================================================================
|
||||
self.clip = 5.0 # Clipping value for normalized data to prevent outliers.
|
||||
|
||||
self.epochs = 30
|
||||
self.log_interval = 100 # Log training status every N batches.
|
||||
self.batch_size = 50 # Batch size per GPU.
|
||||
|
||||
# Number of samples to draw for one "epoch" of training/validation.
|
||||
# This is useful for large datasets where a true epoch is too long.
|
||||
self.n_train_iter = 2000 * self.batch_size
|
||||
self.n_val_iter = 400 * self.batch_size
|
||||
|
||||
# Learning rates for different model components.
|
||||
self.tokenizer_learning_rate = 2e-4
|
||||
self.predictor_learning_rate = 4e-5
|
||||
|
||||
# Gradient accumulation to simulate a larger batch size.
|
||||
self.accumulation_steps = 1
|
||||
|
||||
# AdamW optimizer parameters.
|
||||
self.adam_beta1 = 0.9
|
||||
self.adam_beta2 = 0.95
|
||||
self.adam_weight_decay = 0.1
|
||||
|
||||
# Miscellaneous
|
||||
self.seed = 100 # Global random seed for reproducibility.
|
||||
|
||||
# =================================================================
|
||||
# Experiment Logging & Saving
|
||||
# =================================================================
|
||||
self.use_comet = True # Set to False if you don't want to use Comet ML
|
||||
self.comet_config = {
|
||||
# It is highly recommended to load secrets from environment variables
|
||||
# for security purposes. Example: os.getenv("COMET_API_KEY")
|
||||
"api_key": "YOUR_COMET_API_KEY",
|
||||
"project_name": "Kronos-Finetune-Demo",
|
||||
"workspace": "your_comet_workspace" # TODO: Change to your Comet ML workspace name
|
||||
}
|
||||
self.comet_tag = 'finetune_demo'
|
||||
self.comet_name = 'finetune_demo'
|
||||
|
||||
# Base directory for saving model checkpoints and results.
|
||||
# Using a general 'outputs' directory is a common practice.
|
||||
self.save_path = "./outputs/models"
|
||||
self.tokenizer_save_folder_name = 'finetune_tokenizer_demo'
|
||||
self.predictor_save_folder_name = 'finetune_predictor_demo'
|
||||
self.backtest_save_folder_name = 'finetune_backtest_demo'
|
||||
|
||||
# Path for backtesting results.
|
||||
self.backtest_result_path = "./outputs/backtest_results"
|
||||
|
||||
# =================================================================
|
||||
# Model & Checkpoint Paths
|
||||
# =================================================================
|
||||
# TODO: Update these paths to your pretrained model locations.
|
||||
# These can be local paths or Hugging Face Hub model identifiers.
|
||||
self.pretrained_tokenizer_path = "path/to/your/Kronos-Tokenizer-base"
|
||||
self.pretrained_predictor_path = "path/to/your/Kronos-small"
|
||||
|
||||
# Paths to the fine-tuned models, derived from the save_path.
|
||||
# These will be generated automatically during training.
|
||||
self.finetuned_tokenizer_path = f"{self.save_path}/{self.tokenizer_save_folder_name}/checkpoints/best_model"
|
||||
self.finetuned_predictor_path = f"{self.save_path}/{self.predictor_save_folder_name}/checkpoints/best_model"
|
||||
|
||||
# =================================================================
|
||||
# Backtesting Parameters
|
||||
# =================================================================
|
||||
self.backtest_n_symbol_hold = 50 # Number of symbols to hold in the portfolio.
|
||||
self.backtest_n_symbol_drop = 5 # Number of symbols to drop from the pool.
|
||||
self.backtest_hold_thresh = 5 # Minimum holding period for a stock.
|
||||
self.inference_T = 0.6
|
||||
self.inference_top_p = 0.9
|
||||
self.inference_top_k = 0
|
||||
self.inference_sample_count = 5
|
||||
self.backtest_batch_size = 1000
|
||||
self.backtest_benchmark = self._set_benchmark(self.instrument)
|
||||
|
||||
def _set_benchmark(self, instrument):
|
||||
dt_benchmark = {
|
||||
'csi800': "SH000906",
|
||||
'csi1000': "SH000852",
|
||||
'csi300': "SH000300",
|
||||
}
|
||||
if instrument in dt_benchmark:
|
||||
return dt_benchmark[instrument]
|
||||
else:
|
||||
raise ValueError(f"Benchmark not defined for instrument: {instrument}")
|
||||
@@ -0,0 +1,145 @@
|
||||
import pickle
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from config import Config
|
||||
|
||||
|
||||
class QlibDataset(Dataset):
|
||||
"""
|
||||
A PyTorch Dataset for handling Qlib financial time series data.
|
||||
|
||||
This dataset pre-computes all possible start indices for sliding windows
|
||||
and then randomly samples from them during training/validation.
|
||||
|
||||
Args:
|
||||
data_type (str): The type of dataset to load, either 'train' or 'val'.
|
||||
|
||||
Raises:
|
||||
ValueError: If `data_type` is not 'train' or 'val'.
|
||||
"""
|
||||
|
||||
def __init__(self, data_type: str = 'train'):
|
||||
self.config = Config()
|
||||
if data_type not in ['train', 'val']:
|
||||
raise ValueError("data_type must be 'train' or 'val'")
|
||||
self.data_type = data_type
|
||||
|
||||
# Use a dedicated random number generator for sampling to avoid
|
||||
# interfering with other random processes (e.g., in model initialization).
|
||||
self.py_rng = random.Random(self.config.seed)
|
||||
|
||||
# Set paths and number of samples based on the data type.
|
||||
if data_type == 'train':
|
||||
self.data_path = f"{self.config.dataset_path}/train_data.pkl"
|
||||
self.n_samples = self.config.n_train_iter
|
||||
else:
|
||||
self.data_path = f"{self.config.dataset_path}/val_data.pkl"
|
||||
self.n_samples = self.config.n_val_iter
|
||||
|
||||
with open(self.data_path, 'rb') as f:
|
||||
self.data = pickle.load(f)
|
||||
|
||||
self.window = self.config.lookback_window + self.config.predict_window + 1
|
||||
|
||||
self.symbols = list(self.data.keys())
|
||||
self.feature_list = self.config.feature_list
|
||||
self.time_feature_list = self.config.time_feature_list
|
||||
|
||||
# Pre-compute all possible (symbol, start_index) pairs.
|
||||
self.indices = []
|
||||
print(f"[{data_type.upper()}] Pre-computing sample indices...")
|
||||
for symbol in self.symbols:
|
||||
df = self.data[symbol].reset_index()
|
||||
series_len = len(df)
|
||||
num_samples = series_len - self.window + 1
|
||||
|
||||
if num_samples > 0:
|
||||
# Generate time features and store them directly in the dataframe.
|
||||
df['minute'] = df['datetime'].dt.minute
|
||||
df['hour'] = df['datetime'].dt.hour
|
||||
df['weekday'] = df['datetime'].dt.weekday
|
||||
df['day'] = df['datetime'].dt.day
|
||||
df['month'] = df['datetime'].dt.month
|
||||
# Keep only necessary columns to save memory.
|
||||
self.data[symbol] = df[self.feature_list + self.time_feature_list]
|
||||
|
||||
# Add all valid starting indices for this symbol to the global list.
|
||||
for i in range(num_samples):
|
||||
self.indices.append((symbol, i))
|
||||
|
||||
# The effective dataset size is the minimum of the configured iterations
|
||||
# and the total number of available samples.
|
||||
self.n_samples = min(self.n_samples, len(self.indices))
|
||||
print(f"[{data_type.upper()}] Found {len(self.indices)} possible samples. Using {self.n_samples} per epoch.")
|
||||
|
||||
def set_epoch_seed(self, epoch: int):
|
||||
"""
|
||||
Sets a new seed for the random sampler for each epoch. This is crucial
|
||||
for reproducibility in distributed training.
|
||||
|
||||
Args:
|
||||
epoch (int): The current epoch number.
|
||||
"""
|
||||
epoch_seed = self.config.seed + epoch
|
||||
self.py_rng.seed(epoch_seed)
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Returns the number of samples per epoch."""
|
||||
return self.n_samples
|
||||
|
||||
def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Retrieves a random sample from the dataset.
|
||||
|
||||
Note: The `idx` argument is ignored. Instead, a random index is drawn
|
||||
from the pre-computed `self.indices` list using `self.py_rng`. This
|
||||
ensures random sampling over the entire dataset for each call.
|
||||
|
||||
Args:
|
||||
idx (int): Ignored.
|
||||
|
||||
Returns:
|
||||
tuple[torch.Tensor, torch.Tensor]: A tuple containing:
|
||||
- x_tensor (torch.Tensor): The normalized feature tensor.
|
||||
- x_stamp_tensor (torch.Tensor): The time feature tensor.
|
||||
"""
|
||||
# Select a random sample from the entire pool of indices.
|
||||
random_idx = self.py_rng.randint(0, len(self.indices) - 1)
|
||||
symbol, start_idx = self.indices[random_idx]
|
||||
|
||||
# Extract the sliding window from the dataframe.
|
||||
df = self.data[symbol]
|
||||
end_idx = start_idx + self.window
|
||||
win_df = df.iloc[start_idx:end_idx]
|
||||
|
||||
# Separate main features and time features.
|
||||
x = win_df[self.feature_list].values.astype(np.float32)
|
||||
x_stamp = win_df[self.time_feature_list].values.astype(np.float32)
|
||||
|
||||
# Perform instance-level normalization.
|
||||
x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
|
||||
x = (x - x_mean) / (x_std + 1e-5)
|
||||
x = np.clip(x, -self.config.clip, self.config.clip)
|
||||
|
||||
# Convert to PyTorch tensors.
|
||||
x_tensor = torch.from_numpy(x)
|
||||
x_stamp_tensor = torch.from_numpy(x_stamp)
|
||||
|
||||
return x_tensor, x_stamp_tensor
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Example usage and verification.
|
||||
print("Creating training dataset instance...")
|
||||
train_dataset = QlibDataset(data_type='train')
|
||||
|
||||
print(f"Dataset length: {len(train_dataset)}")
|
||||
|
||||
if len(train_dataset) > 0:
|
||||
try_x, try_x_stamp = train_dataset[100] # Index 100 is ignored.
|
||||
print(f"Sample feature shape: {try_x.shape}")
|
||||
print(f"Sample time feature shape: {try_x_stamp.shape}")
|
||||
else:
|
||||
print("Dataset is empty.")
|
||||
@@ -0,0 +1,120 @@
|
||||
import os
|
||||
import pickle
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import qlib
|
||||
from qlib.config import REG_CN
|
||||
from qlib.data import D
|
||||
from qlib.data.dataset.loader import QlibDataLoader
|
||||
from tqdm import trange
|
||||
|
||||
from config import Config
|
||||
|
||||
|
||||
class QlibDataPreprocessor:
|
||||
"""
|
||||
A class to handle the loading, processing, and splitting of Qlib financial data.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initializes the preprocessor with configuration and data fields."""
|
||||
self.config = Config()
|
||||
self.data_fields = ['open', 'close', 'high', 'low', 'volume', 'vwap']
|
||||
self.data = {} # A dictionary to store processed data for each symbol.
|
||||
|
||||
def initialize_qlib(self):
|
||||
"""Initializes the Qlib environment."""
|
||||
print("Initializing Qlib...")
|
||||
qlib.init(provider_uri=self.config.qlib_data_path, region=REG_CN)
|
||||
|
||||
def load_qlib_data(self):
|
||||
"""
|
||||
Loads raw data from Qlib, processes it symbol by symbol, and stores
|
||||
it in the `self.data` attribute.
|
||||
"""
|
||||
print("Loading and processing data from Qlib...")
|
||||
data_fields_qlib = ['$' + f for f in self.data_fields]
|
||||
cal: np.ndarray = D.calendar()
|
||||
|
||||
# Determine the actual start and end times to load, including buffer for lookback and predict windows.
|
||||
start_index = cal.searchsorted(pd.Timestamp(self.config.dataset_begin_time))
|
||||
end_index = cal.searchsorted(pd.Timestamp(self.config.dataset_end_time))
|
||||
real_start_time = cal[start_index - self.config.lookback_window]
|
||||
|
||||
if cal[end_index] != pd.Timestamp(self.config.dataset_end_time):
|
||||
end_index -= 1
|
||||
real_end_time = cal[end_index + self.config.predict_window]
|
||||
|
||||
# Load data using Qlib's data loader.
|
||||
data_df = QlibDataLoader(config=data_fields_qlib).load(
|
||||
self.config.instrument, real_start_time, real_end_time
|
||||
)
|
||||
data_df = data_df.stack().unstack(level=1) # Reshape for easier access.
|
||||
|
||||
symbol_list = list(data_df.columns)
|
||||
for i in trange(len(symbol_list), desc="Processing Symbols"):
|
||||
symbol = symbol_list[i]
|
||||
symbol_df = data_df[symbol]
|
||||
|
||||
# Pivot the table to have features as columns and datetime as index.
|
||||
symbol_df = symbol_df.reset_index().rename(columns={'level_1': 'field'})
|
||||
symbol_df = pd.pivot(symbol_df, index='datetime', columns='field', values=symbol)
|
||||
symbol_df = symbol_df.rename(columns={f'${field}': field for field in self.data_fields})
|
||||
|
||||
# Calculate amount and select final features.
|
||||
symbol_df['vol'] = symbol_df['volume']
|
||||
symbol_df['amt'] = (symbol_df['open'] + symbol_df['high'] + symbol_df['low'] + symbol_df['close']) / 4 * symbol_df['vol']
|
||||
symbol_df = symbol_df[self.config.feature_list]
|
||||
|
||||
# Filter out symbols with insufficient data.
|
||||
symbol_df = symbol_df.dropna()
|
||||
if len(symbol_df) < self.config.lookback_window + self.config.predict_window + 1:
|
||||
continue
|
||||
|
||||
self.data[symbol] = symbol_df
|
||||
|
||||
def prepare_dataset(self):
|
||||
"""
|
||||
Splits the loaded data into train, validation, and test sets and saves them to disk.
|
||||
"""
|
||||
print("Splitting data into train, validation, and test sets...")
|
||||
train_data, val_data, test_data = {}, {}, {}
|
||||
|
||||
symbol_list = list(self.data.keys())
|
||||
for i in trange(len(symbol_list), desc="Preparing Datasets"):
|
||||
symbol = symbol_list[i]
|
||||
symbol_df = self.data[symbol]
|
||||
|
||||
# Define time ranges from config.
|
||||
train_start, train_end = self.config.train_time_range
|
||||
val_start, val_end = self.config.val_time_range
|
||||
test_start, test_end = self.config.test_time_range
|
||||
|
||||
# Create boolean masks for each dataset split.
|
||||
train_mask = (symbol_df.index >= train_start) & (symbol_df.index <= train_end)
|
||||
val_mask = (symbol_df.index >= val_start) & (symbol_df.index <= val_end)
|
||||
test_mask = (symbol_df.index >= test_start) & (symbol_df.index <= test_end)
|
||||
|
||||
# Apply masks to create the final datasets.
|
||||
train_data[symbol] = symbol_df[train_mask]
|
||||
val_data[symbol] = symbol_df[val_mask]
|
||||
test_data[symbol] = symbol_df[test_mask]
|
||||
|
||||
# Save the datasets using pickle.
|
||||
os.makedirs(self.config.dataset_path, exist_ok=True)
|
||||
with open(f"{self.config.dataset_path}/train_data.pkl", 'wb') as f:
|
||||
pickle.dump(train_data, f)
|
||||
with open(f"{self.config.dataset_path}/val_data.pkl", 'wb') as f:
|
||||
pickle.dump(val_data, f)
|
||||
with open(f"{self.config.dataset_path}/test_data.pkl", 'wb') as f:
|
||||
pickle.dump(test_data, f)
|
||||
|
||||
print("Datasets prepared and saved successfully.")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# This block allows the script to be run directly to perform data preprocessing.
|
||||
preprocessor = QlibDataPreprocessor()
|
||||
preprocessor.initialize_qlib()
|
||||
preprocessor.load_qlib_data()
|
||||
preprocessor.prepare_dataset()
|
||||
@@ -0,0 +1,358 @@
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
import pickle
|
||||
from collections import defaultdict
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from tqdm import trange, tqdm
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
import qlib
|
||||
from qlib.config import REG_CN
|
||||
from qlib.backtest import backtest, executor, CommonInfrastructure
|
||||
from qlib.contrib.evaluate import risk_analysis
|
||||
from qlib.contrib.strategy import TopkDropoutStrategy
|
||||
from qlib.utils import flatten_dict
|
||||
from qlib.utils.time import Freq
|
||||
|
||||
# Ensure project root is in the Python path
|
||||
sys.path.append("../")
|
||||
from config import Config
|
||||
from model.kronos import Kronos, KronosTokenizer, auto_regressive_inference
|
||||
|
||||
|
||||
# =================================================================================
|
||||
# 1. Data Loading and Processing for Inference
|
||||
# =================================================================================
|
||||
|
||||
class QlibTestDataset(Dataset):
|
||||
"""
|
||||
PyTorch Dataset for handling Qlib test data, specifically for inference.
|
||||
|
||||
This dataset iterates through all possible sliding windows sequentially. It also
|
||||
yields metadata like symbol and timestamp, which are crucial for mapping
|
||||
predictions back to the original time series.
|
||||
"""
|
||||
|
||||
def __init__(self, data: dict, config: Config):
|
||||
self.data = data
|
||||
self.config = config
|
||||
self.window_size = config.lookback_window + config.predict_window
|
||||
self.symbols = list(self.data.keys())
|
||||
self.feature_list = config.feature_list
|
||||
self.time_feature_list = config.time_feature_list
|
||||
self.indices = []
|
||||
|
||||
print("Preprocessing and building indices for test dataset...")
|
||||
for symbol in self.symbols:
|
||||
df = self.data[symbol].reset_index()
|
||||
# Generate time features on-the-fly
|
||||
df['minute'] = df['datetime'].dt.minute
|
||||
df['hour'] = df['datetime'].dt.hour
|
||||
df['weekday'] = df['datetime'].dt.weekday
|
||||
df['day'] = df['datetime'].dt.day
|
||||
df['month'] = df['datetime'].dt.month
|
||||
self.data[symbol] = df # Store preprocessed dataframe
|
||||
|
||||
num_samples = len(df) - self.window_size + 1
|
||||
if num_samples > 0:
|
||||
for i in range(num_samples):
|
||||
timestamp = df.iloc[i + self.config.lookback_window - 1]['datetime']
|
||||
self.indices.append((symbol, i, timestamp))
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.indices)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
symbol, start_idx, timestamp = self.indices[idx]
|
||||
df = self.data[symbol]
|
||||
|
||||
context_end = start_idx + self.config.lookback_window
|
||||
predict_end = context_end + self.config.predict_window
|
||||
|
||||
context_df = df.iloc[start_idx:context_end]
|
||||
predict_df = df.iloc[context_end:predict_end]
|
||||
|
||||
x = context_df[self.feature_list].values.astype(np.float32)
|
||||
x_stamp = context_df[self.time_feature_list].values.astype(np.float32)
|
||||
y_stamp = predict_df[self.time_feature_list].values.astype(np.float32)
|
||||
|
||||
# Instance-level normalization, consistent with training
|
||||
x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)
|
||||
x = (x - x_mean) / (x_std + 1e-5)
|
||||
x = np.clip(x, -self.config.clip, self.config.clip)
|
||||
|
||||
return torch.from_numpy(x), torch.from_numpy(x_stamp), torch.from_numpy(y_stamp), symbol, timestamp
|
||||
|
||||
|
||||
# =================================================================================
|
||||
# 2. Backtesting Logic
|
||||
# =================================================================================
|
||||
|
||||
class QlibBacktest:
|
||||
"""
|
||||
A wrapper class for conducting backtesting experiments using Qlib.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Config):
|
||||
self.config = config
|
||||
self.initialize_qlib()
|
||||
|
||||
def initialize_qlib(self):
|
||||
"""Initializes the Qlib environment."""
|
||||
print("Initializing Qlib for backtesting...")
|
||||
qlib.init(provider_uri=self.config.qlib_data_path, region=REG_CN)
|
||||
|
||||
def run_single_backtest(self, signal_series: pd.Series) -> pd.DataFrame:
|
||||
"""
|
||||
Runs a single backtest for a given prediction signal.
|
||||
|
||||
Args:
|
||||
signal_series (pd.Series): A pandas Series with a MultiIndex
|
||||
(instrument, datetime) and prediction scores.
|
||||
Returns:
|
||||
pd.DataFrame: A DataFrame containing the performance report.
|
||||
"""
|
||||
strategy = TopkDropoutStrategy(
|
||||
topk=self.config.backtest_n_symbol_hold,
|
||||
n_drop=self.config.backtest_n_symbol_drop,
|
||||
hold_thresh=self.config.backtest_hold_thresh,
|
||||
signal=signal_series,
|
||||
)
|
||||
executor_config = {
|
||||
"time_per_step": "day",
|
||||
"generate_portfolio_metrics": True,
|
||||
"delay_execution": True,
|
||||
}
|
||||
backtest_config = {
|
||||
"start_time": self.config.backtest_time_range[0],
|
||||
"end_time": self.config.backtest_time_range[1],
|
||||
"account": 100_000_000,
|
||||
"benchmark": self.config.backtest_benchmark,
|
||||
"exchange_kwargs": {
|
||||
"freq": "day", "limit_threshold": 0.095, "deal_price": "open",
|
||||
"open_cost": 0.001, "close_cost": 0.0015, "min_cost": 5,
|
||||
},
|
||||
"executor": executor.SimulatorExecutor(**executor_config),
|
||||
}
|
||||
|
||||
portfolio_metric_dict, _ = backtest(strategy=strategy, **backtest_config)
|
||||
analysis_freq = "{0}{1}".format(*Freq.parse("day"))
|
||||
report, _ = portfolio_metric_dict.get(analysis_freq)
|
||||
|
||||
# --- Analysis and Reporting ---
|
||||
analysis = {
|
||||
"excess_return_without_cost": risk_analysis(report["return"] - report["bench"], freq=analysis_freq),
|
||||
"excess_return_with_cost": risk_analysis(report["return"] - report["bench"] - report["cost"], freq=analysis_freq),
|
||||
}
|
||||
print("\n--- Backtest Analysis ---")
|
||||
print("Benchmark Return:", risk_analysis(report["bench"], freq=analysis_freq), sep='\n')
|
||||
print("\nExcess Return (w/o cost):", analysis["excess_return_without_cost"], sep='\n')
|
||||
print("\nExcess Return (w/ cost):", analysis["excess_return_with_cost"], sep='\n')
|
||||
|
||||
report_df = pd.DataFrame({
|
||||
"cum_bench": report["bench"].cumsum(),
|
||||
"cum_return_w_cost": (report["return"] - report["cost"]).cumsum(),
|
||||
"cum_ex_return_w_cost": (report["return"] - report["bench"] - report["cost"]).cumsum(),
|
||||
})
|
||||
return report_df
|
||||
|
||||
def run_and_plot_results(self, signals: dict[str, pd.DataFrame]):
|
||||
"""
|
||||
Runs backtests for multiple signals and plots the cumulative return curves.
|
||||
|
||||
Args:
|
||||
signals (dict[str, pd.DataFrame]): A dictionary where keys are signal names
|
||||
and values are prediction DataFrames.
|
||||
"""
|
||||
return_df, ex_return_df, bench_df = pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
|
||||
|
||||
for signal_name, pred_df in signals.items():
|
||||
print(f"\nBacktesting signal: {signal_name}...")
|
||||
pred_series = pred_df.stack()
|
||||
pred_series.index.names = ['datetime', 'instrument']
|
||||
pred_series = pred_series.swaplevel().sort_index()
|
||||
report_df = self.run_single_backtest(pred_series)
|
||||
|
||||
return_df[signal_name] = report_df['cum_return_w_cost']
|
||||
ex_return_df[signal_name] = report_df['cum_ex_return_w_cost']
|
||||
if 'return' not in bench_df:
|
||||
bench_df['return'] = report_df['cum_bench']
|
||||
|
||||
# Plotting results
|
||||
fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)
|
||||
return_df.plot(ax=axes[0], title='Cumulative Return with Cost', grid=True)
|
||||
axes[0].plot(bench_df['return'], label=self.config.instrument.upper(), color='black', linestyle='--')
|
||||
axes[0].legend()
|
||||
axes[0].set_ylabel("Cumulative Return")
|
||||
|
||||
ex_return_df.plot(ax=axes[1], title='Cumulative Excess Return with Cost', grid=True)
|
||||
axes[1].legend()
|
||||
axes[1].set_xlabel("Date")
|
||||
axes[1].set_ylabel("Cumulative Excess Return")
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig("../figures/backtest_result_example.png", dpi=200)
|
||||
plt.show()
|
||||
|
||||
|
||||
# =================================================================================
|
||||
# 3. Inference Logic
|
||||
# =================================================================================
|
||||
|
||||
def load_models(config: dict) -> tuple[KronosTokenizer, Kronos]:
|
||||
"""Loads the fine-tuned tokenizer and predictor model."""
|
||||
device = torch.device(config['device'])
|
||||
print(f"Loading models onto device: {device}...")
|
||||
tokenizer = KronosTokenizer.from_pretrained(config['tokenizer_path']).to(device).eval()
|
||||
model = Kronos.from_pretrained(config['model_path']).to(device).eval()
|
||||
return tokenizer, model
|
||||
|
||||
|
||||
def collate_fn_for_inference(batch):
|
||||
"""
|
||||
Custom collate function to handle batches containing Tensors, strings, and Timestamps.
|
||||
|
||||
Args:
|
||||
batch (list): A list of samples, where each sample is the tuple returned by
|
||||
QlibTestDataset.__getitem__.
|
||||
|
||||
Returns:
|
||||
A single tuple containing the batched data.
|
||||
"""
|
||||
# Unzip the list of samples into separate lists for each data type
|
||||
x, x_stamp, y_stamp, symbols, timestamps = zip(*batch)
|
||||
|
||||
# Stack the tensors to create a batch
|
||||
x_batch = torch.stack(x, dim=0)
|
||||
x_stamp_batch = torch.stack(x_stamp, dim=0)
|
||||
y_stamp_batch = torch.stack(y_stamp, dim=0)
|
||||
|
||||
# Return the strings and timestamps as lists
|
||||
return x_batch, x_stamp_batch, y_stamp_batch, list(symbols), list(timestamps)
|
||||
|
||||
|
||||
def generate_predictions(config: dict, test_data: dict) -> dict[str, pd.DataFrame]:
|
||||
"""
|
||||
Runs inference on the test dataset to generate prediction signals.
|
||||
|
||||
Args:
|
||||
config (dict): A dictionary containing inference parameters.
|
||||
test_data (dict): The raw test data loaded from a pickle file.
|
||||
|
||||
Returns:
|
||||
A dictionary where keys are signal types (e.g., 'mean', 'last') and
|
||||
values are DataFrames of predictions (datetime index, symbol columns).
|
||||
"""
|
||||
tokenizer, model = load_models(config)
|
||||
device = torch.device(config['device'])
|
||||
|
||||
# Use the Dataset and DataLoader for efficient batching and processing
|
||||
dataset = QlibTestDataset(data=test_data, config=Config())
|
||||
loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=config['batch_size'] // config['sample_count'],
|
||||
shuffle=False,
|
||||
num_workers=os.cpu_count() // 2,
|
||||
collate_fn=collate_fn_for_inference
|
||||
)
|
||||
|
||||
results = defaultdict(list)
|
||||
with torch.no_grad():
|
||||
for x, x_stamp, y_stamp, symbols, timestamps in tqdm(loader, desc="Inference"):
|
||||
preds = auto_regressive_inference(
|
||||
tokenizer, model, x.to(device), x_stamp.to(device), y_stamp.to(device),
|
||||
max_context=config['max_context'], pred_len=config['pred_len'], clip=config['clip'],
|
||||
T=config['T'], top_k=config['top_k'], top_p=config['top_p'], sample_count=config['sample_count']
|
||||
)
|
||||
|
||||
# The 'close' price is at index 3 in `feature_list`
|
||||
last_day_close = x[:, -1, 3].numpy()
|
||||
signals = {
|
||||
'last': preds[:, -1, 3] - last_day_close,
|
||||
'mean': np.mean(preds[:, :, 3], axis=1) - last_day_close,
|
||||
'max': np.max(preds[:, :, 3], axis=1) - last_day_close,
|
||||
'min': np.min(preds[:, :, 3], axis=1) - last_day_close,
|
||||
}
|
||||
|
||||
for i in range(len(symbols)):
|
||||
for sig_type, sig_values in signals.items():
|
||||
results[sig_type].append((timestamps[i], symbols[i], sig_values[i]))
|
||||
|
||||
print("Post-processing predictions into DataFrames...")
|
||||
prediction_dfs = {}
|
||||
for sig_type, records in results.items():
|
||||
df = pd.DataFrame(records, columns=['datetime', 'instrument', 'score'])
|
||||
pivot_df = df.pivot_table(index='datetime', columns='instrument', values='score')
|
||||
prediction_dfs[sig_type] = pivot_df.sort_index()
|
||||
|
||||
return prediction_dfs
|
||||
|
||||
|
||||
# =================================================================================
|
||||
# 4. Main Execution
|
||||
# =================================================================================
|
||||
|
||||
def main():
|
||||
"""Main function to set up config, run inference, and execute backtesting."""
|
||||
parser = argparse.ArgumentParser(description="Run Kronos Inference and Backtesting")
|
||||
parser.add_argument("--device", type=str, default="cuda:1", help="Device for inference (e.g., 'cuda:0', 'cpu')")
|
||||
args = parser.parse_args()
|
||||
|
||||
# --- 1. Configuration Setup ---
|
||||
base_config = Config()
|
||||
|
||||
# Create a dedicated dictionary for this run's configuration
|
||||
run_config = {
|
||||
'device': args.device,
|
||||
'data_path': base_config.dataset_path,
|
||||
'result_save_path': base_config.backtest_result_path,
|
||||
'result_name': base_config.backtest_save_folder_name,
|
||||
'tokenizer_path': base_config.finetuned_tokenizer_path,
|
||||
'model_path': base_config.finetuned_predictor_path,
|
||||
'max_context': base_config.max_context,
|
||||
'pred_len': base_config.predict_window,
|
||||
'clip': base_config.clip,
|
||||
'T': base_config.inference_T,
|
||||
'top_k': base_config.inference_top_k,
|
||||
'top_p': base_config.inference_top_p,
|
||||
'sample_count': base_config.inference_sample_count,
|
||||
'batch_size': base_config.backtest_batch_size,
|
||||
}
|
||||
|
||||
print("--- Running with Configuration ---")
|
||||
for key, val in run_config.items():
|
||||
print(f"{key:>20}: {val}")
|
||||
print("-" * 35)
|
||||
|
||||
# --- 2. Load Data ---
|
||||
test_data_path = os.path.join(run_config['data_path'], "test_data.pkl")
|
||||
print(f"Loading test data from {test_data_path}...")
|
||||
with open(test_data_path, 'rb') as f:
|
||||
test_data = pickle.load(f)
|
||||
print(test_data)
|
||||
# --- 3. Generate Predictions ---
|
||||
model_preds = generate_predictions(run_config, test_data)
|
||||
|
||||
# --- 4. Save Predictions ---
|
||||
save_dir = os.path.join(run_config['result_save_path'], run_config['result_name'])
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
predictions_file = os.path.join(save_dir, "predictions.pkl")
|
||||
print(f"Saving prediction signals to {predictions_file}...")
|
||||
with open(predictions_file, 'wb') as f:
|
||||
pickle.dump(model_preds, f)
|
||||
|
||||
# --- 5. Run Backtesting ---
|
||||
with open(predictions_file, 'rb') as f:
|
||||
model_preds = pickle.load(f)
|
||||
|
||||
backtester = QlibBacktest(base_config)
|
||||
backtester.run_and_plot_results(model_preds)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,244 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
from time import gmtime, strftime
|
||||
import torch.distributed as dist
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
|
||||
import comet_ml
|
||||
|
||||
# Ensure project root is in path
|
||||
sys.path.append('../')
|
||||
from config import Config
|
||||
from dataset import QlibDataset
|
||||
from model.kronos import KronosTokenizer, Kronos
|
||||
# Import shared utilities
|
||||
from utils.training_utils import (
|
||||
setup_ddp,
|
||||
cleanup_ddp,
|
||||
set_seed,
|
||||
get_model_size,
|
||||
format_time
|
||||
)
|
||||
|
||||
|
||||
def create_dataloaders(config: dict, rank: int, world_size: int):
|
||||
"""
|
||||
Creates and returns distributed dataloaders for training and validation.
|
||||
|
||||
Args:
|
||||
config (dict): A dictionary of configuration parameters.
|
||||
rank (int): The global rank of the current process.
|
||||
world_size (int): The total number of processes.
|
||||
|
||||
Returns:
|
||||
tuple: (train_loader, val_loader, train_dataset, valid_dataset).
|
||||
"""
|
||||
print(f"[Rank {rank}] Creating distributed dataloaders...")
|
||||
train_dataset = QlibDataset('train')
|
||||
valid_dataset = QlibDataset('val')
|
||||
print(f"[Rank {rank}] Train dataset size: {len(train_dataset)}, Validation dataset size: {len(valid_dataset)}")
|
||||
|
||||
train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
|
||||
val_sampler = DistributedSampler(valid_dataset, num_replicas=world_size, rank=rank, shuffle=False)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset, batch_size=config['batch_size'], sampler=train_sampler,
|
||||
num_workers=config.get('num_workers', 2), pin_memory=True, drop_last=True
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
valid_dataset, batch_size=config['batch_size'], sampler=val_sampler,
|
||||
num_workers=config.get('num_workers', 2), pin_memory=True, drop_last=False
|
||||
)
|
||||
return train_loader, val_loader, train_dataset, valid_dataset
|
||||
|
||||
|
||||
def train_model(model, tokenizer, device, config, save_dir, logger, rank, world_size):
|
||||
"""
|
||||
The main training and validation loop for the predictor.
|
||||
"""
|
||||
start_time = time.time()
|
||||
if rank == 0:
|
||||
effective_bs = config['batch_size'] * world_size
|
||||
print(f"Effective BATCHSIZE per GPU: {config['batch_size']}, Total: {effective_bs}")
|
||||
|
||||
train_loader, val_loader, train_dataset, valid_dataset = create_dataloaders(config, rank, world_size)
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
model.parameters(),
|
||||
lr=config['predictor_learning_rate'],
|
||||
betas=(config['adam_beta1'], config['adam_beta2']),
|
||||
weight_decay=config['adam_weight_decay']
|
||||
)
|
||||
scheduler = torch.optim.lr_scheduler.OneCycleLR(
|
||||
optimizer, max_lr=config['predictor_learning_rate'],
|
||||
steps_per_epoch=len(train_loader), epochs=config['epochs'],
|
||||
pct_start=0.03, div_factor=10
|
||||
)
|
||||
|
||||
best_val_loss = float('inf')
|
||||
dt_result = {}
|
||||
batch_idx_global = 0
|
||||
|
||||
for epoch_idx in range(config['epochs']):
|
||||
epoch_start_time = time.time()
|
||||
model.train()
|
||||
train_loader.sampler.set_epoch(epoch_idx)
|
||||
|
||||
train_dataset.set_epoch_seed(epoch_idx * 10000 + rank)
|
||||
valid_dataset.set_epoch_seed(0)
|
||||
|
||||
for i, (batch_x, batch_x_stamp) in enumerate(train_loader):
|
||||
batch_x = batch_x.squeeze(0).to(device, non_blocking=True)
|
||||
batch_x_stamp = batch_x_stamp.squeeze(0).to(device, non_blocking=True)
|
||||
|
||||
# Tokenize input data on-the-fly
|
||||
with torch.no_grad():
|
||||
token_seq_0, token_seq_1 = tokenizer.encode(batch_x, half=True)
|
||||
|
||||
# Prepare inputs and targets for the language model
|
||||
token_in = [token_seq_0[:, :-1], token_seq_1[:, :-1]]
|
||||
token_out = [token_seq_0[:, 1:], token_seq_1[:, 1:]]
|
||||
|
||||
# Forward pass and loss calculation
|
||||
logits = model(token_in[0], token_in[1], batch_x_stamp[:, :-1, :])
|
||||
loss, s1_loss, s2_loss = model.module.head.compute_loss(logits[0], logits[1], token_out[0], token_out[1])
|
||||
|
||||
# Backward pass and optimization
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=3.0)
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
|
||||
# Logging (Master Process Only)
|
||||
if rank == 0 and (batch_idx_global + 1) % config['log_interval'] == 0:
|
||||
lr = optimizer.param_groups[0]['lr']
|
||||
print(
|
||||
f"[Rank {rank}, Epoch {epoch_idx + 1}/{config['epochs']}, Step {i + 1}/{len(train_loader)}] "
|
||||
f"LR {lr:.6f}, Loss: {loss.item():.4f}"
|
||||
)
|
||||
if rank == 0 and logger:
|
||||
lr = optimizer.param_groups[0]['lr']
|
||||
logger.log_metric('train_predictor_loss_batch', loss.item(), step=batch_idx_global)
|
||||
logger.log_metric('train_S1_loss_each_batch', s1_loss.item(), step=batch_idx_global)
|
||||
logger.log_metric('train_S2_loss_each_batch', s2_loss.item(), step=batch_idx_global)
|
||||
logger.log_metric('predictor_learning_rate', lr, step=batch_idx_global)
|
||||
|
||||
batch_idx_global += 1
|
||||
|
||||
# --- Validation Loop ---
|
||||
model.eval()
|
||||
tot_val_loss_sum_rank = 0.0
|
||||
val_batches_processed_rank = 0
|
||||
with torch.no_grad():
|
||||
for batch_x, batch_x_stamp in val_loader:
|
||||
batch_x = batch_x.squeeze(0).to(device, non_blocking=True)
|
||||
batch_x_stamp = batch_x_stamp.squeeze(0).to(device, non_blocking=True)
|
||||
|
||||
token_seq_0, token_seq_1 = tokenizer.encode(batch_x, half=True)
|
||||
token_in = [token_seq_0[:, :-1], token_seq_1[:, :-1]]
|
||||
token_out = [token_seq_0[:, 1:], token_seq_1[:, 1:]]
|
||||
|
||||
logits = model(token_in[0], token_in[1], batch_x_stamp[:, :-1, :])
|
||||
val_loss, _, _ = model.module.head.compute_loss(logits[0], logits[1], token_out[0], token_out[1])
|
||||
|
||||
tot_val_loss_sum_rank += val_loss.item()
|
||||
val_batches_processed_rank += 1
|
||||
|
||||
# Reduce validation metrics
|
||||
val_loss_sum_tensor = torch.tensor(tot_val_loss_sum_rank, device=device)
|
||||
val_batches_tensor = torch.tensor(val_batches_processed_rank, device=device)
|
||||
dist.all_reduce(val_loss_sum_tensor, op=dist.ReduceOp.SUM)
|
||||
dist.all_reduce(val_batches_tensor, op=dist.ReduceOp.SUM)
|
||||
|
||||
avg_val_loss = val_loss_sum_tensor.item() / val_batches_tensor.item() if val_batches_tensor.item() > 0 else 0
|
||||
|
||||
# --- End of Epoch Summary & Checkpointing (Master Process Only) ---
|
||||
if rank == 0:
|
||||
print(f"\n--- Epoch {epoch_idx + 1}/{config['epochs']} Summary ---")
|
||||
print(f"Validation Loss: {avg_val_loss:.4f}")
|
||||
print(f"Time This Epoch: {format_time(time.time() - epoch_start_time)}")
|
||||
print(f"Total Time Elapsed: {format_time(time.time() - start_time)}\n")
|
||||
if logger:
|
||||
logger.log_metric('val_predictor_loss_epoch', avg_val_loss, epoch=epoch_idx)
|
||||
|
||||
if avg_val_loss < best_val_loss:
|
||||
best_val_loss = avg_val_loss
|
||||
save_path = f"{save_dir}/checkpoints/best_model"
|
||||
model.module.save_pretrained(save_path)
|
||||
print(f"Best model saved to {save_path} (Val Loss: {best_val_loss:.4f})")
|
||||
|
||||
dist.barrier()
|
||||
|
||||
dt_result['best_val_loss'] = best_val_loss
|
||||
return dt_result
|
||||
|
||||
|
||||
def main(config: dict):
|
||||
"""Main function to orchestrate the DDP training process."""
|
||||
rank, world_size, local_rank = setup_ddp()
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
set_seed(config['seed'], rank)
|
||||
|
||||
save_dir = os.path.join(config['save_path'], config['predictor_save_folder_name'])
|
||||
|
||||
# Logger and summary setup (master process only)
|
||||
comet_logger, master_summary = None, {}
|
||||
if rank == 0:
|
||||
os.makedirs(os.path.join(save_dir, 'checkpoints'), exist_ok=True)
|
||||
master_summary = {
|
||||
'start_time': strftime("%Y-%m-%dT%H-%M-%S", gmtime()),
|
||||
'save_directory': save_dir,
|
||||
'world_size': world_size,
|
||||
}
|
||||
if config['use_comet']:
|
||||
comet_logger = comet_ml.Experiment(
|
||||
api_key=config['comet_config']['api_key'],
|
||||
project_name=config['comet_config']['project_name'],
|
||||
workspace=config['comet_config']['workspace'],
|
||||
)
|
||||
comet_logger.add_tag(config['comet_tag'])
|
||||
comet_logger.set_name(config['comet_name'])
|
||||
comet_logger.log_parameters(config)
|
||||
print("Comet Logger Initialized.")
|
||||
|
||||
dist.barrier()
|
||||
|
||||
# Model Initialization
|
||||
tokenizer = KronosTokenizer.from_pretrained(config['finetuned_tokenizer_path'])
|
||||
tokenizer.eval().to(device)
|
||||
|
||||
model = Kronos.from_pretrained(config['pretrained_predictor_path'])
|
||||
model.to(device)
|
||||
model = DDP(model, device_ids=[local_rank], find_unused_parameters=False)
|
||||
|
||||
if rank == 0:
|
||||
print(f"Predictor Model Size: {get_model_size(model.module)}")
|
||||
|
||||
# Start Training
|
||||
dt_result = train_model(
|
||||
model, tokenizer, device, config, save_dir, comet_logger, rank, world_size
|
||||
)
|
||||
|
||||
if rank == 0:
|
||||
master_summary['final_result'] = dt_result
|
||||
with open(os.path.join(save_dir, 'summary.json'), 'w') as f:
|
||||
json.dump(master_summary, f, indent=4)
|
||||
print('Training finished. Summary file saved.')
|
||||
if comet_logger: comet_logger.end()
|
||||
|
||||
cleanup_ddp()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Usage: torchrun --standalone --nproc_per_node=NUM_GPUS train_predictor.py
|
||||
if "WORLD_SIZE" not in os.environ:
|
||||
raise RuntimeError("This script must be launched with `torchrun`.")
|
||||
|
||||
config_instance = Config()
|
||||
main(config_instance.__dict__)
|
||||
@@ -0,0 +1,281 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
from time import gmtime, strftime
|
||||
import argparse
|
||||
import datetime
|
||||
import torch.distributed as dist
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
|
||||
import comet_ml
|
||||
|
||||
# Ensure project root is in path
|
||||
sys.path.append("../")
|
||||
from config import Config
|
||||
from dataset import QlibDataset
|
||||
from model.kronos import KronosTokenizer
|
||||
# Import shared utilities
|
||||
from utils.training_utils import (
|
||||
setup_ddp,
|
||||
cleanup_ddp,
|
||||
set_seed,
|
||||
get_model_size,
|
||||
format_time,
|
||||
)
|
||||
|
||||
|
||||
def create_dataloaders(config: dict, rank: int, world_size: int):
|
||||
"""
|
||||
Creates and returns distributed dataloaders for training and validation.
|
||||
|
||||
Args:
|
||||
config (dict): A dictionary of configuration parameters.
|
||||
rank (int): The global rank of the current process.
|
||||
world_size (int): The total number of processes.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing (train_loader, val_loader, train_dataset, valid_dataset).
|
||||
"""
|
||||
print(f"[Rank {rank}] Creating distributed dataloaders...")
|
||||
train_dataset = QlibDataset('train')
|
||||
valid_dataset = QlibDataset('val')
|
||||
print(f"[Rank {rank}] Train dataset size: {len(train_dataset)}, Validation dataset size: {len(valid_dataset)}")
|
||||
|
||||
train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
|
||||
val_sampler = DistributedSampler(valid_dataset, num_replicas=world_size, rank=rank, shuffle=False)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=config['batch_size'],
|
||||
sampler=train_sampler,
|
||||
shuffle=False, # Shuffle is handled by the sampler
|
||||
num_workers=config.get('num_workers', 2),
|
||||
pin_memory=True,
|
||||
drop_last=True
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
valid_dataset,
|
||||
batch_size=config['batch_size'],
|
||||
sampler=val_sampler,
|
||||
shuffle=False,
|
||||
num_workers=config.get('num_workers', 2),
|
||||
pin_memory=True,
|
||||
drop_last=False
|
||||
)
|
||||
print(f"[Rank {rank}] Dataloaders created. Train steps/epoch: {len(train_loader)}, Val steps: {len(val_loader)}")
|
||||
return train_loader, val_loader, train_dataset, valid_dataset
|
||||
|
||||
|
||||
def train_model(model, device, config, save_dir, logger, rank, world_size):
|
||||
"""
|
||||
The main training and validation loop for the tokenizer.
|
||||
|
||||
Args:
|
||||
model (DDP): The DDP-wrapped model to train.
|
||||
device (torch.device): The device for the current process.
|
||||
config (dict): Configuration dictionary.
|
||||
save_dir (str): Directory to save checkpoints.
|
||||
logger (comet_ml.Experiment): Comet logger instance.
|
||||
rank (int): Global rank of the process.
|
||||
world_size (int): Total number of processes.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the trained model and a dictionary of results.
|
||||
"""
|
||||
start_time = time.time()
|
||||
if rank == 0:
|
||||
effective_bs = config['batch_size'] * world_size * config['accumulation_steps']
|
||||
print(f"[Rank {rank}] BATCHSIZE (per GPU): {config['batch_size']}")
|
||||
print(f"[Rank {rank}] Effective total batch size: {effective_bs}")
|
||||
|
||||
train_loader, val_loader, train_dataset, valid_dataset = create_dataloaders(config, rank, world_size)
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
model.parameters(),
|
||||
lr=config['tokenizer_learning_rate'],
|
||||
weight_decay=config['adam_weight_decay']
|
||||
)
|
||||
|
||||
scheduler = torch.optim.lr_scheduler.OneCycleLR(
|
||||
optimizer=optimizer,
|
||||
max_lr=config['tokenizer_learning_rate'],
|
||||
steps_per_epoch=len(train_loader),
|
||||
epochs=config['epochs'],
|
||||
pct_start=0.03,
|
||||
div_factor=10
|
||||
)
|
||||
|
||||
best_val_loss = float('inf')
|
||||
dt_result = {}
|
||||
batch_idx_global_train = 0
|
||||
|
||||
for epoch_idx in range(config['epochs']):
|
||||
epoch_start_time = time.time()
|
||||
model.train()
|
||||
train_loader.sampler.set_epoch(epoch_idx)
|
||||
|
||||
# Set dataset seeds for reproducible sampling
|
||||
train_dataset.set_epoch_seed(epoch_idx * 10000 + rank)
|
||||
valid_dataset.set_epoch_seed(0) # Keep validation sampling consistent
|
||||
|
||||
for i, (ori_batch_x, _) in enumerate(train_loader):
|
||||
ori_batch_x = ori_batch_x.squeeze(0).to(device, non_blocking=True)
|
||||
|
||||
# --- Gradient Accumulation Loop ---
|
||||
current_batch_total_loss = 0.0
|
||||
for j in range(config['accumulation_steps']):
|
||||
start_idx = j * (ori_batch_x.shape[0] // config['accumulation_steps'])
|
||||
end_idx = (j + 1) * (ori_batch_x.shape[0] // config['accumulation_steps'])
|
||||
batch_x = ori_batch_x[start_idx:end_idx]
|
||||
|
||||
# Forward pass
|
||||
zs, bsq_loss, _, _ = model(batch_x)
|
||||
z_pre, z = zs
|
||||
|
||||
# Loss calculation
|
||||
recon_loss_pre = F.mse_loss(z_pre, batch_x)
|
||||
recon_loss_all = F.mse_loss(z, batch_x)
|
||||
recon_loss = recon_loss_pre + recon_loss_all
|
||||
loss = (recon_loss + bsq_loss) / 2 # Assuming w_1=w_2=1
|
||||
|
||||
loss_scaled = loss / config['accumulation_steps']
|
||||
current_batch_total_loss += loss.item()
|
||||
loss_scaled.backward()
|
||||
|
||||
# --- Optimizer Step after Accumulation ---
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=2.0)
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# --- Logging (Master Process Only) ---
|
||||
if rank == 0 and (batch_idx_global_train + 1) % config['log_interval'] == 0:
|
||||
avg_loss = current_batch_total_loss / config['accumulation_steps']
|
||||
print(
|
||||
f"[Rank {rank}, Epoch {epoch_idx + 1}/{config['epochs']}, Step {i + 1}/{len(train_loader)}] "
|
||||
f"LR {optimizer.param_groups[0]['lr']:.6f}, Loss: {avg_loss:.4f}"
|
||||
)
|
||||
if rank == 0 and logger:
|
||||
avg_loss = current_batch_total_loss / config['accumulation_steps']
|
||||
logger.log_metric('train_tokenizer_loss_batch', avg_loss, step=batch_idx_global_train)
|
||||
logger.log_metric(f'train_vqvae_vq_loss_each_batch', bsq_loss.item(), step=batch_idx_global_train)
|
||||
logger.log_metric(f'train_recon_loss_pre_each_batch', recon_loss_pre.item(), step=batch_idx_global_train)
|
||||
logger.log_metric(f'train_recon_loss_each_batch', recon_loss_all.item(), step=batch_idx_global_train)
|
||||
logger.log_metric('tokenizer_learning_rate', optimizer.param_groups[0]["lr"], step=batch_idx_global_train)
|
||||
|
||||
batch_idx_global_train += 1
|
||||
|
||||
# --- Validation Loop ---
|
||||
model.eval()
|
||||
tot_val_loss_sum_rank = 0.0
|
||||
val_sample_count_rank = 0
|
||||
with torch.no_grad():
|
||||
for ori_batch_x, _ in val_loader:
|
||||
ori_batch_x = ori_batch_x.squeeze(0).to(device, non_blocking=True)
|
||||
zs, _, _, _ = model(ori_batch_x)
|
||||
_, z = zs
|
||||
val_loss_item = F.mse_loss(z, ori_batch_x)
|
||||
|
||||
tot_val_loss_sum_rank += val_loss_item.item() * ori_batch_x.size(0)
|
||||
val_sample_count_rank += ori_batch_x.size(0)
|
||||
|
||||
# Reduce validation losses from all processes
|
||||
val_loss_sum_tensor = torch.tensor(tot_val_loss_sum_rank, device=device)
|
||||
val_count_tensor = torch.tensor(val_sample_count_rank, device=device)
|
||||
dist.all_reduce(val_loss_sum_tensor, op=dist.ReduceOp.SUM)
|
||||
dist.all_reduce(val_count_tensor, op=dist.ReduceOp.SUM)
|
||||
|
||||
avg_val_loss = val_loss_sum_tensor.item() / val_count_tensor.item() if val_count_tensor.item() > 0 else 0
|
||||
|
||||
# --- End of Epoch Summary & Checkpointing (Master Process Only) ---
|
||||
if rank == 0:
|
||||
print(f"\n--- Epoch {epoch_idx + 1}/{config['epochs']} Summary ---")
|
||||
print(f"Validation Loss: {avg_val_loss:.4f}")
|
||||
print(f"Time This Epoch: {format_time(time.time() - epoch_start_time)}")
|
||||
print(f"Total Time Elapsed: {format_time(time.time() - start_time)}\n")
|
||||
if logger:
|
||||
logger.log_metric('val_tokenizer_loss_epoch', avg_val_loss, epoch=epoch_idx)
|
||||
|
||||
if avg_val_loss < best_val_loss:
|
||||
best_val_loss = avg_val_loss
|
||||
save_path = f"{save_dir}/checkpoints/best_model"
|
||||
model.module.save_pretrained(save_path)
|
||||
print(f"Best model saved to {save_path} (Val Loss: {best_val_loss:.4f})")
|
||||
if logger:
|
||||
logger.log_model("best_model", save_path)
|
||||
|
||||
dist.barrier() # Ensure all processes finish the epoch before starting the next one.
|
||||
|
||||
dt_result['best_val_loss'] = best_val_loss
|
||||
return model, dt_result
|
||||
|
||||
|
||||
def main(config: dict):
|
||||
"""
|
||||
Main function to orchestrate the DDP training process.
|
||||
"""
|
||||
rank, world_size, local_rank = setup_ddp()
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
set_seed(config['seed'], rank)
|
||||
|
||||
save_dir = os.path.join(config['save_path'], config['tokenizer_save_folder_name'])
|
||||
|
||||
# Logger and summary setup (master process only)
|
||||
comet_logger, master_summary = None, {}
|
||||
if rank == 0:
|
||||
os.makedirs(os.path.join(save_dir, 'checkpoints'), exist_ok=True)
|
||||
master_summary = {
|
||||
'start_time': strftime("%Y-%m-%dT%H-%M-%S", gmtime()),
|
||||
'save_directory': save_dir,
|
||||
'world_size': world_size,
|
||||
}
|
||||
if config['use_comet']:
|
||||
comet_logger = comet_ml.Experiment(
|
||||
api_key=config['comet_config']['api_key'],
|
||||
project_name=config['comet_config']['project_name'],
|
||||
workspace=config['comet_config']['workspace'],
|
||||
)
|
||||
comet_logger.add_tag(config['comet_tag'])
|
||||
comet_logger.set_name(config['comet_name'])
|
||||
comet_logger.log_parameters(config)
|
||||
print("Comet Logger Initialized.")
|
||||
|
||||
dist.barrier() # Ensure save directory is created before proceeding
|
||||
|
||||
# Model Initialization
|
||||
model = KronosTokenizer.from_pretrained(config['pretrained_tokenizer_path'])
|
||||
model.to(device)
|
||||
model = DDP(model, device_ids=[local_rank], find_unused_parameters=False)
|
||||
|
||||
if rank == 0:
|
||||
print(f"Model Size: {get_model_size(model.module)}")
|
||||
|
||||
# Start Training
|
||||
_, dt_result = train_model(
|
||||
model, device, config, save_dir, comet_logger, rank, world_size
|
||||
)
|
||||
|
||||
# Finalize and save summary (master process only)
|
||||
if rank == 0:
|
||||
master_summary['final_result'] = dt_result
|
||||
with open(os.path.join(save_dir, 'summary.json'), 'w') as f:
|
||||
json.dump(master_summary, f, indent=4)
|
||||
print('Training finished. Summary file saved.')
|
||||
if comet_logger:
|
||||
comet_logger.end()
|
||||
|
||||
cleanup_ddp()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Usage: torchrun --standalone --nproc_per_node=NUM_GPUS train_tokenizer.py
|
||||
if "WORLD_SIZE" not in os.environ:
|
||||
raise RuntimeError("This script must be launched with `torchrun`.")
|
||||
|
||||
config_instance = Config()
|
||||
main(config_instance.__dict__)
|
||||
@@ -0,0 +1,118 @@
|
||||
import os
|
||||
import random
|
||||
import datetime
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
|
||||
def setup_ddp():
|
||||
"""
|
||||
Initializes the distributed data parallel environment.
|
||||
|
||||
This function relies on environment variables set by `torchrun` or a similar
|
||||
launcher. It initializes the process group and sets the CUDA device for the
|
||||
current process.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing (rank, world_size, local_rank).
|
||||
"""
|
||||
if not dist.is_available():
|
||||
raise RuntimeError("torch.distributed is not available.")
|
||||
|
||||
dist.init_process_group(backend="nccl")
|
||||
rank = int(os.environ["RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.cuda.set_device(local_rank)
|
||||
print(
|
||||
f"[DDP Setup] Global Rank: {rank}/{world_size}, "
|
||||
f"Local Rank (GPU): {local_rank} on device {torch.cuda.current_device()}"
|
||||
)
|
||||
return rank, world_size, local_rank
|
||||
|
||||
|
||||
def cleanup_ddp():
|
||||
"""Cleans up the distributed process group."""
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def set_seed(seed: int, rank: int = 0):
|
||||
"""
|
||||
Sets the random seed for reproducibility across all relevant libraries.
|
||||
|
||||
Args:
|
||||
seed (int): The base seed value.
|
||||
rank (int): The process rank, used to ensure different processes have
|
||||
different seeds, which can be important for data loading.
|
||||
"""
|
||||
actual_seed = seed + rank
|
||||
random.seed(actual_seed)
|
||||
np.random.seed(actual_seed)
|
||||
torch.manual_seed(actual_seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(actual_seed)
|
||||
# The two lines below can impact performance, so they are often
|
||||
# reserved for final experiments where reproducibility is critical.
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def get_model_size(model: torch.nn.Module) -> str:
|
||||
"""
|
||||
Calculates the number of trainable parameters in a PyTorch model and returns
|
||||
it as a human-readable string.
|
||||
|
||||
Args:
|
||||
model (torch.nn.Module): The PyTorch model.
|
||||
|
||||
Returns:
|
||||
str: A string representing the model size (e.g., "175.0B", "7.1M", "50.5K").
|
||||
"""
|
||||
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
if total_params >= 1e9:
|
||||
return f"{total_params / 1e9:.1f}B" # Billions
|
||||
elif total_params >= 1e6:
|
||||
return f"{total_params / 1e6:.1f}M" # Millions
|
||||
else:
|
||||
return f"{total_params / 1e3:.1f}K" # Thousands
|
||||
|
||||
|
||||
def reduce_tensor(tensor: torch.Tensor, world_size: int, op=dist.ReduceOp.SUM) -> torch.Tensor:
|
||||
"""
|
||||
Reduces a tensor's value across all processes in a distributed setup.
|
||||
|
||||
Args:
|
||||
tensor (torch.Tensor): The tensor to be reduced.
|
||||
world_size (int): The total number of processes.
|
||||
op (dist.ReduceOp, optional): The reduction operation (SUM, AVG, etc.).
|
||||
Defaults to dist.ReduceOp.SUM.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The reduced tensor, which will be identical on all processes.
|
||||
"""
|
||||
rt = tensor.clone()
|
||||
dist.all_reduce(rt, op=op)
|
||||
# Note: `dist.ReduceOp.AVG` is available in newer torch versions.
|
||||
# For compatibility, manual division is sometimes used after a SUM.
|
||||
if op == dist.ReduceOp.AVG:
|
||||
rt /= world_size
|
||||
return rt
|
||||
|
||||
|
||||
def format_time(seconds: float) -> str:
|
||||
"""
|
||||
Formats a duration in seconds into a human-readable H:M:S string.
|
||||
|
||||
Args:
|
||||
seconds (float): The total seconds.
|
||||
|
||||
Returns:
|
||||
str: The formatted time string (e.g., "0:15:32").
|
||||
"""
|
||||
return str(datetime.timedelta(seconds=int(seconds)))
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user