(config, args)
| 15 | |
| 16 | |
| 17 | def set_rollout_engine_config(config, args): |
| 18 | config["cluster"]["node_num"] = args.node_num |
| 19 | config["cluster"]["gpu_per_node"] = args.gpu_per_node |
| 20 | batch_size = config["buffer"]["batch_size"] * config["algorithm"]["repeat_times"] |
| 21 | if config["mode"] == "train": |
| 22 | return |
| 23 | |
| 24 | rollout_model_config = config["explorer"]["rollout_model"] |
| 25 | |
| 26 | if args.engine_type is not None: |
| 27 | rollout_model_config["engine_type"] = args.engine_type |
| 28 | if args.tp_size is not None: |
| 29 | rollout_model_config["tensor_parallel_size"] = args.tp_size |
| 30 | tensor_parallel_size = rollout_model_config["tensor_parallel_size"] |
| 31 | |
| 32 | if tensor_parallel_size > args.gpu_per_node: |
| 33 | assert tensor_parallel_size % args.gpu_per_node == 0, ( |
| 34 | "Please adjust the value of `tensor_parallel_size` so that it is an integer " |
| 35 | "multiple of `gpu_per_node` for cross-node inference." |
| 36 | ) |
| 37 | rollout_model_config["nnodes"] = tensor_parallel_size // args.gpu_per_node |
| 38 | else: |
| 39 | rollout_model_config["nnodes"] = 1 |
| 40 | |
| 41 | if args.engine_num is not None: |
| 42 | rollout_model_config["engine_num"] = args.engine_num |
| 43 | else: # auto set engine_num |
| 44 | opt_explorer_num, opt_ratio_diff = None, float("inf") |
| 45 | total_gpu_num = args.node_num * args.gpu_per_node |
| 46 | |
| 47 | def update_opt_explorer_num(trainer_gpu_num, opt_explorer_num, opt_ratio_diff): |
| 48 | if batch_size % trainer_gpu_num != 0: |
| 49 | return opt_explorer_num, opt_ratio_diff |
| 50 | explorer_gpu_num = total_gpu_num - trainer_gpu_num |
| 51 | if explorer_gpu_num % tensor_parallel_size != 0: |
| 52 | return opt_explorer_num, opt_ratio_diff |
| 53 | explorer_num = explorer_gpu_num // tensor_parallel_size |
| 54 | ratio = explorer_num / trainer_gpu_num |
| 55 | if opt_ratio_diff > abs(ratio - args.explorer_trainer_ratio): |
| 56 | return explorer_num, abs(ratio - args.explorer_trainer_ratio) |
| 57 | return opt_explorer_num, opt_ratio_diff |
| 58 | |
| 59 | if args.node_num == 1: # single node |
| 60 | for trainer_gpu_num in range(1, args.gpu_per_node): |
| 61 | opt_explorer_num, opt_ratio_diff = update_opt_explorer_num( |
| 62 | trainer_gpu_num, opt_explorer_num, opt_ratio_diff |
| 63 | ) |
| 64 | else: # multi node |
| 65 | if tensor_parallel_size <= args.gpu_per_node: |
| 66 | assert args.gpu_per_node % tensor_parallel_size == 0, ( |
| 67 | "Please adjust the value of `tensor_parallel_size` so that it is a divisor " |
| 68 | "of `gpu_per_node`, or an integer multiple of `gpu_per_node` for cross-node inference." |
| 69 | ) |
| 70 | for trainer_node_num in range(1, args.node_num): |
| 71 | trainer_gpu_num = args.gpu_per_node * trainer_node_num |
| 72 | opt_explorer_num, opt_ratio_diff = update_opt_explorer_num( |
| 73 | trainer_gpu_num, opt_explorer_num, opt_ratio_diff |
| 74 | ) |
no test coverage detected