train_torchrun.py

import os
import sys
import json
import glob
import argparse
from easydict import EasyDict as edict

import torch
# import torch.multiprocessing as mp # No longer needed for mp.spawn
import numpy as np
import random
# from peft import LoraModel, LoraConfig

from trellis import models, datasets, trainers # Assuming these are your custom modules
from trellis.utils.dist_utils import setup_dist # Assuming this is your custom module


def find_ckpt(cfg):
    # Load checkpoint
    cfg['load_ckpt'] = None
    if cfg.load_dir != '':
        if cfg.ckpt == 'latest':
            files = glob.glob(os.path.join(cfg.load_dir, 'ckpts', 'misc_*.pt'))
            if len(files) != 0:
                cfg.load_ckpt = max([
                    int(os.path.basename(f).split('step')[-1].split('.')[0])
                    for f in files
                ])
        elif cfg.ckpt == 'none':
            cfg.load_ckpt = None
        else:
            cfg.load_ckpt = int(cfg.ckpt)
    return cfg


def setup_rng(rank):
    # Seed RNGs for reproducibility
    # It's good practice to ensure different ranks get different seeds if necessary,
    # but often a global seed offset by rank is used.
    # The original code used rank, which is fine.
    seed = cfg.get('seed', 42) # Get a base seed from config or use a default
    torch.manual_seed(seed + rank)
    torch.cuda.manual_seed_all(seed + rank)
    np.random.seed(seed + rank)
    random.seed(seed + rank)
    # Ensure determinism if desired (can impact performance)
    # torch.backends.cudnn.deterministic = True
    # torch.backends.cudnn.benchmark = False


def get_model_summary(model):
    model_summary = 'Parameters:\n'
    model_summary += '=' * 128 + '\n'
    model_summary += f'{"Name":<{72}}{"Shape":<{32}}{"Type":<{16}}{"Grad"}\n'
    num_params = 0
    num_trainable_params = 0
    for name, param in model.named_parameters():
        model_summary += f'{name:<{72}}{str(param.shape):<{32}}{str(param.dtype):<{16}}{param.requires_grad}\n'
        num_params += param.numel()
        if param.requires_grad:
            num_trainable_params += param.numel()
    model_summary += '\n'
    model_summary += f'Number of parameters: {num_params}\n'
    model_summary += f'Number of trainable parameters: {num_trainable_params}\n'
    return model_summary


