Repository: pytorch
Issue: [PT2][fp8][FSDP2] compile the function that pre-computes fp8 amax 
Labels: oncall: distributed,triaged,module: fsdp,oncall: pt2,pt2d-triage-nov2024

### 🚀 The feature, motivation and pitch

share repro for @bdhirsh , @tugsbayasgalan on the gaps of torch.compile for FSDP2 fp8 all-gather

for FSDP2 fp8 all-gather, it's criticial to pre-compute amax for `model.parameters()`
```
loss = module(local_inp).sum()
loss.backward()
optim.step()
compute_amaxes_compile(module.parameters())
```

However, `compute_amaxes_compile` is slower than `compute_amaxes_eager`

Sanpshot of cpu overhead
<img width="1360" alt="Screenshot 2024-06-24 at 11 37 23 PM" src="https://github.com/pytorch/pytorch/assets/134637289/cb426bcb-1a2c-476e-95f2-a7342ea00396">

snapshot of gpu time
<img width="1239" alt="Screenshot 2024-06-24 at 11 37 36 PM" src="https://github.com/pytorch/pytorch/assets/134637289/8557e585-0cdd-48ac-bf56-44c5ce0143c6">

repro
* pytorch with nested tensor subclass support, otherwies IMA during torch.compile https://github.com/pytorch/pytorch/pull/127431
* torchao: https://github.com/pytorch/ao
* command `torchrun --standalone --nproc_per_node=2 compile_precompute.py`  

```
# torchrun --standalone --nproc_per_node=2 compile_precompute.py  

import logging
import os
from typing import Callable, Optional, Tuple, List, Tuple
from torchao.float8.float8_linear import Float8Linear
import math
import torch
import torch.distributed as dist
from torch.distributed._composable.fsdp import fully_shard
from torch.distributed._tensor import DTensor
from torchao.float8.float8_utils import EPS
from torchao.float8.fsdp_utils import WeightWithDynamicFloat8CastTensor
from torchao.float8.config import CastConfig, Float8LinearConfig, ScalingType
from torchao.float8.float8_linear_utils import convert_to_float8_training
from torch.testing._internal.distributed._tensor.common_dtensor import (
    ModelArgs,
    Transformer,
    TransformerBlock,
)
import torch.autograd.profiler as profiler


# @torch.compile()
@torch.compile(mode="reduce-overhead")
def compute_amaxes_compile(weights: List[DTensor]):
    abs_weights = torch._foreach_abs(weights)  # S0
    # abs_weights = [torch.abs(p) for p in weights]
    partial_amaxes = [torch.max(a) for a in abs_weights]  # P
    amax_tensor = torch.vstack(partial_amaxes)
    amax_tensor = torch.clamp(amax_tensor, EPS)  # R
    amaxes = torch.split(amax_tensor, 1)  # R
    return amaxes

def compute_amaxes_eager(weights: List[DTensor]):
    max_weights = torch._foreach_norm(weights, ord=math.inf)
    amax_tensor = torch.vstack(max_weights)
    amax_tensor = torch.clamp(amax_tensor, EPS)  # R
    amaxes = torch.split(amax_tensor, 1)  # R
    return amaxes

def main():
    dist.init_process_group(backend="nccl")
    gpu_id = int(os.environ["LOCAL_RANK"])
    device = f"cuda:{gpu_id}"
    torch.cuda.set_device(device)
    torch.manual_seed(42)
    args = ModelArgs(
        n_layers=3,
        dim=768,
        n_heads=12,
        dropout_p=0.0,
        weight_tying=False,
        vocab_size=32,
    )
    module = Transformer(args).cuda()
    module = convert_to_float8_training(
        module,
        config=Float8LinearConfig(
            enable_fsdp_float8_all_gather=True,
            cast_config_weight=CastConfig(scaling_type=ScalingType.DYNAMIC),
        ),
    )
    for submodule in module.modules():
        if isinstance(submodule, TransformerBlock):
            fully_shard(submodule)
    fully_shard(module)
    wait, warmup, active = 0, 1, 2
    num_steps = wait + warmup + active
    rank = torch.distributed.get_rank()

    local_inp = torch.randint(
        0, module.tok_embeddings.weight.size(0), (16, 16), device="cuda"
    )
    loss = module(local_inp).sum()
    loss.backward()
    float8_linears: List[Float8Linear] = [
        m
        for m in module.modules()
        if isinstance(m, Float8Linear)
        and isinstance(m.weight, DTensor)
        and isinstance(m.weight._local_tensor, WeightWithDynamicFloat8CastTensor)
    ]
    weights: List[DTensor] = [float8_linear.weight for float8_linear in float8_linears]
    with torch.profiler.profile(
        activities=[
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.CUDA,
        ],
        schedule=torch.profiler.schedule(
            wait=wait, warmup=warmup, active=active, repeat=1, skip_first=1
        ),
        on_trace_ready=(
            torch.profiler.tensorboard_trace_handler("./profiles")
            if not rank  # only save on rank 0
            else None
        ),
        record_shapes=True,
        profile_memory=False,
        with_stack=False,  # incurs an additional overhead; disable if not needed
        with_flops=False,
        with_modules=False,  # only for torchscript modules at the moment
        experimental_config=torch.profiler._ExperimentalConfig(
            enable_cuda_sync_events=True
        ),
    ) as prof:
        for iter_idx in range(10):
            torch.distributed.barrier()
            start_event = torch.cuda.Event(enable_timing=True)
            end_event = torch.cuda.Event(enable_timing=True)
            with profiler.record_function("compute_amaxes_eager"):
                start_event.record()                
                amaxes_eager = compute_amaxes_eager(weights)
                end_event.record()
            torch.cuda.synchronize()
            eager_gpu_time = start_event.elapsed_time(end_event)
            
            torch.distributed.barrier()
            start_event = torch.cuda.Event(enable_timing=True)
            end_event = torch.cuda.Event(enable_timing=True)
            with profiler.record_function("compute_amaxes_compile"):
                start_event.record()                
                amaxes_compile = compute_amaxes_compile(weights)
                end_event.record()
            torch.cuda.synchronize()
            compile_gpu_time = start_event.elapsed_time(end_event)
            if torch.distributed.get_rank() == 0:
                print(f"{iter_idx=} {eager_gpu_time=} {compile_gpu_time=}")
            prof.step()    


if __name__ == "__main__":
    main()


```

### Alternatives

_No response_

### Additional context

_No response_

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @zhaojuanmao @mrshenli @rohan-varma @chauhang @penguinwu @pritamdamania87 @satgera @gqchen @aazzolini @osalpekar @jiayisuse @XilunWu @tianyu-l @yf225 @ezyang @anijain2305