trainer_misc/__init__.py

from .utils import (
    create_optimizer,
    get_rank,
    get_world_size,
    is_main_process,
    is_dist_avail_and_initialized,
    init_distributed_mode, 
    setup_for_distributed, 
    cosine_scheduler,
    constant_scheduler,
    NativeScalerWithGradNormCount,
    auto_load_model,
    save_model,
)

from .sp_utils import (
    is_sequence_parallel_initialized,
    init_sequence_parallel_group,
    get_sequence_parallel_group,
    get_sequence_parallel_world_size,
    get_sequence_parallel_rank,
    get_sequence_parallel_group_rank,
    get_sequence_parallel_proc_num,
    init_sync_input_group,
    get_sync_input_group,
)

from .communicate import all_to_all
from .fsdp_trainer import train_one_epoch_with_fsdp
from .vae_ddp_trainer import train_one_epoch