| 43 | |
| 44 | |
| 45 | def parse_args(): |
| 46 | parser = argparse.ArgumentParser() |
| 47 | parser.add_argument( |
| 48 | "--step", |
| 49 | type=int, |
| 50 | nargs="+", |
| 51 | choices=(1, 2, 3), |
| 52 | default=(1, 2, 3), |
| 53 | help="Which steps of the ChatGPT example to run", |
| 54 | ) |
| 55 | parser.add_argument( |
| 56 | "--actor-model", |
| 57 | type=lambda x: x.replace("facebook/opt-", ""), |
| 58 | default="1.3b", |
| 59 | choices=("1.3b", "6.7b", "13b", "66b"), |
| 60 | help="Which facebook/opt-* model to use for Actor (step 1)", |
| 61 | ) |
| 62 | parser.add_argument( |
| 63 | "--reward-model", |
| 64 | type=lambda x: x.replace("facebook/opt-", ""), |
| 65 | default="350m", |
| 66 | choices=("350m"), |
| 67 | help="Which facebook/opt-* model to use for Reward (step 2)", |
| 68 | ) |
| 69 | parser.add_argument( |
| 70 | "--actor-zero-stage", |
| 71 | type=str, |
| 72 | default="", |
| 73 | choices=("", "0", "1", "2", "3"), |
| 74 | help="ZeRO stage for step 1 (Actor) training", |
| 75 | ) |
| 76 | parser.add_argument( |
| 77 | "--reward-zero-stage", |
| 78 | type=str, |
| 79 | default="", |
| 80 | choices=("", "0", "1", "2", "3"), |
| 81 | help="ZeRO stage for step 2 (Critic) training", |
| 82 | ) |
| 83 | parser.add_argument( |
| 84 | "--output-dir", |
| 85 | type=lambda x: os.path.abspath(x), |
| 86 | default="./output", |
| 87 | help="Directory for output of each step", |
| 88 | ) |
| 89 | parser.add_argument( |
| 90 | "--deployment-type", |
| 91 | type=str, |
| 92 | default="single_gpu", |
| 93 | choices=("single_gpu", "single_node", "multi_node"), |
| 94 | help="Number of GPUs to run the actor/reward models on", |
| 95 | ) |
| 96 | args = parser.parse_args() |
| 97 | |
| 98 | if args.actor_zero_stage != "" or args.reward_zero_stage != "": |
| 99 | warnings.warn( |
| 100 | "Non-default zero stages may result in OOM errors or worse performance." |
| 101 | ) |
| 102 | |