togethercomputer/ParallelKernelBench_Problems huggingface dataset:
# 25_importance_sampling_loss
import torch
import torch.nn.functional as F
import torch.distributed as dist
from typing import Tuple, Any

def solution(
    hidden_states: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.Tensor,
    old_logprobs: torch.Tensor,
    advantages: torch.Tensor,
    ignore_index: int = -100,
) -> Tuple[torch.Tensor, Any, torch.Tensor, torch.Tensor, torch.Tensor]:
    logits = F.linear(hidden_states, weight)
    logits_flat = logits.view(-1, logits.size(-1))
    labels_flat = labels.view(-1)
    
    per_token_ce = F.cross_entropy(logits_flat, labels_flat, ignore_index=ignore_index, reduction='none')
    
    new_logprobs_flat = -per_token_ce.detach()
    old_logprobs_flat = old_logprobs.view(-1)
    advantages_flat = advantages.view(-1)
    
    valid_mask = (labels_flat != ignore_index)
    n_valid_local = valid_mask.sum().float()
    
    n_valid_global = n_valid_local.clone()
    dist.all_reduce(n_valid_global, op=dist.ReduceOp.SUM)
    n_valid_global_clamped = n_valid_global.clamp(min=1.0)
    
    delta = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0).clamp(min=-20.0, max=20.0)
    ratio = torch.exp(delta)
    
    per_token_pg = -(ratio * advantages_flat).masked_fill(~valid_mask, 0.0)
    local_pg_sum = per_token_pg.sum()
    
    global_pg_sum = local_pg_sum.clone()
    dist.all_reduce(global_pg_sum, op=dist.ReduceOp.SUM)
    true_pg = global_pg_sum / n_valid_global_clamped
    
    w = (ratio.detach() * advantages_flat).masked_fill(~valid_mask, 0.0)
    local_surrogate_sum = (w * per_token_ce).sum()
    
    surrogate = local_surrogate_sum / n_valid_global_clamped
    
    loss = true_pg.detach() + surrogate - surrogate.detach()
    
    ratio_valid = ratio.masked_fill(~valid_mask, 0.0)
    sum_ratio_local = ratio_valid.sum()
    dist.all_reduce(sum_ratio_local, op=dist.ReduceOp.SUM)
    ratio_mean = sum_ratio_local / n_valid_global_clamped
    
    ratio_for_min = ratio.masked_fill(~valid_mask, float('inf'))
    min_ratio_local = ratio_for_min.min() if n_valid_local > 0 else torch.tensor(float('inf'), device=ratio.device)
    dist.all_reduce(min_ratio_local, op=dist.ReduceOp.MIN)
    
    ratio_for_max = ratio.masked_fill(~valid_mask, float('-inf'))
    max_ratio_local = ratio_for_max.max() if n_valid_local > 0 else torch.tensor(float('-inf'), device=ratio.device)
    dist.all_reduce(max_ratio_local, op=dist.ReduceOp.MAX)
    
    k3_local = (ratio - delta - 1.0).masked_fill(~valid_mask, 0.0).sum()
    dist.all_reduce(k3_local, op=dist.ReduceOp.SUM)
    k3_mean = k3_local / n_valid_global_clamped
    
    entropy_local = per_token_ce.detach().masked_fill(~valid_mask, 0.0).sum()
    dist.all_reduce(entropy_local, op=dist.ReduceOp.SUM)
    entropy_mean = entropy_local / n_valid_global_clamped
    
    metrics = torch.stack([ratio_mean, min_ratio_local, max_ratio_local, k3_mean, entropy_mean])
    
    per_token_logprobs = new_logprobs_flat.view_as(labels)
    per_token_loss = per_token_pg.view_as(labels)
    
    return loss, None, per_token_logprobs, per_token_loss, metrics
# 26_moe_token_preprocess
from typing import List, Optional, Tuple

import torch
import torch.distributed as dist


def _preprocess_impl(
    expert_mask: torch.Tensor,
    num_experts: int,
    ep_group: dist.ProcessGroup,
) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]:
    ep_size = ep_group.size()
    num_local_experts = num_experts // ep_size
    rank = dist.get_rank(ep_group)
    num_local_tokens_per_expert = expert_mask.sum(dim=(1, 2))

    input_splits = num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(dim=1).tolist()

    num_global_tokens_per_expert = torch.empty(
        ep_size,
        num_local_tokens_per_expert.size(0),
        dtype=num_local_tokens_per_expert.dtype,
        device=num_local_tokens_per_expert.device,
    )
    dist.all_gather_into_tensor(num_global_tokens_per_expert, num_local_tokens_per_expert, group=ep_group)

    start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts
    num_global_tokens_per_local_expert = num_global_tokens_per_expert[:, start_idx:end_idx].contiguous()

    output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist()

    num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to(
        torch.device("cpu"), non_blocking=True
    )
    num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts).to(
        torch.device("cpu"), non_blocking=True
    )

    return input_splits, output_splits, num_global_tokens_per_local_expert, num_global_sum_tokens_per_local_expert


def solution(
    expert_mask: torch.Tensor,
    num_experts: int,
    group: Optional[dist.ProcessGroup] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    group = group or dist.group.WORLD
    input_splits, output_splits, per_local_expert, sum_per_local_expert = _preprocess_impl(
        expert_mask=expert_mask, num_experts=num_experts, ep_group=group
    )
    input_splits_t = torch.tensor(input_splits, dtype=torch.long)
    output_splits_t = torch.tensor(output_splits, dtype=torch.long)
    return (input_splits_t, output_splits_t, per_local_expert, sum_per_local_expert)
# 27_moe_all2all_primitive
from typing import List, Optional, Union

import torch
import torch.distributed as dist


def solution(
    local_tensor: torch.Tensor,
    input_split_sizes: Optional[Union[List[int], torch.Tensor]] = None,
    output_split_sizes: Optional[Union[List[int], torch.Tensor]] = None,
    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 local_tensor.contiguous()

    local_tensor = local_tensor.contiguous()
    if output_split_sizes is None:
        output = torch.empty_like(local_tensor)
    else:
        out_size = sum(output_split_sizes) if isinstance(output_split_sizes, list) else int(output_split_sizes.sum().item())
        output = torch.empty(
            (out_size, local_tensor.size(1)),
            dtype=local_tensor.dtype,
            device=local_tensor.device,
        )
    dist.all_to_all_single(
        output,
        local_tensor,
        output_split_sizes=output_split_sizes,
        input_split_sizes=input_split_sizes,
        group=group,
    )
    return output
# 30_moe_epgroupgemm_lora_backward
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
# 33_ulysses_all_to_all_tensor_primitive
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()
# 34_ulysses_all_gather_into_tensor_primitive
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
# 35_ulysses_all_gather_variable_primitive
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()
# 36_ulysses_gather_seq_scatter_heads
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
# 37_ulysses_gather_heads_scatter_seq
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)
# 40_ddp
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
# 41_zero1_optimizer_shard
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)
# 42_zero2_optimizer_shard_grad
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)
# 43_fused_adam_grad_unshard_allgather
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
# 44_quantized_grad_allreduce
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)
# 45_reducescatter_fused_rmsnorm
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)