Reward/robometer/train.py

import json
import os
from dataclasses import asdict
import shutil

import torch
import torch.distributed as dist
import yaml
from rich import print as rprint
from rich.panel import Panel
from omegaconf import DictConfig
from hydra.core.config_store import ConfigStore
from hydra import main as hydra_main

from peft import prepare_model_for_kbit_training
from robometer.configs.experiment_configs import (
    ExperimentConfig,
    ModelConfig,
    PEFTConfig,
    DataConfig,
    TrainingConfig,
    LossConfig,
    LoggingConfig,
    SaveBestConfig,
    CustomEvaluationConfig,
)
from robometer.trainers import ReWiNDTrainer, RBMHeadsTrainer
from robometer.data.datasets.helpers import show_available_datasets
from robometer.utils.distributed import is_rank_0
from robometer.utils.logger import rank_0_info
from robometer.utils.timer import _timer
from robometer.utils.save import SaveBestCallback, resolve_checkpoint_path, update_cfg_with_pretrained_ckpt
from robometer.utils.setup_utils import (
    create_training_arguments,
    setup_batch_collator,
    setup_dataset,
    setup_model_and_processor,
    setup_peft_model,
)
from robometer.data.datasets.base import resolve_dataset_keys
from robometer.utils.logger import Logger
from robometer.utils.distributed import banner
from robometer.utils.config_utils import display_config, convert_hydra_to_dataclass
import datasets

datasets.logging.set_verbosity_error()
os.environ["TOKENIZERS_PARALLELISM"] = "false"
torch.autograd.set_detect_anomaly(True)

# Register structured configs with Hydra
cs = ConfigStore.instance()
cs.store(name="base_config", node=ExperimentConfig)
cs.store(group="model", name="model_config", node=ModelConfig)
cs.store(group="peft", name="peft_config", node=PEFTConfig)
cs.store(group="data", name="data_config", node=DataConfig)
cs.store(group="training", name="training_config", node=TrainingConfig)
cs.store(group="loss", name="loss_config", node=LossConfig)
cs.store(group="logging", name="logging_config", node=LoggingConfig)
cs.store(group="logging/save_best", name="save_best_config", node=SaveBestConfig)
cs.store(group="custom_eval", name="custom_eval_config", node=CustomEvaluationConfig)


import torch

torch.set_num_threads(64)
torch.set_num_interop_threads(8)


