About this repository

This is a fork of DistributedShampoo extended with Newton-DB and Chebyshev polynomials to compute matrix inverse roots in the context of DASH (paper, code), a faster and more accurate version of DistributedShampoo.

Optimizers

Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.

Description

Optimizers is a Github repository of PyTorch optimization algorithms. It is designed for external collaboration and development.

Currently includes the optimizers:

• Distributed Shampoo

See the CONTRIBUTING file for how to help out.

License

Optimizers is released under the BSD license.

Installation and Dependencies

Install distributed_shampoo with all dependencies:

git clone git@github.com:facebookresearch/optimizers.git
cd optimizers
pip install .

If you also want to try the examples, replace the last line with pip install ".[examples]".

Usage

After installation, basic usage looks like:

import torch
from distributed_shampoo import AdamPreconditionerConfig, DistributedShampoo

model = ...  # Instantiate model

optim = DistributedShampoo(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    epsilon=1e-8,
    grafting_config=AdamPreconditionerConfig(
        beta2=0.999,
        epsilon=1e-8,
    ),
)

For more, please see the additional documentation here and especially the How to Use section.