train_sft_tsqa.py

#!/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()