106 lines
3.0 KiB
Python
106 lines
3.0 KiB
Python
|
import os
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import yaml
|
||
|
|
import wandb
|
||
|
|
from argparse import ArgumentParser
|
||
|
|
from accelerate import Accelerator, DistributedDataParallelKwargs, DataLoaderConfiguration
|
||
|
|
|
||
|
|
from wall_x.trainer.qwen_vl_act_trainer import QwenVlAct_Trainer
|
||
|
|
|
||
|
|
|
||
|
|
def setup_environment():
|
||
|
|
"""Set up environment variables for training."""
|
||
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
|
|
||
|
|
|
||
|
|
def load_config(config_path):
|
||
|
|
"""Load configuration from YAML file."""
|
||
|
|
with open(config_path, "r") as f:
|
||
|
|
config = yaml.load(f, Loader=yaml.FullLoader)
|
||
|
|
|
||
|
|
# Set model_type in data config if not already set
|
||
|
|
config["data"]["model_type"] = config.get("model_type")
|
||
|
|
|
||
|
|
return config
|
||
|
|
|
||
|
|
|
||
|
|
def setup_accelerator(config):
|
||
|
|
"""Initialize and configure the accelerator for distributed training."""
|
||
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Preparing accelerator")
|
||
|
|
|
||
|
|
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
|
||
|
|
accelerator_dataloader_config = DataLoaderConfiguration(dispatch_batches=False)
|
||
|
|
|
||
|
|
accelerator = Accelerator(
|
||
|
|
kwargs_handlers=[ddp_kwargs],
|
||
|
|
mixed_precision="bf16",
|
||
|
|
dataloader_config=accelerator_dataloader_config,
|
||
|
|
gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1)
|
||
|
|
)
|
||
|
|
|
||
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Accelerator initialization complete")
|
||
|
|
|
||
|
|
return accelerator
|
||
|
|
|
||
|
|
|
||
|
|
def setup_logging(config, accelerator):
|
||
|
|
"""Set up logging with wandb for the main process."""
|
||
|
|
if not accelerator.is_main_process:
|
||
|
|
return None
|
||
|
|
|
||
|
|
# Create save directory if it doesn't exist
|
||
|
|
save_path = config["save_path"]
|
||
|
|
if not os.path.exists(save_path):
|
||
|
|
print(f"Save path {save_path} does not exist, creating directory.")
|
||
|
|
os.makedirs(save_path, exist_ok=True)
|
||
|
|
|
||
|
|
print("Configuration:")
|
||
|
|
print("=" * 50)
|
||
|
|
print(json.dumps(config, indent=2, ensure_ascii=False))
|
||
|
|
print("=" * 50)
|
||
|
|
|
||
|
|
# Initialize wandb logger
|
||
|
|
logger = wandb.init(
|
||
|
|
project=config["log_project"],
|
||
|
|
name=config["log_name"],
|
||
|
|
save_code=False,
|
||
|
|
force=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
return logger
|
||
|
|
|
||
|
|
|
||
|
|
def main(args):
|
||
|
|
"""Main training function."""
|
||
|
|
setup_environment()
|
||
|
|
|
||
|
|
# Load configuration
|
||
|
|
config = load_config(args.config)
|
||
|
|
|
||
|
|
# Set up accelerator
|
||
|
|
accelerator = setup_accelerator(config)
|
||
|
|
|
||
|
|
# Set up logging
|
||
|
|
logger = setup_logging(config, accelerator)
|
||
|
|
|
||
|
|
# Initialize trainer
|
||
|
|
trainer = QwenVlAct_Trainer(
|
||
|
|
config=config,
|
||
|
|
logger=logger,
|
||
|
|
accelerator=accelerator,
|
||
|
|
seed=args.seed,
|
||
|
|
data_config_path=args.config,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Start training
|
||
|
|
trainer.fit()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
parser = ArgumentParser(description="Training script for Wall-X model")
|
||
|
|
parser.add_argument("--config", type=str, required=True, help="Path to configuration YAML file")
|
||
|
|
parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
main(args)
|