Add training params
(opt)
| 313 | |
| 314 | |
| 315 | def add_training_params(opt): |
| 316 | """Add training params""" |
| 317 | opt.add_argument("--seq_length", |
| 318 | type=int, |
| 319 | default=2048, |
| 320 | help="sequence length, default is 2048.") |
| 321 | opt.add_argument("--vocab_size", |
| 322 | type=int, |
| 323 | default=40000, |
| 324 | help="vocabulary size, default is 40000.") |
| 325 | opt.add_argument("--embedding_size", |
| 326 | type=int, |
| 327 | default=16384, |
| 328 | help="embedding table size, default is 16384.") |
| 329 | opt.add_argument("--num_layers", |
| 330 | type=int, |
| 331 | default=64, |
| 332 | help="total layers, default is 64.") |
| 333 | opt.add_argument("--num_heads", |
| 334 | type=int, |
| 335 | default=128, |
| 336 | help="head size, default is 128.") |
| 337 | opt.add_argument("--stage_num", |
| 338 | type=int, |
| 339 | default=1, |
| 340 | help="Pipeline stage num, default is 1.") |
| 341 | opt.add_argument("--micro_size", |
| 342 | type=int, |
| 343 | default=1, |
| 344 | help="Pipeline micro_size, default is 1.") |
| 345 | opt.add_argument("--eod_reset", |
| 346 | type=int, |
| 347 | default=1, |
| 348 | help="Enable eod mask, default is 1.") |
| 349 | opt.add_argument("--warmup_step", |
| 350 | type=int, |
| 351 | default=2000, |
| 352 | help="Warmup step, default is 2000.") |
| 353 | opt.add_argument("--decay_steps", |
| 354 | type=int, |
| 355 | default=200000, |
| 356 | help="Decay step, default is 200000.") |
| 357 | opt.add_argument("--optimizer", |
| 358 | type=str, |
| 359 | default="adam", |
| 360 | choices=["adam", "lamb"], |
| 361 | help="select which optimizer to be used, default adam") |
| 362 | opt.add_argument("--opt_offload", |
| 363 | type=int, default=0, |
| 364 | help="Enable optimizer status offload to host CPU, default is 0") |
| 365 | opt.add_argument("--use_moe", |
| 366 | type=int, default=0, |
| 367 | help="Use moe, default is 0") |
| 368 | opt.add_argument("--per_dp_dim_expert_num", |
| 369 | type=int, default=1, |
| 370 | help="Expert nums in one data parallel dim, only effective when applying moe, default is 1") |
| 371 | opt.add_argument("--eod_id", |
| 372 | type=int, default=50256, |