Main entry point for training. Validates config, creates/initializes model(s), and kicks off worker process(es).
(config: DictConfig)
| 47 | |
| 48 | @hydra.main(version_base=None, config_path="config", config_name="config") |
| 49 | def main(config: DictConfig): |
| 50 | """Main entry point for training. Validates config, creates/initializes model(s), and kicks off worker process(es).""" |
| 51 | |
| 52 | # Resolve hydra references, e.g. so we don't re-compute the run directory |
| 53 | OmegaConf.resolve(config) |
| 54 | |
| 55 | missing_keys: Set[str] = OmegaConf.missing_keys(config) |
| 56 | if missing_keys: |
| 57 | raise ValueError(f"Got missing keys in config:\n{missing_keys}") |
| 58 | |
| 59 | if config.eval_every % config.batch_size != 0: |
| 60 | print('WARNING: eval_every must be divisible by batch_size') |
| 61 | print('Setting eval_every to', config.eval_every - config.eval_every % config.batch_size) |
| 62 | config.eval_every = config.eval_every - config.eval_every % config.batch_size |
| 63 | |
| 64 | if 'FSDP' in config.trainer and config.fsdp_port is None: |
| 65 | free_port = get_open_port() |
| 66 | print('no FSDP port specified; using open port for FSDP:', free_port) |
| 67 | config.fsdp_port = free_port |
| 68 | |
| 69 | print(OmegaConf.to_yaml(config)) |
| 70 | |
| 71 | config_path = os.path.join(config.local_run_dir, 'config.yaml') |
| 72 | with open(config_path, 'w') as f: |
| 73 | OmegaConf.save(config, f) |
| 74 | |
| 75 | print('=' * 80) |
| 76 | print(f'Writing to {socket.gethostname()}:{config.local_run_dir}') |
| 77 | print('=' * 80) |
| 78 | |
| 79 | os.environ['XDG_CACHE_HOME'] = get_local_dir(config.local_dirs) |
| 80 | print('building policy') |
| 81 | model_kwargs = {'device_map': 'balanced'} if config.trainer == 'BasicTrainer' else {} |
| 82 | policy_dtype = getattr(torch, config.model.policy_dtype) |
| 83 | policy = transformers.AutoModelForCausalLM.from_pretrained( |
| 84 | config.model.name_or_path, cache_dir=get_local_dir(config.local_dirs), low_cpu_mem_usage=True, torch_dtype=policy_dtype, **model_kwargs) |
| 85 | disable_dropout(policy) |
| 86 | |
| 87 | if config.loss.name in {'dpo', 'ipo'}: |
| 88 | print('building reference model') |
| 89 | reference_model_dtype = getattr(torch, config.model.reference_dtype) |
| 90 | reference_model = transformers.AutoModelForCausalLM.from_pretrained( |
| 91 | config.model.name_or_path, cache_dir=get_local_dir(config.local_dirs), low_cpu_mem_usage=True, torch_dtype=reference_model_dtype, **model_kwargs) |
| 92 | disable_dropout(reference_model) |
| 93 | else: |
| 94 | reference_model = None |
| 95 | |
| 96 | if config.model.archive is not None: |
| 97 | state_dict = torch.load(config.model.archive, map_location='cpu') |
| 98 | step, metrics = state_dict['step_idx'], state_dict['metrics'] |
| 99 | print(f'loading pre-trained weights at step {step} from {config.model.archive} with metrics {json.dumps(metrics, indent=2)}') |
| 100 | policy.load_state_dict(state_dict['state']) |
| 101 | if config.loss.name in {'dpo', 'ipo'}: |
| 102 | reference_model.load_state_dict(state_dict['state']) |
| 103 | print('loaded pre-trained weights') |
| 104 | |
| 105 | if 'FSDP' in config.trainer: |
| 106 | world_size = torch.cuda.device_count() |
no test coverage detected