Speedup sync_float8_amax_and_scale_history

[drisspg]

Summary

Once per training iteration we need to sync the amax_history and update scales for all float8 linear module. Currently this is done by iterating over all child modules, finding the linears and syncing. The picture below is the eager performance of calling this on one linear_fp8_module.

Code Pointer: https://github.com/pytorch-labs/float8_experimental/blob/f6302af11bd29855bb3cbef93209ddffdd21dd2b/float8_experimental/float8_linear.py#L310

I also tried compiling just this function to see if this could be reduced.

As well I tried compiling this function with reduce-overhead and the warning skipping cudagraphs due to ['mutated inputs'] was raised

[drisspg]

Step one in improving Perf: #211

[drisspg]

Example showing that reduce overhead is being skipped for mutation on input:

from float8_experimental.float8_linear_utils import (
    swap_linear_with_float8_linear,
    sync_float8_amax_and_scale_history,
    get_float8_layers,
)
from float8_experimental.float8_linear import Float8Linear


import torch
import torch.nn as nn


my_module = nn.Sequential(
    nn.Linear(16, 32, bias=True), nn.ReLU(), nn.Linear(32, 16, bias=True)
).to("cuda")

# Swap the linear layers with float8 linear layers
swap_linear_with_float8_linear(my_module, Float8Linear)


inpt = torch.randn(16, 16, device="cuda", dtype=torch.float32, requires_grad=True)

target = torch.rand(16, 16, device="cuda", dtype=torch.float32, requires_grad=False)


sync_func = torch.compile(
    sync_float8_amax_and_scale_history, mode="reduce-overhead", fullgraph=True
)
fp8_layers = get_float8_layers(my_module)

optimizer = torch.optim.SGD(my_module.parameters(), lr=0.1)

for _ in range(10):
    out = my_module(inpt)
    loss = torch.nn.functional.mse_loss(out, target)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    sync_func(my_module, fp8_layers)

    print("Loss: ", loss.item())

[vkuzo]

going to call this done, but feel free to reopen if not