def train(cfg: ExperimentConfig):
    timing_raw = {}

    run_name = cfg.training.exp_name
    if cfg.debug:
        run_name += "_debug"
        cfg.training.logging_steps = 1
        cfg.training.eval_steps = 5
        # cfg.data.eval_subset_size = 100
        cfg.training.custom_eval_steps = 5
        cfg.logging.save_best.save_every = 5
        cfg.data.dataloader_num_workers = 0
        cfg.data.dataloader_persistent_workers = False

        # cfg.custom_eval.num_examples_per_quality_pr = 1
        # cfg.custom_eval.policy_ranking_max_tasks = 10

    # Set memory management
    torch.backends.cudnn.benchmark = True
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    checkpoint_to_load = cfg.training.load_from_checkpoint or cfg.training.resume_from_checkpoint
    if checkpoint_to_load:
        rank_0_info(f"Loading model from checkpoint: {checkpoint_to_load}")
    update_cfg_with_pretrained_ckpt(cfg, checkpoint_to_load)

    banner("Setting up model and processor")
    with _timer("time/setup_model_and_processor", timing_raw=timing_raw):
        tokenizer, processor, rbm_model = setup_model_and_processor(
            cfg.model,
            hf_model_id=checkpoint_to_load or "",
            peft_config=cfg.peft,
        )

    # Apply PEFT if enabled
    if cfg.model.use_peft:
        peft_rbm_model = setup_peft_model(rbm_model, cfg.peft)
    else:
        peft_rbm_model = rbm_model
        rank_0_info("PEFT not enabled, using full model")

    if cfg.model.quantization:
        peft_rbm_model = prepare_model_for_kbit_training(peft_rbm_model)

    output_dir = os.path.join(cfg.training.output_dir, run_name)

    training_args = create_training_arguments(cfg.training, output_dir)

    # Handle output directory existence (works with accelerate/distributed training)
    overwrite_output_dir = getattr(cfg.training, "overwrite_output_dir", False)

    # Check if distributed training is initialized (for proper synchronization)
    # This is important for accelerate/FSDP setups where multiple processes run
    dist_initialized = dist.is_available() and dist.is_initialized()

    # Check if output directory exists (only on rank 0 to avoid race conditions)
    if is_rank_0() and os.path.exists(output_dir):
        if overwrite_output_dir:
            rank_0_info(f"Output directory {output_dir} already exists. Overwriting (overwrite_output_dir=True)...")
            shutil.rmtree(output_dir)
        else:
            raise ValueError(
                f"Output directory {output_dir} already exists. "
                f"Set overwrite_output_dir=True in config to overwrite it, or use a different output directory."
            )

    # Synchronize all processes before creating directory (important for distributed training)
    # This ensures rank 0 finishes checking/removing before other processes try to create it
    if dist_initialized:
        dist.barrier()

    banner("Creating output directory", f"Logging to: {output_dir}")
    # Create output directory (all processes need to do this for distributed training)
    # os.makedirs is safe to call multiple times (exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    # Synchronize after directory creation to ensure all processes see it
    if dist_initialized:
        dist.barrier()

    # Initialize logger (works with wandb/tensorboard)
    log_to = cfg.logging.log_to
    log_level = cfg.logging.log_level
    logger = Logger(log_to=log_to, output_dir=output_dir, is_main_process=is_rank_0(), log_level=log_level)
    config_save_path = os.path.join(output_dir, "config.yaml")
    config_dict = asdict(cfg)
    with open(config_save_path, "w") as f:
        yaml.dump(config_dict, f, default_flow_style=False, indent=2)
    rank_0_info(f"Saved training config to: {config_save_path}")

    # Try to load existing wandb info if resuming training
    wandb_info_path = os.path.join(output_dir, "wandb_info.json")
    resume_id = None
    if os.path.exists(wandb_info_path):
        try:
            with open(wandb_info_path) as f:
                wandb_info = json.load(f)
            resume_id = wandb_info.get("wandb_id")
            if resume_id:
                rank_0_info(f"Found existing wandb run ID: {resume_id}, will resume run")
        except Exception as e:
            rank_0_info(f"Could not load wandb info: {e}")

    # Initialize wandb via logger if requested
    if "wandb" in (cfg.logging.log_to or []) and is_rank_0():
        # Convert config to dict for wandb using dataclass asdict
        config_dict = asdict(cfg)
        logger.init_wandb(
            project=cfg.logging.wandb_project,
            entity=cfg.logging.wandb_entity,
            name=run_name,
            config=config_dict,
            notes=cfg.logging.wandb_notes,
            mode=cfg.logging.wandb_mode,
            resume_id=resume_id,
        )
        if resume_id:
            rank_0_info(f"Wandb resumed run: {run_name} (ID: {resume_id})")
        else:
            rank_0_info(f"Wandb initialized: {run_name}")
        if cfg.logging.wandb_notes:
            rank_0_info(f"Wandb notes: {cfg.logging.wandb_notes}")

    logger.write_wandb_info(output_dir, run_name)

    # Use the shared utilities for batch collator and dataset

    if is_rank_0():
        show_available_datasets()

    banner("Resolving dataset keys")
    cfg.data.train_datasets = resolve_dataset_keys(cfg.data.train_datasets, split="train")
    rank_0_info(f"Resolved train datasets: {cfg.data.train_datasets}")

    if cfg.data.eval_datasets:
        cfg.data.eval_datasets = resolve_dataset_keys(cfg.data.eval_datasets, split="eval")
        rank_0_info(f"Resolved eval datasets: {cfg.data.eval_datasets}")

    # Resolve custom evaluation dataset keys once (replace in place)
    for eval_type in cfg.custom_eval.eval_types:
        datasets = getattr(cfg.custom_eval, eval_type, None)
        if datasets:
            resolved = resolve_dataset_keys(datasets, split="eval")
            setattr(cfg.custom_eval, eval_type, resolved)
            rank_0_info(f"Resolved {eval_type} datasets: {resolved}")

    rank_0_info("Dataset keys resolved")

    banner("Setting up training and evaluation datasets and collator")
    with _timer("time/setup_data", timing_raw=timing_raw):
        batch_collator = setup_batch_collator(processor, tokenizer, cfg, is_eval=False)
        train_dataset = setup_dataset(cfg.data)
        num_train_samples = len(train_dataset)
        rank_0_info(f"Training dataset created with {num_train_samples} samples")
        rank_0_info(f"=" * 100)

    # Set up evaluation dataset if evaluation is enabled
    eval_dataset = None
    if cfg.training.do_eval:
        if cfg.data.eval_subset_size is not None:
            dataset_kwargs = {"max_samples": cfg.data.eval_subset_size}
        else:
            dataset_kwargs = {}

        eval_dataset = setup_dataset(cfg.data, is_eval=True, **dataset_kwargs)
        num_eval_samples = len(eval_dataset)
        rank_0_info(f"Evaluation dataset created with {num_eval_samples} samples")

    banner("Setting up trainer", f"Trainer class: {cfg.trainer_cls}")
    trainer_cls = {
        "rbm_heads": RBMHeadsTrainer,
        "rewind_transformer": ReWiNDTrainer,
        "rewind_scale_transformer": ReWiNDTrainer,
    }[cfg.trainer_cls]

    # Add SaveBestCallback to automatically save and upload best models
    save_best_cfg = cfg.logging.save_best
    save_callback = SaveBestCallback(
        **asdict(save_best_cfg),
        base_model=cfg.model.base_model_id,
    )

    trainer = trainer_cls(
        model=peft_rbm_model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=batch_collator,
        config=cfg,
        logger=logger,
        callbacks=[save_callback],
    )

    # Set trainer reference in the callback so it can access trainer methods
    save_callback.setup_trainer_reference(trainer)

    # Debug: Check if callback was added
    rank_0_info(f"🔧 DEBUG: Trainer callbacks: {[type(cb).__name__ for cb in trainer.callback_handler.callbacks]}")

    metrics_info = []
    for name, is_better in zip(save_best_cfg.metric_names, save_best_cfg.greater_is_better):
        direction = "↗️ higher" if is_better else "↘️ lower"
        metrics_info.append(f"{name} ({direction})")

    rank_0_info(f"💾 SaveBest monitoring: {', '.join(metrics_info)}")
    rank_0_info(f"📁 Keeping top {save_best_cfg.keep_top_k} checkpoint(s) and upload(s)")

    if is_rank_0():
        print("\n" + "=" * 80)
        print("--- PRE-TRAINING FSDP DIAGNOSTICS ---")
        # The Trainer creates its own Accelerator instance. Let's check its state.
        if hasattr(trainer, "accelerator"):
            print("Trainer's Accelerator object found.")
            fsdp_plugin = getattr(trainer.accelerator.state, "fsdp_plugin", None)
            if fsdp_plugin:
                print("FSDP Plugin found in Accelerator state.")
                # This is the configuration the accelerator will ACTUALLY use for wrapping.
                print(f"VERIFY: Actual FSDP plugin config being used: {fsdp_plugin}")
            else:
                print("ERROR: FSDP Plugin NOT found in the Trainer's accelerator state!")
        else:
            print("ERROR: Trainer has no 'accelerator' attribute yet. This check needs to be later.")
        print("=" * 80 + "\n")

    # log timing_raw via logger
    if is_rank_0():
        logger.log_scalars(timing_raw)

    rank_0_info(f"Timing raw: {timing_raw}")

    # Full resume: restore optimizer state and step counter (load_from_checkpoint only loads weights at setup)
    hub_token = (save_best_cfg.hub_token if save_best_cfg else None) or os.environ.get("HF_TOKEN")
    resume_path = (
        resolve_checkpoint_path(cfg.training.resume_from_checkpoint, hub_token=hub_token)
        if cfg.training.resume_from_checkpoint
        else None
    )
    if resume_path:
        rank_0_info(f"Resuming training from checkpoint: {resume_path}")
    else:
        rank_0_info("Training from step 0 (no resume)")

    # Restore random state from checkpoint only when doing full resume
    if resume_path and os.path.isdir(resume_path):
        random_state_file = os.path.join(resume_path, "dataset_random_state.json")
        if os.path.exists(random_state_file):
            try:
                with open(random_state_file, "r") as f:
                    random_state = json.load(f)
                # Handle RepeatedDataset wrapper if present
                train_dataset = train_dataset.dataset if hasattr(train_dataset, "dataset") else train_dataset
                if hasattr(train_dataset, "set_random_state"):
                    train_dataset.set_random_state(random_state)
                    rank_0_info(f"Restored dataset random state from {random_state_file}")
                else:
                    rank_0_info(f"Dataset does not support random state restoration")
            except Exception as e:
                rank_0_info(f"Could not restore random state: {e}")
        else:
            rank_0_info(f"No dataset_random_state.json found in checkpoint, starting with fresh random state")

    if cfg.debug:
        rank_0_info("🐛 DEBUG MODE: eval_steps=2, custom_eval_steps=2, eval_subset_size=10")

    trainer.train(resume_from_checkpoint=resume_path)
    trainer.save_model(cfg.training.output_dir)
    rank_0_info(f"Training complete! Check {cfg.training.output_dir} for checkpoints and final model.")


@hydra_main(version_base=None, config_path="robometer/configs", config_name="config")
def main(cfg: DictConfig):
    banner("Starting Robometer Training")

    # Convert Hydra config to dataclass
    exp_cfg = convert_hydra_to_dataclass(cfg, ExperimentConfig)

    # Display the configuration in a nice Rich format
    display_config(exp_cfg)

    if exp_cfg.mode == "train":
        if is_rank_0():
            rprint(Panel.fit("🚀 Starting Robometer Training", style="bold green"))
        train(exp_cfg)
    else:
        raise ValueError(f"Unknown mode: {exp_cfg.mode}. Must be 'train' or 'evaluate'")


if __name__ == "__main__":
    main()