forked from NVIDIA/Megatron-LM
-
Notifications
You must be signed in to change notification settings - Fork 0
/
arguments.py
25 lines (22 loc) · 813 Bytes
/
arguments.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
def add_ammo_args(parser):
"""Add additional arguments for ammo."""
group = parser.add_argument_group(title="ammo-generic")
group.add_argument(
"--ammo-load-classic-megatron-to-mcore",
action="store_true",
help="Load a classic megatron-lm checkpoint to a new megatron-core model.",
)
group.add_argument(
"--ammo-convert-te-to-local-spec",
action="store_true",
help="Load a megatron-core transformer-engine checkpoint to a model with local spec.",
)
group.add_argument(
"--ammo-quant-cfg",
type=str,
default=None,
choices=["int8_sq", "fp8", "int4_awq", "None"],
help="Algorithms supported by atq.quantize.",
)
return parser