Main function for each worker process (may be only 1 for BasicTrainer/TensorParallelTrainer).
(rank: int, world_size: int, config: DictConfig, policy: nn.Module, reference_model: Optional[nn.Module] = None)
| 19 | |
| 20 | |
| 21 | def worker_main(rank: int, world_size: int, config: DictConfig, policy: nn.Module, reference_model: Optional[nn.Module] = None): |
| 22 | """Main function for each worker process (may be only 1 for BasicTrainer/TensorParallelTrainer).""" |
| 23 | if 'FSDP' in config.trainer: |
| 24 | init_distributed(rank, world_size, port=config.fsdp_port) |
| 25 | |
| 26 | if config.debug: |
| 27 | wandb.init = lambda *args, **kwargs: None |
| 28 | wandb.log = lambda *args, **kwargs: None |
| 29 | |
| 30 | if rank == 0 and config.wandb.enabled: |
| 31 | os.environ['WANDB_CACHE_DIR'] = get_local_dir(config.local_dirs) |
| 32 | wandb.init( |
| 33 | entity=config.wandb.entity, |
| 34 | project=config.wandb.project, |
| 35 | config=OmegaConf.to_container(config), |
| 36 | dir=get_local_dir(config.local_dirs), |
| 37 | name=config.exp_name, |
| 38 | ) |
| 39 | |
| 40 | TrainerClass = getattr(trainers, config.trainer) |
| 41 | print(f'Creating trainer on process {rank} with world size {world_size}') |
| 42 | trainer = TrainerClass(policy, config, config.seed, config.local_run_dir, reference_model=reference_model, rank=rank, world_size=world_size) |
| 43 | |
| 44 | trainer.train() |
| 45 | trainer.save() |
| 46 | |
| 47 | |
| 48 | @hydra.main(version_base=None, config_path="config", config_name="config") |
no test coverage detected