def main(cfg): # local_rank is no longer passed as an argument
    # Set up distributed training using environment variables set by torchrun
    local_rank = int(os.environ.get('LOCAL_RANK', 0))
    rank = int(os.environ.get('RANK', 0))
    world_size = int(os.environ.get('WORLD_SIZE', 1))

    if world_size > 1:
        torch.cuda.set_device(local_rank) # Crucial for torchrun
        # setup_dist will use env:// method or an explicitly passed master_addr/port
        # If setup_dist is designed to use 'env://', master_addr and master_port from cfg might not be needed for it.
        # For safety, we can keep them in cfg if setup_dist needs them,
        # but torchrun sets MASTER_ADDR and MASTER_PORT in the environment.
        master_addr = os.environ.get('MASTER_ADDR', 'localhost')
        master_port = os.environ.get('MASTER_PORT', '12345') # Default if not set
        print(f"Rank {rank}: Initializing distributed training. local_rank={local_rank}, world_size={world_size}, master_addr={master_addr}, master_port={master_port}")
        setup_dist(rank, local_rank, world_size, master_addr, master_port)
    else:
        print("Rank 0: Running in single GPU mode.")
        if torch.cuda.is_available():
             torch.cuda.set_device(local_rank) # for consistency, local_rank is 0

    # Seed rngs
    setup_rng(rank) # Seed with global rank

    if cfg.random_cond_gt:
        cfg.dataset.args.random_cond_gt = True
    if cfg.coords_aug_size is not None:
        cfg.dataset.args.coords_aug_size = cfg.coords_aug_size
    if cfg.feats_aug_grid_size is not None:
        cfg.dataset.args.feats_aug_grid_size = cfg.feats_aug_grid_size
    if cfg.feats_aug_ratio is not None:
        cfg.dataset.args.feats_aug_ratio = cfg.feats_aug_ratio
    if cfg.voxel_aug_ratio is not None:
        cfg.dataset.args.voxel_aug_ratio = cfg.voxel_aug_ratio
    if cfg.adapt_simple_edit_data:
        cfg.dataset.args.adapt_simple_edit_data = True
    if cfg.mixamo_data_repeat_ratio is not None:
        cfg.dataset.args.mixamo_data_repeat_ratio = cfg.mixamo_data_repeat_ratio
    if cfg.random_ori_edit is not None:
        cfg.dataset.args.random_ori_edit = cfg.random_ori_edit
    if cfg.simple_edit_data_if_filtered:
        cfg.dataset.args.simple_edit_data_if_filtered = True
        
    # Load data
    default_data_dir = 'data'  # not used in the dataset init
    dataset = getattr(datasets, cfg.dataset.name)(default_data_dir, **cfg.dataset.args)

    # Build model
    if cfg.ori_ss_latents_weights is not None:
        cfg.models.denoiser.args.ori_ss_latents_weights = cfg.ori_ss_latents_weights
    if cfg.feats_3d_t is not None:
        cfg.models.denoiser.args.feats_3d_t = cfg.feats_3d_t
    model_dict = {
        name: getattr(models, model.name)(**model.args).cuda() # .cuda() will use the device set by set_device
        for name, model in cfg.models.items()
    }
    # if cfg.lora:
    #     lora_config = LoraConfig(
    #         r=cfg.lora_rank,
    #         lora_alpha=cfg.lora_alpha,
    #         lora_dropout=cfg.lora_dropout,
    #         target_modules=["to_qkv", "to_out", "to_q", "to_kv"], # "mlp"
    #         exclude_modules='.*editing.*'
    #     )
    #     original_denoiser = model_dict['denoiser']
    #     model_dict['denoiser'] = LoraModel(original_denoiser, lora_config, "default")

    #     for name, param in model_dict['denoiser'].named_parameters():
    #         if 'editing' in name:
    #             param.requires_grad = True

    #     cfg.trainer.args.lora = True

    if cfg.lr is not None:
        cfg.trainer.args.optimizer.args.lr = cfg.lr

    if cfg.batch_size_per_gpu is not None:
        cfg.trainer.args.batch_size_per_gpu = cfg.batch_size_per_gpu

    if cfg.batch_split is not None:
        cfg.trainer.args.batch_split = cfg.batch_split

    if cfg.max_steps is not None:
        cfg.trainer.args.max_steps = cfg.max_steps

    if cfg.train_only_editing_weights:
        for name, param in model_dict['denoiser'].named_parameters():
            if 'editing' in name:
                param.requires_grad = True
            else:
                param.requires_grad = False

    if cfg.debug:
        cfg.trainer.args.max_steps = 100
        cfg.trainer.args.i_print = 1
        cfg.trainer.args.i_log = 1
        cfg.trainer.args.i_sample = 10
        cfg.trainer.args.i_save = 100
        cfg.trainer.args.init_sample = False
        cfg.trainer.args.init_dataset_vis = False

    if cfg.no_sample_images:
        cfg.trainer.args.no_sample_images = True

    # Model summary
    if rank == 0:
        for name, backbone in model_dict.items():
            model_summary = get_model_summary(backbone)
            print(f'\n\nBackbone: {name}\n' + model_summary)
            with open(os.path.join(cfg.output_dir, f'{name}_model_summary.txt'), 'w') as fp:
                print(model_summary, file=fp)

    # Build trainer
    trainer = getattr(trainers, cfg.trainer.name)(
        model_dict, dataset, **cfg.trainer.args,
        output_dir=cfg.output_dir, load_dir=cfg.load_dir, step=cfg.load_ckpt
    )

    # Train
    if not cfg.tryrun:
        if cfg.profile:
            trainer.profile()
        else:
            trainer.run()

    if world_size > 1:
        torch.distributed.barrier() # Ensure all processes finish before exiting
        torch.distributed.destroy_process_group()


