[Update] Support SVD method to calculate M^{-1/p}

[kozistr]

Problem (Why?)

Theoretically, Schur-Newton method is faster than SVD method to calculate M^{-1/p}. However, the inefficiency of the loop and others, in some cases, SVD is much faster than that. So, supports SVD method too. (related to #100)

Solution (What/How?)

• calculate the power of the matrix with SVD
  • compute on GPU (not CPU this time)
• perform batch SVD when the shapes of the pre-conditioners are all same. (maybe later, RaggedTensor could be used here.)
• rename the new Shampoo optimizer to ScalableShampoo. (actually, not scalable though)
• implement the original Shampoo optimizer.

Benchmark

tested on i7700K + GTX 1060 6GB.

backbone: resmlp_12_distilled_224, bs: 16

x4.325 faster (A -> B)

• AdamP: 3.73 iter / s
• (old) Shampoo: 25s / iter
• Scalable Shampoo w/ Schur-Newton (block size = 256): 1.73s / iter -> A
• Scalable Shampoo w/ SVD (block size = 256): 1.59 iter / s
• Scalable Shampoo w/ SVD (block size = 512): 2.50 iter / s -> B

backbone: mixer_s32_224, bs: 8

x5.408 slower (A -> B)

• AdamP: 3.85 iter / s
• Scalable Shampoo w/ Schur-Newton (block size = 256): 1.68s / iter
• Scalable Shampoo w/ Schur-Newton (block size = 512): 1.05 iter / s -> A
• Scalable Shampoo w/ SVD (block size = 256): 5.15s / iter -> B
• Scalable Shampoo w/ SVD (block size = 512): 7.01s / iter

backbone: mixer_b16_224, bs: 2

x3.292 slower (A -> B)

• AdamP: 3.15 iter / s
• (old) Shampoo: over 2 mins / iter
• Scalable Shampoo w/ Schur-Newton (block size = 256): 5.11s / iter -> A
• Scalable Shampoo w/ SVD (block size = 256): 16.82s / s -> B
• Scalable Shampoo w/ SVD (block size = 512): 32.47s / iter

code

    from timm import create_model
    from tqdm import tqdm

    model = create_model(backbone, pretrained=False, num_classes=1)
    model.train()
    model.cuda()

    optimizer = load_optimizer('scalableshampoo')(
        model.parameters(), 
        start_preconditioning_step=1,
        block_size=block_size,
        use_svd=use_svd,
    )

    inp = torch.randn((bs, 3, 224, 224), dtype=torch.float32).cuda()
    y = torch.randn((bs, 1), dtype=torch.float32).cuda()

    for _ in tqdm(range(100)):
        optimizer.zero_grad()

        torch.nn.functional.binary_cross_entropy_with_logits(model(inp), y).backward()

        optimizer.step()

Other changes (bug fixes, small refactors)

nope

Notes

nope

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.71%. Comparing base (de06f63) to head (01b5c5a).
Report is 1686 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #103   +/-   ##
=======================================
  Coverage   99.70%   99.71%           
=======================================
  Files          39       39           
  Lines        3034     3125   +91     
=======================================
+ Hits         3025     3116   +91     
  Misses          9        9           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.