(args, rank, current_time)
| 176 | |
| 177 | |
| 178 | def prepare_configs(args, rank, current_time): |
| 179 | base_path = os.path.dirname(os.path.abspath(__file__)) |
| 180 | |
| 181 | current_time_str = time.strftime("%Y%m%d-%H%M%S", time.localtime(current_time)) |
| 182 | run_path = os.path.join(base_path, "runs", current_time_str) |
| 183 | config_path = os.path.join(run_path, "config.yaml") |
| 184 | if rank == 0: |
| 185 | os.makedirs(run_path) |
| 186 | |
| 187 | with open(os.path.join(base_path, "config", f"{args.dataset}-template.yaml")) as f: |
| 188 | config = yaml.safe_load(f) |
| 189 | |
| 190 | config["name"] += f"-{current_time_str}" |
| 191 | config["checkpoint_root_dir"] = os.path.join(run_path, "checkpoints") |
| 192 | set_rollout_engine_config(config, args) |
| 193 | config["model"]["model_path"] = ( |
| 194 | args.model_path |
| 195 | or config["model"]["model_path"] |
| 196 | or os.environ.get(MODEL_PATH_ENV_VAR, "Qwen/Qwen2.5-1.5B-Instruct") |
| 197 | ) |
| 198 | if ALGORITHM_TYPE.get(config["algorithm"]["algorithm_type"]).use_critic: |
| 199 | config["model"]["critic_model_path"] = ( |
| 200 | args.critic_model_path |
| 201 | or config["model"].get("critic_model_path") |
| 202 | or config["model"]["model_path"] |
| 203 | ) |
| 204 | if args.critic_lr: |
| 205 | config["trainer"]["trainer_config"]["critic"]["optim"]["lr"] = args.critic_lr |
| 206 | if args.dataset == "alfworld": |
| 207 | print( |
| 208 | "Warning: The current benchmark script of ALFWorld only supports GRPO; the SFT stage will be supported soon." |
| 209 | ) |
| 210 | taskset_config = config["buffer"]["explorer_input"]["taskset"] |
| 211 | taskset_config["path"] = check_taskset_path( |
| 212 | args.dataset, |
| 213 | args.taskset_path or os.environ.get("TASKSET_PATH") or taskset_config["path"], |
| 214 | ) |
| 215 | eval_taskset_config = config["buffer"]["explorer_input"]["eval_tasksets"] |
| 216 | if len(eval_taskset_config) > 0: |
| 217 | # TODO: support separately set path for eval taskset |
| 218 | for eval_taskset_config in eval_taskset_config: |
| 219 | eval_taskset_config["path"] = taskset_config["path"] |
| 220 | if args.lr: |
| 221 | config["algorithm"]["optimizer"]["lr"] = args.lr |
| 222 | if args.sync_interval: |
| 223 | config["synchronizer"]["sync_interval"] = args.sync_interval |
| 224 | if args.sync_offset: |
| 225 | config["synchronizer"]["sync_offset"] = args.sync_offset |
| 226 | if args.sync_style: |
| 227 | config["synchronizer"]["sync_style"] = args.sync_style |
| 228 | if args.trainer_strategy: |
| 229 | config["trainer"]["trainer_strategy"] = args.trainer_strategy |
| 230 | if args.monitor: |
| 231 | config["monitor"]["monitor_type"] = args.monitor |
| 232 | |
| 233 | with open(config_path, "w") as f: |
| 234 | yaml.dump(config, f, allow_unicode=True, sort_keys=False) |
| 235 | return config_path |
no test coverage detected