170 lines
6.8 KiB
Python
170 lines
6.8 KiB
Python
import argparse
|
|
import inspect
|
|
import os
|
|
import pathlib
|
|
import random
|
|
import sys
|
|
from typing import Tuple
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
|
|
REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
import deep_gemm
|
|
from deep_gemm.utils.math import ceil_div
|
|
|
|
|
|
def init_test_dist(local_rank_arg: int = None) -> Tuple[int, int, dist.ProcessGroup]:
|
|
local_rank = local_rank_arg if local_rank_arg is not None else int(os.environ.get('LOCAL_RANK', '0'))
|
|
rank = int(os.environ.get('RANK', '0'))
|
|
world_size = int(os.environ.get('WORLD_SIZE', '1'))
|
|
master_addr = os.environ.get('MASTER_ADDR', '127.0.0.1')
|
|
master_port = int(os.environ.get('MASTER_PORT', '8362'))
|
|
|
|
torch.cuda.set_device(local_rank)
|
|
sig = inspect.signature(dist.init_process_group)
|
|
params = {
|
|
'backend': 'nccl',
|
|
'init_method': f'tcp://{master_addr}:{master_port}',
|
|
'world_size': world_size,
|
|
'rank': rank,
|
|
}
|
|
if 'device_id' in sig.parameters:
|
|
params['device_id'] = torch.device(f'cuda:{local_rank}')
|
|
dist.init_process_group(**params)
|
|
torch.set_default_device('cuda')
|
|
return rank, world_size, dist.new_group(list(range(world_size)))
|
|
|
|
|
|
def interleaved_to_natural_n(n: torch.Tensor, half_n: int, gran: int = 8) -> torch.Tensor:
|
|
pair_group = n // (2 * gran)
|
|
offset = n - pair_group * (2 * gran)
|
|
gate_n = pair_group * gran + offset
|
|
up_n = half_n + pair_group * gran + offset - gran
|
|
return torch.where(offset < gran, gate_n, up_n)
|
|
|
|
|
|
def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
|
|
torch.manual_seed(2027 + dist.get_rank())
|
|
l1_weights = (torch.randn(
|
|
(num_experts_per_rank, intermediate_hidden * 2, hidden),
|
|
dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn)
|
|
l2_weights = (torch.randn(
|
|
(num_experts_per_rank, hidden, intermediate_hidden),
|
|
dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn)
|
|
|
|
l1_weights_sf = torch.empty(
|
|
(num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128),
|
|
dtype=torch.float32, device='cuda')
|
|
for expert in range(num_experts_per_rank):
|
|
for n_group in range(l1_weights_sf.shape[1]):
|
|
for k_group in range(l1_weights_sf.shape[2]):
|
|
l1_weights_sf[expert, n_group, k_group] = 0.75 + 0.125 * n_group + 0.0625 * k_group
|
|
|
|
l2_weights_sf = torch.ones(
|
|
(num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128),
|
|
dtype=torch.float32, device='cuda')
|
|
return deep_gemm.transform_weights_for_mega_moe(
|
|
(l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf))
|
|
|
|
|
|
def reference_l1_accum(buffer: deep_gemm.SymmBuffer,
|
|
l1_weights: torch.Tensor,
|
|
l1_weights_sf: torch.Tensor,
|
|
hidden: int,
|
|
intermediate_hidden: int) -> torch.Tensor:
|
|
block_m = 128
|
|
block_n = 128
|
|
x_fp8 = buffer.l1_acts[:block_m]
|
|
x_sf = buffer.l1_acts_sf[:block_m, :hidden // 128]
|
|
w_fp8 = l1_weights[0, :block_n]
|
|
|
|
n = torch.arange(block_n, device='cuda')
|
|
natural_n = interleaved_to_natural_n(n, intermediate_hidden)
|
|
n_groups = natural_n // 128
|
|
|
|
expected = torch.zeros((block_m, block_n), dtype=torch.float32, device='cuda')
|
|
for k_group in range(hidden // 128):
|
|
k_start = k_group * 128
|
|
k_end = k_start + 128
|
|
x = x_fp8[:, k_start:k_end].to(torch.float32)
|
|
w = w_fp8[:, k_start:k_end].to(torch.float32)
|
|
sfa = x_sf[:, k_group].to(torch.float32)
|
|
sfb = l1_weights_sf[0, n_groups, k_group].to(torch.float32)
|
|
expected += (x @ w.t()) * sfa[:, None] * sfb[None, :]
|
|
return expected
|
|
|
|
|
|
def run_case(args: argparse.Namespace, group: dist.ProcessGroup, rank_idx: int, num_ranks: int) -> None:
|
|
hidden = args.hidden
|
|
intermediate_hidden = args.intermediate_hidden
|
|
num_tokens = args.num_tokens
|
|
num_topk = 1
|
|
num_experts = args.num_experts if args.num_experts is not None else num_ranks
|
|
num_experts_per_rank = num_experts // num_ranks
|
|
assert num_experts % num_ranks == 0
|
|
assert num_experts_per_rank >= 1
|
|
|
|
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
|
|
group, num_experts, args.num_max_tokens_per_rank, num_topk,
|
|
hidden, intermediate_hidden)
|
|
weights = make_weights(num_experts_per_rank, hidden, intermediate_hidden)
|
|
|
|
torch.manual_seed(1234 + rank_idx)
|
|
x = (torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn)
|
|
x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') + 0.75
|
|
local_global_expert = rank_idx * num_experts_per_rank
|
|
topk_idx = torch.full((num_tokens, num_topk), local_global_expert, dtype=torch.long, device='cuda')
|
|
topk_weights = torch.ones((num_tokens, num_topk), dtype=torch.float32, device='cuda')
|
|
|
|
buffer.x[:num_tokens].copy_(x)
|
|
buffer.x_sf[:num_tokens].copy_(x_sf)
|
|
buffer.topk_idx[:num_tokens].copy_(topk_idx)
|
|
buffer.topk_weights[:num_tokens].copy_(topk_weights)
|
|
torch.cuda.synchronize()
|
|
dist.barrier(group=group)
|
|
|
|
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
|
|
deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer)
|
|
torch.cuda.synchronize()
|
|
|
|
expected = reference_l1_accum(buffer, weights[0][0], weights[0][1], hidden, intermediate_hidden)
|
|
actual = buffer.l1_accum_debug
|
|
torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=args.rtol, atol=args.atol)
|
|
|
|
dist.barrier(group=group)
|
|
if rank_idx == 0:
|
|
max_diff = (actual - expected).abs().max().item()
|
|
print(f'[PASSED] Phase 3 L1 WGMMA single tile, max_diff={max_diff:.6f}', flush=True)
|
|
buffer.destroy()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 3 L1 WGMMA single-tile correctness')
|
|
parser.add_argument('--num-tokens', type=int, default=128)
|
|
parser.add_argument('--num-max-tokens-per-rank', type=int, default=384)
|
|
parser.add_argument('--hidden', type=int, default=256)
|
|
parser.add_argument('--intermediate-hidden', type=int, default=128)
|
|
parser.add_argument('--num-experts', type=int, default=None)
|
|
parser.add_argument('--local-rank', type=int, default=None)
|
|
parser.add_argument('--atol', type=float, default=1e-2)
|
|
parser.add_argument('--rtol', type=float, default=5e-2)
|
|
args = parser.parse_args()
|
|
|
|
rank_idx, num_ranks, group = init_test_dist(args.local_rank)
|
|
assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
|
|
assert args.num_tokens == 128
|
|
assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
|
|
|
|
random.seed(5678 + rank_idx)
|
|
run_case(args, group, rank_idx, num_ranks)
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|