#!/usr/bin/env python """Train native ITFormer on the Time-MQA TSQA train split.""" from __future__ import annotations import argparse import json import random from pathlib import Path import numpy as np import torch from accelerate import Accelerator from transformers import AutoTokenizer from dataset.tsqa_dataset import ITFormerTSQACollator, ITFormerTSQADataset from EXP.exp_instruct import Exp_Instruct from models.TimeLanguageModel import TLMConfig from utils.accelerate_compat import patch_accelerate_unwrap_model def str2bool(value): return str(value).strip().lower() in {"1", "true", "yes", "y", "on"} def parse_args(): parser = argparse.ArgumentParser(description="Train ITFormer on Time-MQA TSQA.") parser.add_argument("--train_path", required=True) parser.add_argument("--eval_path", required=True) parser.add_argument( "--llm_model_path", default="/mnt/share01/sqk/models/qwen2.5-7b-instruct", ) parser.add_argument("--load_ts_encoder", required=True) parser.add_argument("--output_dir", required=True) parser.add_argument("--max_train_samples", type=int, default=0) parser.add_argument("--max_eval_samples", type=int, default=0) parser.add_argument("--max_seq_length", type=int, default=4096) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--model", default="TimeSeriesEncoder") parser.add_argument("--d_model", type=int, default=512) parser.add_argument("--n_heads", type=int, default=8) parser.add_argument("--e_layers", type=int, default=4) parser.add_argument("--patch_len", type=int, default=60) parser.add_argument("--stride", type=int, default=60) parser.add_argument("--input_len", type=int, default=600) parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--freeze_ts_model", type=str2bool, default=True) parser.add_argument("--pretrain", type=str2bool, default=False) parser.add_argument("--min_mask_ratio", type=float, default=0.7) parser.add_argument("--max_mask_ratio", type=float, default=0.8) parser.add_argument("--it_d_model", type=int, default=896) parser.add_argument("--it_n_heads", type=int, default=16) parser.add_argument("--it_layers", type=int, default=2) parser.add_argument("--it_dropout", type=float, default=0.1) parser.add_argument("--itformer_legacy_double_residual", action="store_true") parser.add_argument("--prefix_num", type=int, default=25) parser.add_argument("--adapter_type", default="itformer", choices=["itformer"]) parser.add_argument( "--llm_attn_implementation", default=None, choices=["eager", "sdpa", "flash_attention_2"], ) parser.add_argument( "--llm_torch_dtype", default=None, choices=["float16", "bfloat16", "float32"], ) parser.add_argument("--use_lora", type=str2bool, default=True) parser.add_argument("--lora_r", type=int, default=16) parser.add_argument("--lora_alpha", type=int, default=32) parser.add_argument("--lora_dropout", type=float, default=0.05) parser.add_argument( "--lora_target_modules", nargs="+", default=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], ) parser.add_argument("--gradient_checkpointing", type=str2bool, default=True) parser.add_argument("--do_train", type=str2bool, default=True) parser.add_argument("--per_device_train_batch_size", type=int, default=1) parser.add_argument("--per_device_eval_batch_size", type=int, default=1) parser.add_argument("--gradient_accumulation_steps", type=int, default=2) parser.add_argument("--learning_rate", type=float, default=5e-5) parser.add_argument("--max_grad_norm", type=float, default=1.0) parser.add_argument("--num_train_epochs", type=float, default=2.0) parser.add_argument("--weight_decay", type=float, default=1e-6) parser.add_argument("--fp16", type=str2bool, default=False) parser.add_argument("--bf16", action="store_true") parser.add_argument("--dataloader_pin_memory", type=str2bool, default=True) parser.add_argument("--dataloader_num_workers", type=int, default=4) parser.add_argument("--save_steps", type=int, default=100) parser.add_argument("--save_total_limit", type=int, default=2) parser.add_argument("--logging_steps", type=int, default=5) parser.add_argument("--eval_steps", type=int, default=999999999) parser.add_argument("--eval_stragy", default="no") parser.add_argument("--shuffle", type=str2bool, default=True) parser.add_argument("--report_to", default="none") parser.add_argument("--mode", default="train") return parser.parse_args() def main(): patch_accelerate_unwrap_model() args = parse_args() accelerator = Accelerator(device_placement=True) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.seed) tokenizer = AutoTokenizer.from_pretrained(args.llm_model_path) tokenizer.padding_side = "left" if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token if "<|image_pad|>" not in tokenizer.get_vocab(): tokenizer.add_tokens(["<|image_pad|>"]) train_dataset = ITFormerTSQADataset( args.train_path, tokenizer, prefix_num=args.prefix_num, input_len=args.input_len, max_samples=args.max_train_samples, include_answer=True, max_seq_length=args.max_seq_length, ) eval_dataset = ITFormerTSQADataset( args.eval_path, tokenizer, prefix_num=args.prefix_num, input_len=args.input_len, max_samples=args.max_eval_samples, include_answer=False, ) if accelerator.is_main_process: print(f"TSQA train samples: {len(train_dataset)}") print( "TSQA overlength train samples skipped: " f"{len(train_dataset.skipped_overlength)} " f"(max_seq_length={args.max_seq_length})" ) print(f"TSQA eval samples reserved for post-training eval: {len(eval_dataset)}") tlm_config = TLMConfig( llm_model_path=args.llm_model_path, freeze_ts_model=args.freeze_ts_model, ts_pad_num=args.prefix_num, llm_attn_implementation=args.llm_attn_implementation, llm_torch_dtype=args.llm_torch_dtype, use_lora=args.use_lora, lora_r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout, lora_target_modules=args.lora_target_modules, gradient_checkpointing=args.gradient_checkpointing, ) trainer = Exp_Instruct( args, train_dataset=train_dataset, eval_dataset=None, tlm_config=tlm_config, ) trainer.data_collator = ITFormerTSQACollator(tokenizer) if accelerator.is_main_process: output = Path(args.output_dir) output.mkdir(parents=True, exist_ok=True) skip_path = output / "skipped_overlength_train.jsonl" with skip_path.open("w", encoding="utf-8") as handle: for row in train_dataset.skipped_overlength: handle.write(json.dumps(row, ensure_ascii=False) + "\n") print(f"Overlength skip manifest: {skip_path}") if len(tokenizer) != trainer.model.llm_model.config.vocab_size: trainer.model.llm_model.resize_token_embeddings(len(tokenizer)) trainable_params = sum( parameter.numel() for parameter in trainer.model.parameters() if parameter.requires_grad ) total_params = sum(parameter.numel() for parameter in trainer.model.parameters()) if accelerator.is_main_process: print( f"ITFormer + LoRA trainable parameters: {trainable_params:,} / " f"{total_params:,} ({100 * trainable_params / total_params:.4f}%)" ) trainer.train(resume_from_checkpoint=False) trainer.save_model(args.output_dir) if trainer.args.should_save: tokenizer.save_pretrained(args.output_dir) run_config = vars(args).copy() run_config.update( { "actual_train_samples": len(train_dataset), "actual_eval_samples": len(eval_dataset), "skipped_overlength_train_samples": len( train_dataset.skipped_overlength ), } ) output = Path(args.output_dir) output.mkdir(parents=True, exist_ok=True) with (output / "tsqa_run_config.json").open("w", encoding="utf-8") as handle: json.dump(run_config, handle, ensure_ascii=False, indent=2) print(f"Saved final ITFormer checkpoint: {output}") if __name__ == "__main__": main()