willychan21/ParallelKernelBench_Problems huggingface dataset:
from typing import Optional, Tuple

import torch
import torch.distributed as dist


def solution(
    grad_fc1_1_lora_A: torch.Tensor,
    grad_fc1_2_lora_A: torch.Tensor,
    grad_fc2_lora_B: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    group = group or dist.group.WORLD
    dist.all_reduce(grad_fc1_1_lora_A, op=dist.ReduceOp.SUM, group=group)
    dist.all_reduce(grad_fc1_2_lora_A, op=dist.ReduceOp.SUM, group=group)
    dist.all_reduce(grad_fc2_lora_B, op=dist.ReduceOp.SUM, group=group)
    return grad_fc1_1_lora_A, grad_fc1_2_lora_A, grad_fc2_lora_B
from typing import Optional

import torch
import torch.distributed as dist


def solution(
    x: torch.Tensor,
    scatter_dim: int,
    gather_dim: int,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return x.contiguous()

    x = x.contiguous()
    input_list = [t.contiguous() for t in torch.tensor_split(x, world_size, scatter_dim)]
    output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
    dist.all_to_all(output_list, input_list, group=group)
    return torch.cat(output_list, dim=gather_dim).contiguous()
from typing import Optional

import torch
import torch.distributed as dist


def solution(
    x: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return x.contiguous()

    x = x.contiguous()
    dim_size = list(x.size())
    dim_size[0] = dim_size[0] * world_size
    output = torch.empty(dim_size, dtype=x.dtype, device=x.device)
    dist.all_gather_into_tensor(output, x, group=group)
    return output
from typing import List, Optional, Tuple

import torch
import torch.distributed as dist


def solution(
    x: torch.Tensor,
    gather_dim: int,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return x.contiguous()

    device = x.device
    dtype = x.dtype
    x = x.contiguous()
    x_size = torch.tensor(x.size(), dtype=torch.int64, device=device)
    size_list = [torch.zeros(x_size.size(), dtype=torch.int64, device=device) for _ in range(world_size)]
    dist.all_gather(size_list, x_size, group=group)
    tensor_list = [
        torch.empty(torch.Size(size_list[i].tolist()), dtype=dtype, device=device)
        for i in range(world_size)
    ]
    dist.all_gather(tensor_list, x, group=group)
    return torch.cat(tensor_list, dim=gather_dim).contiguous()
from typing import Optional

import torch
import torch.distributed as dist
from torch.distributed import ProcessGroup


def _all_to_all(
    local_input: torch.Tensor,
    scatter_dim: int,
    gather_dim: int,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    seq_world_size = dist.get_world_size(group)
    input_list = [t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)]
    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]
    dist.all_to_all(output_list, input_list, group=group)
    return torch.cat(output_list, dim=gather_dim).contiguous()


def _all_to_all_single(
    x: torch.Tensor,
    scatter_dim: int,
    gather_dim: int,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    sp_world_size = dist.get_world_size(group)
    assert scatter_dim <= 1 and gather_dim <= 1
    if scatter_dim != 0:
        gather_dim_bef = x.shape[gather_dim]
        scatter_dim_bef = x.shape[scatter_dim]
        x = (
            x.reshape(
                [gather_dim_bef, sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:])
            )
            .transpose(0, 1)
            .reshape(
                [gather_dim_bef * sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:])
            )
            .contiguous()
        )
    output = torch.empty_like(x)
    dist.all_to_all_single(output, x.contiguous(), group=group)
    if scatter_dim == 0:
        output = torch.cat(output.split(x.size(0) // sp_world_size), dim=gather_dim)
    return output


def _all_to_all_tensor(
    x: torch.Tensor,
    scatter_dim: int,
    gather_dim: int,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    if scatter_dim <= 1 and gather_dim <= 1:
        return _all_to_all_single(x, scatter_dim, gather_dim, group)
    return _all_to_all(x, scatter_dim, gather_dim, group)


def solution(
    x: torch.Tensor,
    seq_dim: int,
    head_dim: int,
    group: Optional[ProcessGroup] = None,
    unpadded_dim_size: int = 0,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    sp_world = dist.get_world_size(group)
    x = _all_to_all_tensor(x, scatter_dim=head_dim, gather_dim=seq_dim, group=group)
    if unpadded_dim_size and unpadded_dim_size % sp_world != 0:
        padding_size = x.size(seq_dim) - unpadded_dim_size
        slc = [slice(None)] * x.dim()
        slc[seq_dim] = slice(0, -padding_size)
        x = x[tuple(slc)]
    return x
from typing import Optional

import torch
import torch.distributed as dist
from torch.distributed import ProcessGroup


def _pad_tensor(x: torch.Tensor, dim: int, padding_size: int, padding_value: int = 0) -> torch.Tensor:
    shape = list(x.shape)
    shape[dim] = padding_size
    pad = torch.full(shape, padding_value, dtype=x.dtype, device=x.device)
    return torch.cat([x, pad], dim=dim)


def _all_to_all(
    local_input: torch.Tensor,
    scatter_dim: int,
    gather_dim: int,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    seq_world_size = dist.get_world_size(group)
    input_list = [t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)]
    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]
    dist.all_to_all(output_list, input_list, group=group)
    return torch.cat(output_list, dim=gather_dim).contiguous()


def _all_to_all_single(
    x: torch.Tensor,
    scatter_dim: int,
    gather_dim: int,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    sp_world_size = dist.get_world_size(group)
    assert scatter_dim <= 1 and gather_dim <= 1
    if scatter_dim != 0:
        gather_dim_bef = x.shape[gather_dim]
        scatter_dim_bef = x.shape[scatter_dim]
        x = (
            x.reshape(
                [gather_dim_bef, sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:])
            )
            .transpose(0, 1)
            .reshape(
                [gather_dim_bef * sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:])
            )
            .contiguous()
        )
    output = torch.empty_like(x)
    dist.all_to_all_single(output, x.contiguous(), group=group)
    if scatter_dim == 0:
        output = torch.cat(output.split(x.size(0) // sp_world_size), dim=gather_dim)
    return output


def _all_to_all_tensor(
    x: torch.Tensor,
    scatter_dim: int,
    gather_dim: int,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    if scatter_dim <= 1 and gather_dim <= 1:
        return _all_to_all_single(x, scatter_dim, gather_dim, group)
    return _all_to_all(x, scatter_dim, gather_dim, group)


def solution(
    x: torch.Tensor,
    seq_dim: int,
    head_dim: int,
    group: Optional[ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    dim_size = x.size(seq_dim)
    sp_world = dist.get_world_size(group)
    if dim_size % sp_world != 0:
        padding_size = sp_world - (dim_size % sp_world)
        x = _pad_tensor(x, seq_dim, padding_size)
    return _all_to_all_tensor(x, scatter_dim=seq_dim, gather_dim=head_dim, group=group)
from __future__ import annotations

import math

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors


def solution(
    X_local: Tensor,
    y_local: Tensor,
    W1: Tensor,
    b1: Tensor,
    W2: Tensor,
    b2: Tensor,
    exp_avg_W1: Tensor,
    exp_avg_b1: Tensor,
    exp_avg_W2: Tensor,
    exp_avg_b2: Tensor,
    exp_avg_sq_W1: Tensor,
    exp_avg_sq_b1: Tensor,
    exp_avg_sq_W2: Tensor,
    exp_avg_sq_b2: Tensor,
    lr: float,
    beta1: float,
    beta2: float,
    eps: float,
    step: int,
) -> tuple[Tensor, ...]:
    world_size = dist.get_world_size()

    params = [W1, b1, W2, b2]
    exp_avg = [exp_avg_W1, exp_avg_b1, exp_avg_W2, exp_avg_b2]
    exp_avg_sq = [exp_avg_sq_W1, exp_avg_sq_b1, exp_avg_sq_W2, exp_avg_sq_b2]

    flat_params = _flatten_dense_tensors(params)
    dist.broadcast(flat_params, src=0)
    broadcast_params = _unflatten_dense_tensors(flat_params, params)
    params = [t.detach().requires_grad_(True) for t in broadcast_params]

    flat_m = _flatten_dense_tensors(exp_avg)
    dist.broadcast(flat_m, src=0)
    exp_avg = list(_unflatten_dense_tensors(flat_m, exp_avg))

    flat_v = _flatten_dense_tensors(exp_avg_sq)
    dist.broadcast(flat_v, src=0)
    exp_avg_sq = list(_unflatten_dense_tensors(flat_v, exp_avg_sq))

    h = F.relu(F.linear(X_local, params[0], params[1]))
    out = F.linear(h, params[2], params[3])
    loss = F.mse_loss(out, y_local)
    loss.backward()

    grads = [p.grad for p in params]
    flat_grad = _flatten_dense_tensors(grads)
    dist.all_reduce(flat_grad, op=dist.ReduceOp.SUM)
    flat_grad.div_(world_size)
    avg_grads = _unflatten_dense_tensors(flat_grad, grads)
    for p, g in zip(params, avg_grads):
        p.grad.copy_(g)

    assert step >= 1
    bc1 = 1.0 - math.pow(beta1, step)
    bc2 = 1.0 - math.pow(beta2, step)

    for p, m_buf, v_buf in zip(params, exp_avg, exp_avg_sq):
        g = p.grad
        m_buf.mul_(beta1).add_(g, alpha=1.0 - beta1)
        v_buf.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)
        m_hat = m_buf / bc1
        v_hat = v_buf / bc2
        denom = v_hat.sqrt().add(eps)
        p.data.add_(m_hat.div(denom).mul(-lr))

    out_tensors = tuple(list(params) + exp_avg + exp_avg_sq)
    return out_tensors
from __future__ import annotations

import math

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors


def solution(
    X_local: Tensor,
    y_local: Tensor,
    W1: Tensor,
    b1: Tensor,
    W2: Tensor,
    b2: Tensor,
    exp_avg_part: Tensor,
    exp_avg_sq_part: Tensor,
    lr: float,
    beta1: float,
    beta2: float,
    eps: float,
    step: int,
) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
    world_size = dist.get_world_size()
    rank = dist.get_rank()

    templates = [W1, b1, W2, b2]

    flat_p = _flatten_dense_tensors(templates)
    dist.broadcast(flat_p, src=0)

    param_views = _unflatten_dense_tensors(flat_p, templates)
    params = [t.detach().requires_grad_(True) for t in param_views]

    part = exp_avg_part.numel()
    assert flat_p.numel() == part * world_size

    m_part = exp_avg_part.clone()
    v_part = exp_avg_sq_part.clone()

    h = F.relu(F.linear(X_local, params[0], params[1]))
    out = F.linear(h, params[2], params[3])
    loss = F.mse_loss(out, y_local)
    loss.backward()

    flat_g = _flatten_dense_tensors([p.grad for p in params])
    dist.all_reduce(flat_g, op=dist.ReduceOp.SUM)
    flat_g.div_(world_size)

    start = rank * part
    g_part = flat_g[start : start + part]
    w_part = flat_p[start : start + part].clone()

    assert step >= 1
    bc1 = 1.0 - math.pow(beta1, step)
    bc2 = 1.0 - math.pow(beta2, step)

    m_part.mul_(beta1).add_(g_part, alpha=1.0 - beta1)
    v_part.mul_(beta2).addcmul_(g_part, g_part, value=1.0 - beta2)
    m_hat = m_part / bc1
    v_hat = v_part / bc2
    w_part.add_(m_hat.div(v_hat.sqrt().add(eps)).mul(-lr))

    gathered = torch.empty_like(flat_p)
    dist.all_gather_into_tensor(gathered, w_part.contiguous())
    flat_p.copy_(gathered)

    out_params = _unflatten_dense_tensors(flat_p, templates)
    return (*out_params, m_part, v_part)
from __future__ import annotations

import math

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors


def solution(
    X_local: Tensor,
    y_local: Tensor,
    W1: Tensor,
    b1: Tensor,
    W2: Tensor,
    b2: Tensor,
    exp_avg_part: Tensor,
    exp_avg_sq_part: Tensor,
    lr: float,
    beta1: float,
    beta2: float,
    eps: float,
    step: int,
) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
    world_size = dist.get_world_size()
    rank = dist.get_rank()

    templates = [W1, b1, W2, b2]
    flat_p = _flatten_dense_tensors(templates)
    dist.broadcast(flat_p, src=0)

    param_views = _unflatten_dense_tensors(flat_p, templates)
    params = [t.detach().requires_grad_(True) for t in param_views]

    part = exp_avg_part.numel()
    assert flat_p.numel() == part * world_size

    m_part = exp_avg_part.clone()
    v_part = exp_avg_sq_part.clone()

    h = F.relu(F.linear(X_local, params[0], params[1]))
    out = F.linear(h, params[2], params[3])
    loss = F.mse_loss(out, y_local)
    loss.backward()

    flat_g = _flatten_dense_tensors([p.grad for p in params]).contiguous()
    g_part = torch.empty(part, dtype=flat_g.dtype, device=flat_g.device)
    dist.reduce_scatter_tensor(g_part, flat_g, op=dist.ReduceOp.SUM)
    g_part.div_(world_size)

    start = rank * part
    w_part = flat_p[start : start + part].clone()

    assert step >= 1
    bc1 = 1.0 - math.pow(beta1, step)
    bc2 = 1.0 - math.pow(beta2, step)

    m_part.mul_(beta1).add_(g_part, alpha=1.0 - beta1)
    v_part.mul_(beta2).addcmul_(g_part, g_part, value=1.0 - beta2)
    m_hat = m_part / bc1
    v_hat = v_part / bc2
    w_part.add_(m_hat.div(v_hat.sqrt().add(eps)).mul(-lr))

    gathered = torch.empty_like(flat_p)
    dist.all_gather_into_tensor(gathered, w_part.contiguous())
    flat_p.copy_(gathered)

    out_params = _unflatten_dense_tensors(flat_p, templates)
    return (*out_params, m_part, v_part)
from __future__ import annotations

import math

import torch
import torch.distributed as dist
from torch import Tensor


@torch.no_grad()
def solution(
    grad_shard: Tensor,
    master_shard: Tensor,
    exp_avg: Tensor,
    exp_avg_sq: Tensor,
    lr: float,
    beta1: float,
    beta2: float,
    eps: float,
    step: int,
) -> Tensor:
    world_size = dist.get_world_size()

    assert step >= 1
    p = grad_shard.numel()
    assert p > 0

    m = exp_avg.clone()
    v = exp_avg_sq.clone()
    w = master_shard.clone()

    bc1 = 1.0 - math.pow(beta1, step)
    bc2 = 1.0 - math.pow(beta2, step)

    g = grad_shard
    m.mul_(beta1).add_(g, alpha=1.0 - beta1)
    v.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)
    m_hat = m / bc1
    v_hat = v / bc2
    w.add_(m_hat.div(v_hat.sqrt().add(eps)).mul(-lr))

    gathered = torch.empty(world_size * p, dtype=w.dtype, device=w.device)
    dist.all_gather_into_tensor(gathered, w.contiguous())
    return gathered
from __future__ import annotations

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor


@torch.no_grad()
def _block_int8_quant_dequant(x_flat: Tensor, block_size: int) -> Tensor:
    n = x_flat.numel()
    if n == 0:
        return x_flat.clone()
    flat = x_flat.contiguous().reshape(-1)
    pad = (-n) % block_size
    if pad:
        flat = F.pad(flat, (0, pad))
    nb = flat.numel() // block_size
    xv = flat.view(nb, block_size)
    scales = xv.abs().amax(dim=1).float().clamp(min=1e-8) / 127.0
    q = (xv.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
    out = (q.float() * scales.unsqueeze(1)).reshape(-1)
    return out[:n]


@torch.no_grad()
def solution(
    flat_grad: Tensor,
    block_size: int,
) -> Tensor:
    assert block_size >= 1

    world_size = dist.get_world_size()
    orig_shape = flat_grad.shape
    x = flat_grad.reshape(-1)

    rec = _block_int8_quant_dequant(x, block_size)
    acc = rec.float()
    dist.all_reduce(acc, op=dist.ReduceOp.SUM)
    acc.div_(world_size)

    return acc.to(dtype=flat_grad.dtype).reshape(orig_shape)
from __future__ import annotations

import torch
import torch.distributed as dist
from torch import Tensor


@torch.no_grad()
def solution(
    rs_input_1d: Tensor,
    gamma: Tensor,
    eps: float,
) -> Tensor:
    world_size = dist.get_world_size()
    n = rs_input_1d.numel()
    chunk = n // world_size

    hidden = gamma.numel()
    assert chunk % hidden == 0, f"chunk ({chunk}) must divide hidden ({hidden})"
    rows = chunk // hidden

    out_flat = torch.empty(chunk, dtype=rs_input_1d.dtype, device=rs_input_1d.device)
    dist.reduce_scatter_tensor(out_flat, rs_input_1d.contiguous(), op=dist.ReduceOp.SUM)
    out_flat.div_(world_size)

    x = out_flat.view(rows, hidden).float()
    gn = gamma.float()
    rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True).add(eps))
    y = x * rms * gn
    return y.to(dtype=rs_input_1d.dtype)
from __future__ import annotations

import math

import torch
from torch import Tensor


@torch.no_grad()
def solution(
    flat_param_shard: Tensor,
    flat_grad_shard: Tensor,
    exp_avg_shard: Tensor,
    exp_avg_sq_shard: Tensor,
    lr: float,
    beta1: float,
    beta2: float,
    eps: float,
    weight_decay: float,
    step: int,
) -> tuple[Tensor, Tensor, Tensor]:
    assert step >= 1

    m = exp_avg_shard.clone()
    v = exp_avg_sq_shard.clone()
    g = flat_grad_shard
    theta = flat_param_shard.clone()

    m.mul_(beta1).add_(g, alpha=1.0 - beta1)
    v.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)
    bc1 = 1.0 - math.pow(beta1, step)
    bc2 = 1.0 - math.pow(beta2, step)
    m_hat = m / bc1
    v_hat = v / bc2
    denom = v_hat.sqrt().add(eps)

    theta.add_(m_hat.div(denom), alpha=-lr)
    theta.add_(flat_param_shard, alpha=-lr * weight_decay)

    return theta, m, v
from __future__ import annotations

import math
from typing import Sequence

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors


def solution(
    X_local: Tensor,
    y_local: Tensor,
    flat_param_shard: Tensor,
    param_shapes: Sequence[tuple[int, ...]],
    exp_avg_shard: Tensor,
    exp_avg_sq_shard: Tensor,
    lr: float,
    beta1: float,
    beta2: float,
    eps: float,
    weight_decay: float,
    step: int,
) -> tuple[Tensor, Tensor, Tensor]:
    assert step >= 1

    world_size = dist.get_world_size()
    p = flat_param_shard.numel()

    device = flat_param_shard.device
    dtype = flat_param_shard.dtype

    templates = [torch.zeros(shape, dtype=dtype, device=device) for shape in param_shapes]
    full_flat = torch.empty(world_size * p, dtype=dtype, device=device)
    dist.all_gather_into_tensor(full_flat, flat_param_shard.contiguous())

    params_f = _unflatten_dense_tensors(full_flat, templates)
    params = [t.detach().requires_grad_(True) for t in params_f]

    h = F.relu(F.linear(X_local, params[0], params[1]))
    out = F.linear(h, params[2], params[3])
    loss = F.mse_loss(out, y_local)
    loss.backward()

    flat_g = _flatten_dense_tensors([x.grad for x in params])
    g_shard = torch.empty(p, dtype=flat_g.dtype, device=flat_g.device)
    dist.reduce_scatter_tensor(g_shard, flat_g.contiguous(), op=dist.ReduceOp.SUM)
    g_shard.div_(world_size)

    m = exp_avg_shard.clone()
    v = exp_avg_sq_shard.clone()
    theta = flat_param_shard.clone()

    m.mul_(beta1).add_(g_shard, alpha=1.0 - beta1)
    v.mul_(beta2).addcmul_(g_shard, g_shard, value=1.0 - beta2)
    bc1 = 1.0 - math.pow(beta1, step)
    bc2 = 1.0 - math.pow(beta2, step)
    m_hat = m / bc1
    v_hat = v / bc2
    denom = v_hat.sqrt().add(eps)

    theta.add_(m_hat.div(denom), alpha=-lr)
    theta.add_(flat_param_shard, alpha=-lr * weight_decay)

    return theta, m, v
from __future__ import annotations

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor


def _make_tp_fsdp_groups(n_tp: int, n_fsdp: int, rank: int):
    tp_group = None
    fsdp_group = None
    for j in range(n_fsdp):
        ranks = [j * n_tp + ii for ii in range(n_tp)]
        g = dist.new_group(ranks)
        if rank in ranks:
            tp_group = g
    for i in range(n_tp):
        ranks = [jj * n_tp + i for jj in range(n_fsdp)]
        g = dist.new_group(ranks)
        if rank in ranks:
            fsdp_group = g
    assert tp_group is not None and fsdp_group is not None
    return tp_group, fsdp_group


def _gather_fsdp_concat_dim0(shard: Tensor, fsdp_group, parts: int) -> Tensor:
    lst = [torch.empty_like(shard) for _ in range(parts)]
    dist.all_gather(lst, shard.contiguous(), group=fsdp_group)
    return torch.cat(lst, dim=0)


def _gather_fsdp_concat_dim1(shard: Tensor, fsdp_group, parts: int) -> Tensor:
    lst = [torch.empty_like(shard) for _ in range(parts)]
    dist.all_gather(lst, shard.contiguous(), group=fsdp_group)
    return torch.cat(lst, dim=1)


@torch.no_grad()
def solution(
    x_local: Tensor,
    W1_shard: Tensor,
    W2_shard: Tensor,
    W3_shard: Tensor,
    n_tp: int,
    n_fsdp: int,
) -> Tensor:
    world_size = dist.get_world_size()
    rank = dist.get_rank()
    assert world_size == n_tp * n_fsdp, (
        f"world_size ({world_size}) must equal n_tp * n_fsdp ({n_tp} * {n_fsdp})"
    )

    tp_group, fsdp_group = _make_tp_fsdp_groups(n_tp, n_fsdp, rank)

    W1 = _gather_fsdp_concat_dim0(W1_shard, fsdp_group, n_fsdp)
    W2 = _gather_fsdp_concat_dim0(W2_shard, fsdp_group, n_fsdp)
    W3 = _gather_fsdp_concat_dim1(W3_shard, fsdp_group, n_fsdp)

    x1 = x_local @ W1
    x2 = x_local @ W2
    z = F.silu(x1) * x2
    y_partial = z @ W3

    y = y_partial.clone()
    dist.all_reduce(y, op=dist.ReduceOp.SUM, group=tp_group)
    return y
from __future__ import annotations

import torch
import torch.distributed as dist
from torch import Tensor

_FP8_E4M3_MAX = 448.0


@torch.no_grad()
def _update_amax_history(amax_history: Tensor, cur_abs_max: Tensor) -> Tensor:
    out = torch.roll(amax_history, shifts=-1, dims=0)
    out[-1] = cur_abs_max.to(dtype=out.dtype)
    return out


@torch.no_grad()
def _fp8_round_trip_bf16(x: Tensor, scale: Tensor) -> Tensor:
    xf = x.float()
    qs = xf / scale
    q = qs.to(torch.float8_e4m3fn)
    return (q.float() * scale).to(dtype=x.dtype)


@torch.no_grad()
def solution(flat_grads: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]:
    world_size = dist.get_world_size()
    n = flat_grads.numel()
    shard_elems = n // world_size

    cur_abs_max = flat_grads.abs().max().to(torch.float32)
    updated_hist = _update_amax_history(amax_history, cur_abs_max)

    scale = updated_hist.max().clamp(min=1e-12).to(torch.float32) / _FP8_E4M3_MAX
    recon = _fp8_round_trip_bf16(flat_grads, scale)

    out_shard = torch.empty(shard_elems, dtype=flat_grads.dtype, device=flat_grads.device)
    dist.reduce_scatter_tensor(out_shard, recon.contiguous(), op=dist.ReduceOp.SUM)
    out_shard.div_(world_size)

    return out_shard, updated_hist