Text generate arguments.
(parser)
| 1451 | |
| 1452 | |
| 1453 | def _add_zero_args(parser): |
| 1454 | """Text generate arguments.""" |
| 1455 | |
| 1456 | group = parser.add_argument_group("ZeRO configurations", "configurations") |
| 1457 | group.add_argument("--zero-stage", type=int, default=1.0) |
| 1458 | group.add_argument( |
| 1459 | "--zero-reduce-scatter", |
| 1460 | action="store_true", |
| 1461 | help="Use reduce scatter if specified", |
| 1462 | ) |
| 1463 | group.add_argument( |
| 1464 | "--zero-contigious-gradients", |
| 1465 | action="store_true", |
| 1466 | help="Use contigious memory optimizaiton if specified", |
| 1467 | ) |
| 1468 | group.add_argument("--zero-reduce-bucket-size", type=int, default=0.0) |
| 1469 | group.add_argument("--zero-allgather-bucket-size", type=int, default=0.0) |
| 1470 | group.add_argument( |
| 1471 | "--remote-device", |
| 1472 | type=str, |
| 1473 | default="none", |
| 1474 | choices=["none", "cpu", "nvme"], |
| 1475 | help="Remote device for ZeRO-3 initialized parameters.", |
| 1476 | ) |
| 1477 | group.add_argument( |
| 1478 | "--use-pin-memory", |
| 1479 | action="store_true", |
| 1480 | help="Use pinned CPU memory for ZeRO-3 initialized model parameters.", |
| 1481 | ) |
| 1482 | return parser |
| 1483 | |
| 1484 | |
| 1485 | def _add_memoryopt_args(parser): |