| 1518 | |
| 1519 | |
| 1520 | def _add_activation_checkpoint_args(parser): |
| 1521 | group = parser.add_argument_group( |
| 1522 | "Activation Checkpointing", "Checkpointing Configurations" |
| 1523 | ) |
| 1524 | group.add_argument( |
| 1525 | "--deepspeed-activation-checkpointing", |
| 1526 | action="store_true", |
| 1527 | help="uses activation checkpointing from deepspeed", |
| 1528 | ) |
| 1529 | group.add_argument( |
| 1530 | "--partition-activations", |
| 1531 | action="store_true", |
| 1532 | help="partition Activations across GPUs before checkpointing.", |
| 1533 | ) |
| 1534 | group.add_argument( |
| 1535 | "--contigious-checkpointing", |
| 1536 | action="store_true", |
| 1537 | help="Contigious memory checkpointing for activatoins.", |
| 1538 | ) |
| 1539 | group.add_argument( |
| 1540 | "--checkpoint-in-cpu", |
| 1541 | action="store_true", |
| 1542 | help="Move the activation checkpoints to CPU.", |
| 1543 | ) |
| 1544 | group.add_argument( |
| 1545 | "--synchronize-each-layer", |
| 1546 | action="store_true", |
| 1547 | help="does a synchronize at the beginning and end of each checkpointed layer.", |
| 1548 | ) |
| 1549 | group.add_argument( |
| 1550 | "--profile-backward", |
| 1551 | action="store_true", |
| 1552 | help="Enables backward pass profiling for checkpointed layers.", |
| 1553 | ) |
| 1554 | return parser |