Skip to content

llm.kfac

/home/runner/work/llm-pytorch/llm-pytorch/llm/kfac.py

add_kfac_options() 

add_kfac_options(parser: argparse.ArgumentParser) -> None

Add CLI arguments for K-FAC.

Parameters:

  • parser (ArgumentParser) –

    Parser object to add K-FAC arguments to.

Source code in llm/kfac.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def add_kfac_options(parser: argparse.ArgumentParser) -> None:
    """Add CLI arguments for K-FAC.

    Args:
        parser: Parser object to add K-FAC arguments to.
    """
    args_str = ' '.join(sys.argv)
    group = parser.add_argument_group(
        title='K-FAC',
        description='K-FAC preconditioning options',
    )
    group.add_argument(
        '--kfac',
        action='store_true',
        help='Use K-FAC for preconditioning',
    )
    group.add_argument(
        '--kfac-factor-update-steps',
        type=int,
        required=bool(re.search(r'--kfac($|\s)', args_str)),
        help='Steps between updating Kronecker factors',
    )
    group.add_argument(
        '--kfac-inv-update-steps',
        type=int,
        required=bool(re.search(r'--kfac($|\s)', args_str)),
        help='Steps between recomputing inverses/eigen decompositions',
    )
    group.add_argument(
        '--kfac-damping',
        type=float,
        default=0.001,
        help='K-FAC damping value',
    )
    group.add_argument(
        '--kfac-factor-decay',
        type=float,
        default=0.95,
        help='K-FAC factor decay rate',
    )
    group.add_argument(
        '--kfac-kl-clip',
        type=float,
        default=0.001,
        help='K-FAC KL-clip',
    )