Modded GPT pseudoinverse improvement

[rohan-anil]

https://x.com/_arohan_/status/2064216022232826080

The main change that improved things was replacing:

  inv_power_L = L.pow(-1.0 / root)

with a truncated inverse-power version:

  positive_eigenvalue_mask = L > 1e-15
  inv_power_L = torch.zeros_like(L)
  inv_power_L[positive_eigenvalue_mask] = L[positive_eigenvalue_mask].pow(-1.0 / root)

My read is that very small eigenvalues were getting amplified too aggressively by the inverse root, and truncating them behaves more like a pseudoinverse / rank-truncated inverse. That seems to work better for this problem. This problem has been dominating the twitter conversations around optimizers.

I saw a few related stability options in the codebase. Is there another configuration you would recommend trying, or should we add explicit support for this pseudoinverse-style truncated inverse?

[rohan-anil]

Here is how the stats look like. Early in training it has junk

[rohan-anil]

Tagging @tsunghsienlee @hjmshi :)

[anana10c]

Hi @rohan-anil! We actually have this exact option implemented as PseudoInverseConfig which can be set inside the amortized_computation_config, e.g.:

optimizer = DistributedShampoo(
    model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999),
    epsilon=0.0,
    ...,
    preconditioner_config=RootInvShampooPreconditionerConfig(
        amortized_computation_config=EigenConfig(
            rank_deficient_stability_config=PseudoInverseConfig(
                rank_atol=1e-15,
                rank_rtol=0.0,
            ),
        ),
    ),
)

(You can see the implementation here: https://github.com/facebookresearch/optimizers/blob/main/distributed_shampoo/preconditioner/matrix_functions.py#L200)

However, note that this option forces epsilon to be 0.0, as we originally envisioned it to be an alternative to using epsilon / Tikhonov regularization for stability purposes. Do you think this makes sense?

Sorry that we haven't updated the README with details about this feature, we will add that!

[rohan-anil]

Thanks Anna! I did find it --the difference is very stark for speed run:

before

epsilon=1e-12
preconditioner_config=RootInvShampooPreconditionerConfig() # implicit default

after

epsilon=0.0
preconditioner_config=RootInvShampooPreconditionerConfig(
amortized_computation_config=EigenConfig(
rank_deficient_stability_config=PseudoInverseConfig(),
),
)

[rohan-anil]

For @anana10c @hjmshi @tsunghsienlee

In case you are curious in reproducing speed run results, and optimizing the code for smaller problems: here is the full reproduction trace. I really like this distributed_shampoo package, this is the first time I am actually using it although have been rooting for it since I used to use my Jax implementations in my previous life.

sh-origpinv-s3375-lr1em2-wd010-b9em1-ge15-pf1-near1-record-130563a2-b40b-43c1-8fb4-b7a3bbfa5969.txt