if __name__ == '__main__':
    # Arguments and config
    parser = argparse.ArgumentParser()
    ## config
    parser.add_argument('--config', type=str, required=True, help='Experiment config file')
    ## io and resume
    parser.add_argument('--output_dir', type=str, required=True, help='Output directory')
    parser.add_argument('--load_dir', type=str, default='', help='Load directory, default to output_dir')
    parser.add_argument('--ckpt', type=str, default='latest', help='Checkpoint step to resume training, default to latest')
    parser.add_argument('--data_dir', type=str, default='/path_to_3DEditVerse/', help='Data directory')
    parser.add_argument('--auto_retry', type=int, default=0, help='Number of retries on error (simplified for torchrun)') # Max retries for main function
    parser.add_argument('--seed', type=int, default=42, help='Base random seed.')
    ## dubug
    parser.add_argument('--tryrun', action='store_true', help='Try run without training')
    parser.add_argument('--profile', action='store_true', help='Profile training')
    ## training
    parser.add_argument('--lr', type=float, default=None, help='Learning rate')
    parser.add_argument('--batch_size_per_gpu', type=int, default=None, help='Batch size per gpu')
    parser.add_argument('--batch_split', type=int, default=None, help='Batch split')
    parser.add_argument('--max_steps', type=int, default=None, help='Max steps')
    parser.add_argument('--debug', action='store_true', help='Debug mode')
    parser.add_argument('--train_only_editing_weights', action='store_true', help='Train only editing weights')
    parser.add_argument('--ori_ss_latents_weights', type=float, default=None, help='Weight for ori ss latents fusing with noising latents')
    parser.add_argument('--feats_3d_t', type=float, nargs=2, default=None, help='Feats 3d t')  # [0.1, 0.9]
    parser.add_argument('--no_sample_images', action='store_true', help='No sample images')
    ## dataset
    parser.add_argument('--random_cond_gt', action='store_true', help='Use random cond gt')
    parser.add_argument('--coords_aug_size', type=int, default=None, help='Coords aug size')
    parser.add_argument('--feats_aug_grid_size', type=int, nargs='*', default=None, help='Feats aug grid size')
    parser.add_argument('--feats_aug_ratio', type=float, nargs=2, default=None, help='Feats aug ratio')
    parser.add_argument('--voxel_aug_ratio', type=float, nargs=2, default=None, help='Voxel aug ratio')
    parser.add_argument('--adapt_simple_edit_data', action='store_true', help='Adapt simple edit data')
    parser.add_argument('--mixamo_data_repeat_ratio', type=float, default=None, help='Mixamo data repeat ratio')
    parser.add_argument('--random_ori_edit', type=float, default=None, help='Random ori edit data')
    parser.add_argument('--simple_edit_data_if_filtered', action='store_true', help='Simple edit data if filtered')
    opt = parser.parse_args()

    opt.load_dir = opt.load_dir if opt.load_dir != '' else opt.output_dir
    # opt.num_gpus is not used to launch processes anymore with torchrun.
    # It can be kept if cfg.num_gpus is used elsewhere in the logic,
    # otherwise, it's informational.
    # The actual number of GPUs used per node is determined by `torchrun --nproc_per_node`.

    ## Load config
    def replace_data_dir_placeholders(value):
        if isinstance(value, str):
            return value.replace('/path_to_3DEditVerse', opt.data_dir)
        if isinstance(value, list):
            return [replace_data_dir_placeholders(v) for v in value]
        if isinstance(value, dict):
            return {k: replace_data_dir_placeholders(v) for k, v in value.items()}
        return value

    config_from_file = replace_data_dir_placeholders(json.load(open(opt.config, 'r')))
    
    ## Combine arguments and config
    cfg = edict()
    cfg.update(opt.__dict__) # Command line args take precedence
    cfg.update(config_from_file) # Then update with file config (potentially overwriting CLI defaults if not specified in CLI)
    # To ensure CLI overrides file config for shared keys:
    # temp_cfg = edict(config_from_file)
    # temp_cfg.update(opt.__dict__) # CLI overrides file
    # cfg = temp_cfg

    # Update cfg with command line arguments again to ensure they have priority
    # This makes CLI args override json config values.
    for key, value in opt.__dict__.items():
        # Only update if the arg was actually provided or is not the default for action='store_true'
        if value is not None:
             # For argparse arguments with defaults, they will always be in opt.__dict__.
             # For 'action=store_true', default is False. If specified, it's True.
             # This logic ensures CLI args effectively override config file values.
            is_default_argparse = False
            for action in parser._actions:
                if action.dest == key:
                    if action.default == value and not isinstance(action, argparse._StoreTrueAction) and not isinstance(action, argparse._StoreFalseAction):
                        # Check if the value is the default AND it wasn't explicitly set by user
                        # This part is tricky without checking sys.argv. A simpler approach is just to override.
                        pass # Simpler to just let CLI override.
            cfg[key] = value


    # Get rank for file operations (like saving config)
    # These env vars are set by torchrun.
    # Use 0 if not in a distributed environment (e.g. world_size=1)
    current_rank = int(os.environ.get('RANK', 0))
    world_size_for_setup = int(os.environ.get('WORLD_SIZE', 1))

    if current_rank == 0: # Only master process should create dirs and save initial files
        print('\n\nConfig:')
        print('=' * 80)
        # Use a serializable dictionary for printing/saving
        # edict can sometimes have issues with json.dump if it contains non-standard types.
        config_to_print_save = dict(cfg)
        print(json.dumps(config_to_print_save, indent=4))

        os.makedirs(cfg.output_dir, exist_ok=True)
        ## Save command and config
        with open(os.path.join(cfg.output_dir, 'command.txt'), 'w') as fp:
            print(' '.join(['python'] + sys.argv), file=fp) # This will show the torchrun command if applicable
        with open(os.path.join(cfg.output_dir, 'config.json'), 'w') as fp:
            json.dump(config_to_print_save, fp, indent=4)

    # Run
    # The auto_retry logic needs to be within the main call if it's for application-level retries.
    # If it was for process launch failures, torchrun/scheduler handles that.
    # For simplicity, if an error occurs in `main` and `auto_retry` is > 0, we can try rerunning `main`.
    # Note: This simplistic retry doesn't reset CUDA state or other global states perfectly.
    # A more robust retry would be at the job submission level.

    cfg = find_ckpt(cfg) # Find checkpoint before potentially entering retry loop

    if cfg.auto_retry > 0 and world_size_for_setup > 1: # Only attempt retry if distributed
        print(f"Warning: auto_retry ({cfg.auto_retry}) within torchrun script has limited effect and might not recover from all errors. Job-level retry is preferred.")

    for rty in range(cfg.auto_retry + 1): # +1 because range is exclusive at the end, so 0 retries means 1 attempt
        try:
            # cfg = find_ckpt(cfg) # Moved outside loop; typically don't want to re-find checkpoint on retry unless intended
            main(cfg)
            break # Success
        except Exception as e:
            print(f"Error during main execution: {e}")
            if rty < cfg.auto_retry:
                print(f"Retrying ({rty + 1}/{cfg.auto_retry})...")
                if world_size_for_setup > 1: # If distributed, wait a bit before retrying
                    torch.distributed.barrier() # Wait for all processes to hit the error before retrying
                    # A small delay might be useful in some cases
                    # import time
                    # time.sleep(5)
            else:
                print("Max retries reached. Failing.")
                if world_size_for_setup > 1 and torch.distributed.is_initialized():
                    torch.distributed.destroy_process_group()
                raise e # Re-raise the exception if max retries are exhausted