Memory optimization arguments.
(parser)
| 1483 | |
| 1484 | |
| 1485 | def _add_memoryopt_args(parser): |
| 1486 | """Memory optimization arguments.""" |
| 1487 | |
| 1488 | group = parser.add_argument_group("Memory optimizations", "configurations") |
| 1489 | group.add_argument( |
| 1490 | "--scattered-embeddings", |
| 1491 | action="store_true", |
| 1492 | help="Save memory by scattering embedding activations. " |
| 1493 | "Introduces dropout differences across MP configurations.", |
| 1494 | ) |
| 1495 | group.add_argument( |
| 1496 | "--split-transformers", |
| 1497 | action="store_true", |
| 1498 | help="Save memory by splitting transformer layers into two parts, " |
| 1499 | "allowing for more frequent activation checkpoint savings.", |
| 1500 | ) |
| 1501 | group.add_argument( |
| 1502 | "--memory-centric-tiled-linear", |
| 1503 | action="store_true", |
| 1504 | help="Save memory by tiling with deepspeed.zero.TiledLinear.", |
| 1505 | ) |
| 1506 | group.add_argument( |
| 1507 | "--tile-factor", |
| 1508 | type=int, |
| 1509 | default=1, |
| 1510 | help="Make all linear layers the same size of [hidden/tile_factor, hidden/tile_factor]. " |
| 1511 | "Must be enabled with --memory-centric-tiled-linear. " |
| 1512 | "Example A: if tile_factor=1, the qkv layer [hidden, 3* hidden] would be converted into [1,3] tiles of size [hidden,hidden]. " |
| 1513 | "Example B: if tile_factor=2, the intermediate layer [4*hidden, hidden] will be converted into [8, 2] tiles of size [hidden/2, hidden/2]. " |
| 1514 | "Default is 1.", |
| 1515 | ) |
| 1516 | |
| 1517 | return parser |
| 1518 | |
| 1519 | |
| 1520 | def _add_activation_checkpoint_args(parser): |