[float8] fuse abs/max with torch.linalg.vector_norm

[weifengpy]

torch.linalg.vector_norm is faster

torch.max(torch.abs) are slower because of 3 kernels and launching overhead

generated profiler traces with following code

import torch
import os
import contextlib
@contextlib.contextmanager
def enable_profiling(enable=False):
    if not enable:
        torch_profiler = contextlib.nullcontext()
        yield None
    else:
        trace_dir = "./profilers"
        rank = 0
        def trace_handler(prof):
            curr_trace_dir_name = "iteration_" + str(prof.step_num)
            curr_trace_dir = os.path.join(trace_dir, curr_trace_dir_name)
            if not os.path.exists(curr_trace_dir):
                os.makedirs(curr_trace_dir, exist_ok=True)
            prof.export_chrome_trace(f"{curr_trace_dir}/rank{rank}_trace.json")
            # torch.distributed.barrier()
        if not os.path.exists(trace_dir):
            os.makedirs(trace_dir, exist_ok=True)
        warmup, active = 1, 2
        wait = 1
        with torch.profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active),
            on_trace_ready=trace_handler,
            record_shapes=True,
        ) as torch_profiler:
            yield torch_profiler

with enable_profiling(True) as torch_profiler:
    tensor = torch.rand(256, 1024, device="cuda")
    for i in range(10):
        with torch.profiler.record_function("vector_norm"):
            # amax = torch.max(torch.abs(tensor))
            torch.linalg.vector_norm(tensor, ord=float("inf"))
        torch_profiler.step()
        

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/905

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit bf49f41 with merge base 53b6b78 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vkuzo]

torch.compile performance is important in this code path, thoughts on how does this PR impact performance when torch.compile is on?

[weifengpy]

torch.compile performance is important in this code path, thoughts on how does this PR impact performance when torch.compile is on?

good question! I can benchmark on inductor output code and update again

[y-sq]

When float8 is used, we usually compile the model. Was the issue of three kernels initially observed from the un-compiled fsdp_pre_all_gather region in a compiled model?

Besides the additional kernels for amax scaling factor, do we also see more kernels for cast_tp_fp8 part?

Is it possible to add an option in fp8-all-gather to explicitly compile the casting parts?