diff --git a/benchmarks/benchmark_dispatcher_npu.py b/benchmarks/benchmark_dispatcher_npu.py new file mode 100644 index 00000000..eecb1aac --- /dev/null +++ b/benchmarks/benchmark_dispatcher_npu.py @@ -0,0 +1,463 @@ +""" +NPU AllToAll Dispatcher Benchmark +torchrun --nproc_per_node=8 benchmarks/benchmark_dispatcher_npu.py +""" + +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() + +import argparse +import os +import sys +import threading +import time +from pathlib import Path + +import psutil +import torch +import torch.distributed as dist + +REPO_ROOT = Path(__file__).resolve().parent +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) +if str(REPO_ROOT.parent) not in sys.path: + sys.path.insert(0, str(REPO_ROOT.parent)) + +from steptronoss.core.parallel_state import PM +from steptronoss.exp.base_exp import ParallelConfig +from steptronoss.model.ep_dispatcher.token_dispatcher import TokenDispatcher +from steptronoss.utils.optimizable import set_optimization + +# PARAM_SETS = [ +# { +# "ep_size": 8, +# "num_experts": 8, +# "seq_len": 4096, +# "hidden": 256, +# "topk": 2, +# "warmup": 5, +# "iters": 20, +# "check": True, +# }, +# ] + +PARAM_SETS = [ + { + "ep_size": 8, + "num_experts": 288, + "seq_len": 4096, + "hidden": 4096, + "topk": 8, + "warmup": 5, + "iters": 20, + "check": True, + }, +] + + +def _sync() -> None: + torch.npu.synchronize() + + +def _reset_peak_memory_stats() -> None: + if hasattr(torch.npu, "empty_cache"): + torch.npu.empty_cache() + if hasattr(torch.npu, "reset_peak_memory_stats"): + torch.npu.reset_peak_memory_stats() + + +def _peak_memory_stats_mb() -> tuple[float | None, float | None]: + peak_allocated = None + peak_reserved = None + if hasattr(torch.npu, "max_memory_allocated"): + peak_allocated = float(torch.npu.max_memory_allocated()) / (1024**2) + if hasattr(torch.npu, "max_memory_reserved"): + peak_reserved = float(torch.npu.max_memory_reserved()) / (1024**2) + return peak_allocated, peak_reserved + + +def _measure_host_memory(fn): + proc = psutil.Process() + stop = threading.Event() + peak_rss_mb = 0.0 + peak_host_used_mb = 0.0 + peak_host_percent = 0.0 + + def poll(): + nonlocal peak_rss_mb, peak_host_used_mb, peak_host_percent + while not stop.is_set(): + try: + rss_mb = proc.memory_info().rss / (1024**2) + vm = psutil.virtual_memory() + used_mb = vm.used / (1024**2) + percent = float(vm.percent) + peak_rss_mb = max(peak_rss_mb, rss_mb) + peak_host_used_mb = max(peak_host_used_mb, used_mb) + peak_host_percent = max(peak_host_percent, percent) + except Exception: + pass + stop.wait(0.05) + + thread = threading.Thread(target=poll, daemon=True) + thread.start() + try: + result = fn() + finally: + stop.set() + thread.join(timeout=1.0) + return result, peak_rss_mb, peak_host_used_mb, peak_host_percent + + +def _init_dist_and_mesh(ep_size: int) -> bool: + if not hasattr(torch, "npu") or not torch.npu.is_available(): + raise RuntimeError("NPU is required for benchmark_dispatcher_npu.py") + if not dist.is_available(): + raise RuntimeError("torch.distributed is not available") + + did_init = False + if dist.is_initialized(): + if dist.get_backend() != "hccl": + raise RuntimeError("Benchmark requires HCCL backend") + else: + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29500") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", str(ep_size)) + os.environ.setdefault("LOCAL_RANK", "0") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.npu.set_device(local_rank) + dist.init_process_group(backend="hccl") + did_init = True + + if dist.get_world_size() != ep_size: + raise RuntimeError(f"Benchmark assumes WORLD_SIZE={ep_size}") + + torch.npu.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + + PM.initialize(backend="hccl") + parallel_cfg = PM._cur_cfg or None + if parallel_cfg is None or PM.size_of("EP") != ep_size: + parallel_cfg = ParallelConfig() + parallel_cfg.tensor_model_parallel_size = 1 + parallel_cfg.pipeline_model_parallel_size = 1 + parallel_cfg.context_parallel_size = 1 + parallel_cfg.expert_model_parallel_size = ep_size + parallel_cfg.expert_tensor_parallel_size = 1 + parallel_cfg.virtual_pipeline_model_parallel_size = 1 + PM.set_mesh(parallel_cfg) + + return did_init + + +def _cleanup_dist(did_init: bool) -> None: + if did_init and dist.is_initialized(): + dist.destroy_process_group() + PM._all_groups.clear() + PM.parallels.clear() + PM.all_parallels.clear() + PM._stack.clear() + PM._cur_cfg = None + PM._rng_seeds.clear() + PM.rng_states.clear() + + +def _make_inputs(seq_len: int, hidden: int, topk: int, num_experts: int, device: torch.device): + hidden_states = torch.randn((seq_len, hidden), device=device, dtype=torch.bfloat16) + token_expert_ids = torch.randint(0, num_experts, (seq_len, topk), device=device, dtype=torch.int64) + if topk > 1: + for i in range(1, topk): + clash = token_expert_ids[:, i] == token_expert_ids[:, 0] + token_expert_ids[clash, i] = (token_expert_ids[clash, i] + i) % num_experts + token_expert_weights = torch.rand((seq_len, topk), device=device, dtype=torch.float32) + token_expert_weights = token_expert_weights / token_expert_weights.sum(dim=1, keepdim=True) + return hidden_states, token_expert_ids, token_expert_weights + + +def _max_abs_diff(x: torch.Tensor, y: torch.Tensor) -> float: + if x.dtype in (torch.int32, torch.int64, torch.int16, torch.int8, torch.uint8): + return float((x.to(torch.int64) - y.to(torch.int64)).abs().max().item()) + return float((x.float() - y.float()).abs().max().item()) + + +def _check_close(x: torch.Tensor, y: torch.Tensor, rtol: float, atol: float) -> bool: + try: + torch.testing.assert_close(x, y, rtol=rtol, atol=atol) + return True + except Exception: + return False + + +def _dispatcher_loss(output: torch.Tensor, recv_prob: torch.Tensor) -> torch.Tensor: + # Cover both hidden-state restore and prob communication paths in backward. + return output.float().sum() + recv_prob.float().sum() + + +def _run_dispatcher( + backend: str | None, + hidden_states: torch.Tensor, + token_expert_ids: torch.Tensor, + token_expert_weights: torch.Tensor, + num_experts: int, + warmup: int, + iters: int, +): + set_optimization(**{"TokenDispatcher": backend}) + dist.barrier() + _sync() + dist.barrier() + + dispatcher = TokenDispatcher("EP", num_experts=num_experts) + recv_hidden = recv_idx = recv_prob = out = None + _reset_peak_memory_stats() + for _ in range(warmup): + recv_hidden, recv_idx, recv_prob = dispatcher.dispatch(hidden_states, token_expert_ids, token_expert_weights) + out = dispatcher.combine(recv_hidden) + _sync() + dist.barrier() + + _reset_peak_memory_stats() + + def timed_run(): + nonlocal recv_hidden, recv_idx, recv_prob, out + start = time.perf_counter() + for _ in range(iters): + recv_hidden, recv_idx, recv_prob = dispatcher.dispatch( + hidden_states, token_expert_ids, token_expert_weights + ) + out = dispatcher.combine(recv_hidden) + _sync() + dist.barrier() + return (time.perf_counter() - start) * 1000.0 / iters + + ms_per_iter, peak_rss_mb, peak_host_used_mb, peak_host_percent = _measure_host_memory(timed_run) + peak_allocated_mb, peak_reserved_mb = _peak_memory_stats_mb() + + grad_hidden_states = hidden_states.detach().clone().requires_grad_(True) + grad_token_expert_weights = token_expert_weights.detach().clone().requires_grad_(True) + dispatcher_bw = TokenDispatcher("EP", num_experts=num_experts) + recv_prob_bw = out_bw = None + grad_hidden = grad_token_expert_weights_out = None + + for _ in range(warmup): + recv_hidden_bw, _, recv_prob_bw = dispatcher_bw.dispatch( + grad_hidden_states, + token_expert_ids, + grad_token_expert_weights, + ) + out_bw = dispatcher_bw.combine(recv_hidden_bw) + _dispatcher_loss(out_bw, recv_prob_bw).backward() + grad_hidden_states.grad = None + grad_token_expert_weights.grad = None + _sync() + dist.barrier() + + start = time.perf_counter() + for _ in range(iters): + recv_hidden_bw, _, recv_prob_bw = dispatcher_bw.dispatch( + grad_hidden_states, + token_expert_ids, + grad_token_expert_weights, + ) + out_bw = dispatcher_bw.combine(recv_hidden_bw) + _dispatcher_loss(out_bw, recv_prob_bw).backward() + grad_hidden = grad_hidden_states.grad.detach().clone() + grad_token_expert_weights_out = grad_token_expert_weights.grad.detach().clone() + grad_hidden_states.grad = None + grad_token_expert_weights.grad = None + _sync() + dist.barrier() + total_ms_per_iter = (time.perf_counter() - start) * 1000.0 / iters + + return ( + ms_per_iter, + total_ms_per_iter, + out, + recv_hidden, + recv_idx, + recv_prob, + out_bw.detach().clone(), + grad_hidden, + grad_token_expert_weights_out, + peak_allocated_mb, + peak_reserved_mb, + peak_rss_mb, + peak_host_used_mb, + peak_host_percent, + bool(getattr(dispatcher, "last_dispatch_used_fused_permute", False)), + bool(getattr(dispatcher, "last_combine_used_fused_unpermute", False)), + ) + + +def main() -> int: + parser = argparse.ArgumentParser() + defaults = dict(PARAM_SETS[0]) + + parser.add_argument("--ep-size", type=int, default=defaults["ep_size"]) + parser.add_argument("--num-experts", type=int, default=defaults["num_experts"]) + parser.add_argument("--seq-len", type=int, default=defaults["seq_len"]) + parser.add_argument("--hidden", type=int, default=defaults["hidden"]) + parser.add_argument("--topk", type=int, default=defaults["topk"]) + parser.add_argument("--warmup", type=int, default=defaults["warmup"]) + parser.add_argument("--iters", type=int, default=defaults["iters"]) + parser.add_argument("--rtol", type=float, default=1e-3) + parser.add_argument("--atol", type=float, default=1e-3) + parser.add_argument("--check", action="store_true", default=defaults["check"]) + args = parser.parse_args() + + did_init = False + try: + did_init = _init_dist_and_mesh(args.ep_size) + num_experts = args.num_experts + device = torch.device("npu", int(os.environ.get("LOCAL_RANK", "0"))) + hidden_states, token_expert_ids, token_expert_weights = _make_inputs( + args.seq_len, args.hidden, args.topk, num_experts, device + ) + + ( + base_fw_ms, + base_total_ms, + base_out, + base_recv_hidden, + base_recv_idx, + base_recv_prob, + base_out_bw, + base_grad_hidden, + base_grad_prob, + base_peak_alloc_mb, + base_peak_reserved_mb, + base_peak_rss_mb, + base_peak_host_used_mb, + base_peak_host_percent, + base_fused_permute, + base_fused_unpermute, + ) = _run_dispatcher( + None, + hidden_states, + token_expert_ids, + token_expert_weights, + num_experts, + args.warmup, + args.iters, + ) + + ( + npu_fw_ms, + npu_total_ms, + npu_out, + npu_recv_hidden, + npu_recv_idx, + npu_recv_prob, + npu_out_bw, + npu_grad_hidden, + npu_grad_prob, + npu_peak_alloc_mb, + npu_peak_reserved_mb, + npu_peak_rss_mb, + npu_peak_host_used_mb, + npu_peak_host_percent, + npu_fused_permute, + npu_fused_unpermute, + ) = _run_dispatcher( + "npu_alltoall", + hidden_states, + token_expert_ids, + token_expert_weights, + num_experts, + args.warmup, + args.iters, + ) + + torch.testing.assert_close(npu_recv_hidden, base_recv_hidden, rtol=args.rtol, atol=args.atol) + torch.testing.assert_close(npu_recv_idx, base_recv_idx, rtol=0, atol=0) + torch.testing.assert_close(npu_recv_prob, base_recv_prob, rtol=args.rtol, atol=args.atol) + torch.testing.assert_close(npu_out, base_out, rtol=args.rtol, atol=args.atol) + torch.testing.assert_close(npu_out_bw, base_out_bw, rtol=args.rtol, atol=args.atol) + torch.testing.assert_close(npu_grad_hidden, base_grad_hidden, rtol=args.rtol, atol=args.atol) + torch.testing.assert_close(npu_grad_prob, base_grad_prob, rtol=args.rtol, atol=args.atol) + + if PM.rank_in("EP") == 0: + total_tokens = args.seq_len * args.ep_size + base_bw_ms = base_total_ms - base_fw_ms + npu_bw_ms = npu_total_ms - npu_fw_ms + base_fw_tokens_per_s = total_tokens / (base_fw_ms / 1000.0) + base_bw_tokens_per_s = total_tokens / (base_bw_ms / 1000.0) + base_total_tokens_per_s = total_tokens / (base_total_ms / 1000.0) + npu_fw_tokens_per_s = total_tokens / (npu_fw_ms / 1000.0) + npu_bw_tokens_per_s = total_tokens / (npu_bw_ms / 1000.0) + npu_total_tokens_per_s = total_tokens / (npu_total_ms / 1000.0) + print( + f"[dispatcher_npu] ep={args.ep_size} num_experts={num_experts} " + f"seq_len={args.seq_len} hidden={args.hidden} topk={args.topk}" + ) + print( + f"[dispatcher_npu] baseline: fw_ms={base_fw_ms:.3f}, bw_ms={base_bw_ms:.3f}, " + f"total_ms={base_total_ms:.3f}, fw_tokens/s={base_fw_tokens_per_s:,.0f}, " + f"bw_tokens/s={base_bw_tokens_per_s:,.0f}, total_tokens/s={base_total_tokens_per_s:,.0f}" + ) + print( + f"[dispatcher_npu] baseline_peak_mem: allocated_mb={base_peak_alloc_mb:.1f}, " + f"reserved_mb={base_peak_reserved_mb:.1f}" + ) + print( + f"[dispatcher_npu] baseline_host_mem: process_rss_mb={base_peak_rss_mb:.1f}, " + f"host_used_mb={base_peak_host_used_mb:.1f}, host_used_percent={base_peak_host_percent:.1f}" + ) + print( + f"[dispatcher_npu] baseline_fused_ops: permute={base_fused_permute}, unpermute={base_fused_unpermute}" + ) + print( + f"[dispatcher_npu] npu_alltoall: fw_ms={npu_fw_ms:.3f}, bw_ms={npu_bw_ms:.3f}, " + f"total_ms={npu_total_ms:.3f}, fw_tokens/s={npu_fw_tokens_per_s:,.0f}, " + f"bw_tokens/s={npu_bw_tokens_per_s:,.0f}, total_tokens/s={npu_total_tokens_per_s:,.0f}" + ) + print( + f"[dispatcher_npu] speedup_vs_baseline: fw_throughput={npu_fw_tokens_per_s / base_fw_tokens_per_s:.2f}x, " + f"total_throughput={npu_total_tokens_per_s / base_total_tokens_per_s:.2f}x" + ) + print( + f"[dispatcher_npu] npu_alltoall_peak_mem: allocated_mb={npu_peak_alloc_mb:.1f}, " + f"reserved_mb={npu_peak_reserved_mb:.1f}" + ) + print( + f"[dispatcher_npu] npu_alltoall_host_mem: process_rss_mb={npu_peak_rss_mb:.1f}, " + f"host_used_mb={npu_peak_host_used_mb:.1f}, host_used_percent={npu_peak_host_percent:.1f}" + ) + print( + f"[dispatcher_npu] npu_alltoall_fused_ops: permute={npu_fused_permute}, unpermute={npu_fused_unpermute}" + ) + print("metric, close, max_abs_diff") + print( + f"recv_hidden, {_check_close(npu_recv_hidden, base_recv_hidden, args.rtol, args.atol)}, " + f"{_max_abs_diff(npu_recv_hidden, base_recv_hidden):.6f}" + ) + print( + f"recv_idx, {_check_close(npu_recv_idx, base_recv_idx, 0.0, 0.0)}, {_max_abs_diff(npu_recv_idx, base_recv_idx):.6f}" + ) + print( + f"recv_prob, {_check_close(npu_recv_prob, base_recv_prob, args.rtol, args.atol)}, " + f"{_max_abs_diff(npu_recv_prob, base_recv_prob):.6f}" + ) + print( + f"combine_out, {_check_close(npu_out, base_out, args.rtol, args.atol)}, {_max_abs_diff(npu_out, base_out):.6f}" + ) + print( + f"combine_out_bw, {_check_close(npu_out_bw, base_out_bw, args.rtol, args.atol)}, " + f"{_max_abs_diff(npu_out_bw, base_out_bw):.6f}" + ) + print( + f"grad_hidden, {_check_close(npu_grad_hidden, base_grad_hidden, args.rtol, args.atol)}, " + f"{_max_abs_diff(npu_grad_hidden, base_grad_hidden):.6f}" + ) + print( + f"grad_prob, {_check_close(npu_grad_prob, base_grad_prob, args.rtol, args.atol)}, " + f"{_max_abs_diff(npu_grad_prob, base_grad_prob):.6f}" + ) + return 0 + finally: + set_optimization(**{"TokenDispatcher": None}) + _cleanup_dist(did_init) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_grouped_gemm_npu.py b/benchmarks/benchmark_grouped_gemm_npu.py new file mode 100644 index 00000000..b1f8b363 --- /dev/null +++ b/benchmarks/benchmark_grouped_gemm_npu.py @@ -0,0 +1,224 @@ +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() + +import argparse +import time +from pathlib import Path + +import torch + +REPO_ROOT = Path(__file__).resolve().parent +if str(REPO_ROOT.parent) not in __import__("sys").path: + __import__("sys").path.insert(0, str(REPO_ROOT.parent)) + + +DEFAULT_PARAM_SETS = [ + { + "name": "moe_like_large", + "group_size": 36, + "batch_size": 3256, + "k": 4096, + "n": 2560, + "dtype": "bf16", + "warmup": 20, + "iters": 20, + "trans_b": True, + }, +] + + +def _dtype_from_name(name: str) -> torch.dtype: + table = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + } + if name not in table: + raise ValueError(f"Unsupported dtype: {name}") + return table[name] + + +def _sync(): + torch.npu.synchronize() + + +def _build_inputs(params: dict[str, object], device: torch.device): + dtype = _dtype_from_name(params["dtype"]) + group_size = int(params["group_size"]) + batch_size = int(params["batch_size"]) + k = int(params["k"]) + n = int(params["n"]) + trans_b = bool(params["trans_b"]) + + batch_sizes = torch.full((group_size,), batch_size, device=device, dtype=torch.int64) + total_m = int(batch_sizes.sum().item()) + + a = torch.randn(total_m, k, device=device, dtype=dtype, requires_grad=True) + if trans_b: + b = torch.randn(group_size, n, k, device=device, dtype=dtype, requires_grad=True) + else: + b = torch.randn(group_size, k, n, device=device, dtype=dtype, requires_grad=True) + return a, b, batch_sizes + + +def _run_baseline( + mat_a_flat: torch.Tensor, mat_b: torch.Tensor, batch_sizes: torch.Tensor, trans_b: bool +) -> torch.Tensor: + batch_sizes_list = batch_sizes.tolist() + outputs = [] + start = 0 + for i, size in enumerate(batch_sizes_list): + rhs = mat_b[i].t() if trans_b else mat_b[i] + outputs.append(mat_a_flat[start : start + size] @ rhs) + start += size + if outputs: + return torch.cat(outputs, dim=0) + return mat_a_flat.new_zeros((0, mat_b.shape[1] if trans_b else mat_b.shape[2])) + + +def _run_npu_gmm_v2( + mat_a_flat: torch.Tensor, mat_b: torch.Tensor, batch_sizes: torch.Tensor, trans_b: bool +) -> torch.Tensor: + try: + from mindspeed.ops.gmm import npu_gmm_v2 + except Exception as exc: + raise ImportError("from mindspeed.ops.gmm import npu_gmm_v2 failed.") from exc + + if mat_a_flat.shape[0] == 0: + return mat_a_flat.new_zeros((0, mat_b.shape[1] if trans_b else mat_b.shape[2])) + + weight = mat_b.transpose(-1, -2) if trans_b else mat_b + if batch_sizes.device.type != "npu": + batch_sizes = batch_sizes.to(device=mat_a_flat.device) + batch_sizes = batch_sizes.to(dtype=torch.int64) + return npu_gmm_v2(mat_a_flat, weight, bias=None, group_list=batch_sizes, group_type=0) + + +def _time_forward(fn, warmup: int, iters: int) -> tuple[float, torch.Tensor]: + out = None + for _ in range(warmup): + with torch.no_grad(): + out = fn() + _sync() + + start = time.perf_counter() + for _ in range(iters): + with torch.no_grad(): + out = fn() + _sync() + return (time.perf_counter() - start) * 1000.0 / iters, out + + +def _time_forward_backward( + fn, a: torch.Tensor, b: torch.Tensor, warmup: int, iters: int +) -> tuple[float, torch.Tensor, torch.Tensor, torch.Tensor]: + out = None + for _ in range(warmup): + out = fn() + out.sum().backward() + a.grad = None + b.grad = None + _sync() + + start = time.perf_counter() + for _ in range(iters): + out = fn() + out.sum().backward() + grad_a = a.grad.detach().clone() + grad_b = b.grad.detach().clone() + a.grad = None + b.grad = None + _sync() + total_ms = (time.perf_counter() - start) * 1000.0 / iters + return total_ms, out.detach().clone(), grad_a, grad_b + + +def _max_abs_diff(x: torch.Tensor, y: torch.Tensor) -> float: + return float((x.float() - y.float()).abs().max().item()) + + +def _check_close(x: torch.Tensor, y: torch.Tensor, rtol: float, atol: float) -> bool: + try: + torch.testing.assert_close(x, y, rtol=rtol, atol=atol) + return True + except Exception: + return False + + +def _bench_one(params: dict[str, object], rtol: float, atol: float): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + raise RuntimeError("NPU is not available.") + + device = torch.device("npu") + warmup = int(params["warmup"]) + iters = int(params["iters"]) + trans_b = bool(params["trans_b"]) + + a_base, b_base, batch_sizes = _build_inputs(params, device) + a_npu = a_base.detach().clone().requires_grad_(True) + b_npu = b_base.detach().clone().requires_grad_(True) + + fw_ms_base, ref_out = _time_forward( + lambda: _run_baseline(a_base, b_base, batch_sizes, trans_b), + warmup=warmup, + iters=iters, + ) + total_ms_base, ref_out_bw, ref_da, ref_db = _time_forward_backward( + lambda: _run_baseline(a_base, b_base, batch_sizes, trans_b), + a=a_base, + b=b_base, + warmup=warmup, + iters=iters, + ) + + fw_ms_npu, out_npu = _time_forward( + lambda: _run_npu_gmm_v2(a_npu, b_npu, batch_sizes, trans_b), + warmup=warmup, + iters=iters, + ) + total_ms_npu, out_npu_bw, da_npu, db_npu = _time_forward_backward( + lambda: _run_npu_gmm_v2(a_npu, b_npu, batch_sizes, trans_b), + a=a_npu, + b=b_npu, + warmup=warmup, + iters=iters, + ) + + print( + f"[npu_grouped_gemm] name={params['name']} group={params['group_size']} " + f"batch={params['batch_size']} k={params['k']} n={params['n']} " + f"dtype={params['dtype']} trans_b={trans_b}" + ) + print("backend, fw_ms, bw_ms, total_ms") + print(f"baseline, {fw_ms_base:.3f}, {total_ms_base - fw_ms_base:.3f}, {total_ms_base:.3f}") + print(f"npu_gmm_v2, {fw_ms_npu:.3f}, {total_ms_npu - fw_ms_npu:.3f}, {total_ms_npu:.3f}") + print( + "speedup_vs_baseline, " + f"fw={fw_ms_base / fw_ms_npu:.2f}x, " + f"bw={(total_ms_base - fw_ms_base) / (total_ms_npu - fw_ms_npu):.2f}x, " + f"total={total_ms_base / total_ms_npu:.2f}x" + ) + print("metric, close, max_abs_diff") + print(f"forward, {_check_close(out_npu, ref_out, rtol, atol)}, {_max_abs_diff(out_npu, ref_out):.6f}") + print( + f"forward_bw_run, {_check_close(out_npu_bw, ref_out_bw, rtol, atol)}, {_max_abs_diff(out_npu_bw, ref_out_bw):.6f}" + ) + print(f"grad_a, {_check_close(da_npu, ref_da, rtol, atol)}, {_max_abs_diff(da_npu, ref_da):.6f}") + print(f"grad_b, {_check_close(db_npu, ref_db, rtol, atol)}, {_max_abs_diff(db_npu, ref_db):.6f}") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--rtol", type=float, default=1e-2) + parser.add_argument("--atol", type=float, default=1e-2) + args = parser.parse_args() + + torch.npu.set_device(0) + for params in DEFAULT_PARAM_SETS: + _bench_one(params, rtol=args.rtol, atol=args.atol) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/ASCEND.md b/docs/ASCEND.md new file mode 100644 index 00000000..e8ebb41f --- /dev/null +++ b/docs/ASCEND.md @@ -0,0 +1,296 @@ +# Ascend Support + +This note documents Ascend/NPU support in StepTronOSS. It focuses on what is +available, how it is enabled, and what runtime constraints matter in practice. + +## Overview + +StepTronOSS provides three main pieces of Ascend support: + +- manual runtime patching through `steptronoss.utils.npu_patch.apply_npu_patch()` +- a native `AttentionCore="npu-flash-attn"` alternative +- native `grouped_gemm="npu_gmm"` and `TokenDispatcher="npu_alltoall"` + alternatives + +The Ascend workflow also includes: + +- benchmark scripts for grouped GEMM and token dispatch +- NPU-specific SFT experiment entrypoints +- a small real-data preparation script for Ascend smoke tests +- unit tests for dispatcher layout logic and grouped GEMM correctness + +## Tested Environment + +The user-facing examples in this document were validated with the following +Ascend software stack: + +- CANN `8.3.RC2` +- `torch_npu` `2.8` +- MindSpeed `v2.2.0_core_r0.12.1` + +All other dependencies are expected to stay aligned with the project's `uv` +environment. + +## Runtime Activation + +Ascend runtime patching is enabled manually. Importing `steptronoss` alone does +not activate NPU patches. + +To enable the Ascend runtime patch layer, call `apply_npu_patch()` before +importing modules that depend on patched NPU behavior: + +```python +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() +``` + +Behavior: + +- if `torch_npu` is unavailable, the patch stays inactive and the normal CUDA + path remains unchanged +- if `torch_npu` is available and `apply_npu_patch()` is called, StepTron: + - marks the runtime as NPU-backed + - imports `torch_npu.contrib.transfer_to_npu` + - optionally replaces `torch.compile` with an identity wrapper + - patches several CUDA-specific helper paths to use NPU-safe implementations + +`apply_npu_patch()` is no longer responsible for registering +`npu-flash-attn`, `npu_gmm`, or `npu_alltoall`. Those alternatives are +declared directly in the corresponding `@optimizable(...)` definitions and +become available when those modules are imported. + +## What Gets Patched + +When the NPU patch is active, the runtime patches the following: + +- CUDA RNG state setters + Redirected to `torch.npu` generator state handling. +- grad clipping and zero counting helpers + Adjusted to work on NPU tensors instead of CUDA-only tensor types. +- fp32 <-> fp16/bf16 conversion helpers + Relaxed to match NPU tensor dtypes. + +## Native NPU Alternatives + +The NPU optimization backends are registered natively through +`@optimizable(...)`, not through `apply_npu_patch()`: + +- `AttentionCore` + Registers `npu-flash-attn`, implemented by `NpuFlashAttention` in + `steptronoss/model/common/attention_core.py` and backed by Hugging Face NPU + flash attention integration. +- `grouped_gemm` + Registers `npu_gmm`, backed by `mindspeed.ops.gmm.npu_gmm_v2`. +- `TokenDispatcher` + Registers `npu_alltoall`, backed by + `steptronoss/model/ep_dispatcher/npu_alltoall_dispatcher.py`. + +In practice, most Ascend training and benchmark entrypoints still call +`apply_npu_patch()` early because they rely on the runtime helper patches in +addition to selecting these alternatives. + +## `npu_alltoall` Dispatcher + +The new dispatcher is a tensorized EP routing backend for NPU. + +Key behavior: + +- builds a unique token-rank routing map per token, so repeated experts on the + same remote rank are deduplicated before communication +- uses `all_to_all_single` for hidden states, token ids, and routing weights +- preserves backward through communication with a custom autograd-enabled + `all_to_all_single` +- on the fast path, tries to use MindSpeed fused + `npu_moe_token_permute` / `npu_moe_token_unpermute` +- on unsupported cases, falls back to `index_select` / `index_add`-based logic + +Fast-path requirements: + +- device type must be `npu` +- hidden states must be `torch.bfloat16` +- MindSpeed token permute/unpermute ops must import successfully + +The dispatcher exposes two observability flags on the instance: + +- `last_dispatch_used_fused_permute` +- `last_combine_used_fused_unpermute` + +These are useful in benchmarks to confirm whether the fused path was actually +hit. + +## `npu_gmm` Grouped GEMM + +The grouped GEMM backend wraps `mindspeed.ops.gmm.npu_gmm_v2`. + +Notes: + +- the implementation accepts the same semantic inputs as the default + `grouped_gemm` +- `batch_sizes` is normalized internally to `torch.int64` +- if `batch_sizes` is not already on NPU, it is moved to the input device +- empty input is handled explicitly and returns an empty output with the correct + shape + +Recommended runtime assumptions: + +- run in bf16 +- keep `batch_sizes` CPU-visible or easy to cast to `torch.int64` + +## Environment Variables + +The Ascend runtime uses or recognizes these environment variables: + +- `STEPTRON_ENABLE_TORCH_COMPILE_ON_NPU=1` + Keep real `torch.compile` on NPU. Without this, the patch layer replaces + `torch.compile` with a no-op wrapper to avoid unsupported compile paths. +- `STEPTRON_SFT_DATALOADER_WORKERS` + Controls async dataloader worker count in the new Qwen3 NPU SFT configs. +- `PYTORCH_NPU_ALLOC_CONF=expandable_segments:True` + Used in the new multi-node Muon SFT resource config. +- `CUDA_DEVICE_MAX_CONNECTIONS=1` + Also set by that resource config for launch consistency with existing + training assumptions. + +## How To Enable NPU Optimizations + +Select the NPU alternatives through the optimization registry. For typical +Ascend runs, you will also call `apply_npu_patch()` early for the runtime patch +layer: + +```python +from steptronoss.utils.npu_patch import apply_npu_patch +from steptronoss.utils.optimizable import set_optimization + +apply_npu_patch() + +set_optimization( + TokenDispatcher="npu_alltoall", + grouped_gemm="npu_gmm", + AttentionCore="npu-flash-attn", +) +``` + +This pattern is used by the NPU experiment entrypoints below. + +## Experiment Entry Points + +The following experiment scripts are available: + +- `playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py` + Single-node toy SFT debug path for the Step3.5 Flash runtime. +- `playground/sft/step3/npu/step3p5_flash_sft_step3_data_muon_npu.py` + Larger Muon-based Step3.5 Flash SFT config for Ascend. +- `playground/sft/qwen3/npu/qwen3_1p7b_sft_step3_data_npu.py` + Qwen3-1.7B SFT config pointing to a compiled real-data dataset. +- `playground/sft/qwen3/npu/qwen3_30a3b_sft_step3_data_npu.py` + Qwen3-30A3B SFT config pointing to the same compiled dataset. + +Example single-node toy launch: + +```bash +torchrun --standalone --nproc-per-node=8 \ + playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py +``` + +These NPU experiment entrypoints call `apply_npu_patch()` near the top of the +file before importing modules that depend on patched NPU runtime behavior. + +The larger Step3.5 Muon config uses a `TorchrunResourceConfig` with +`replica=4` and `gpu=16`, so it should be launched with the repo's usual +multi-node torchrun workflow rather than treated as a one-machine smoke test. + +## Real-Data Preparation Script + +`playground/sft/qwen3/npu/prepare_ascend_qwen3_real_sft_dataset.py` prepares a +small real SFT dataset for Ascend validation. + +It currently: + +- streams `HuggingFaceH4/ultrachat_200k` from ModelScope +- normalizes messages into StepTron conversation format +- keeps loss only on assistant turns +- filters invalid, empty, too-short, and too-long samples +- tokenizes with a local Qwen3 tokenizer path +- writes `dialogs.json`, `stats.json`, and compiled dataset shards + +Example: + +```bash +python playground/sft/qwen3/npu/prepare_ascend_qwen3_real_sft_dataset.py \ + --model-path /oss/model_zoo/Qwen3-1.7B-Base/ \ + --output-root /oss/steptronoss_data/qwen3_1p7b_sft_real_ultrachat +``` + +This script does not call `apply_npu_patch()`, so dataset preparation does not +depend on NPU runtime activation. + +## Tests + +The following tests are available: + +- `tests/test_npu_alltoall_dispatcher.py` + Verifies rank bucketing, deduplication, sort contract, and token id dtype + selection. +- `tests/test_grouped_gemm_npu.py` + Verifies forward/backward equivalence of `npu_gmm` against a reference grouped + matmul implementation. + +Suggested commands: + +```bash +pytest tests/test_npu_alltoall_dispatcher.py +pytest tests/test_grouped_gemm_npu.py +``` + +The grouped GEMM test is skipped unless an actual NPU runtime is available. + +## Benchmarks + +The following benchmark scripts are available: + +- `benchmarks/benchmark_grouped_gemm_npu.py` +- `benchmarks/benchmark_dispatcher_npu.py` + +Example commands: + +```bash +python benchmarks/benchmark_grouped_gemm_npu.py +torchrun --nproc_per_node=8 benchmarks/benchmark_dispatcher_npu.py +``` + +The dispatcher benchmark: + +- initializes HCCL distributed state +- builds an EP mesh through `PM` +- compares dispatcher alternatives through `set_optimization` +- measures forward, backward, correctness, and memory + +Both NPU benchmark scripts call `apply_npu_patch()` explicitly before using NPU +backends because the benchmarks exercise the runtime patch layer as well as the +registered NPU alternatives. + +## Constraints And Fallbacks + +Keep these constraints in mind: + +- `npu_gmm` requires MindSpeed `npu_gmm_v2` +- the fused dispatcher path requires MindSpeed token permute/unpermute ops +- HCCL is expected for distributed NPU runs +- the fused dispatcher path is currently bf16-only +- if optional MindSpeed ops are unavailable, the dispatcher still works through + fallback gather/scatter logic, but performance expectations change + +## Practical Checklist + +Before reporting an Ascend regression, confirm: + +- `torch_npu` imports and `torch.npu.is_available()` is true +- if the run depends on patched NPU runtime helpers, `apply_npu_patch()` has + been called early enough +- `set_optimization(...)` selects `TokenDispatcher="npu_alltoall"` and + `grouped_gemm="npu_gmm"`, and selects + `AttentionCore="npu-flash-attn"` when using the NPU flash attention path +- distributed runs use HCCL +- hidden states use bf16 if you expect the fused dispatcher fast path +- MindSpeed optional ops import successfully if you expect peak performance diff --git a/docs/ASCEND_ZH.md b/docs/ASCEND_ZH.md new file mode 100644 index 00000000..3ca6ef0f --- /dev/null +++ b/docs/ASCEND_ZH.md @@ -0,0 +1,281 @@ +# Ascend 支持说明 + +本文档说明 StepTronOSS 中的 Ascend/NPU 支持，重点覆盖可用能力、启用方式、运行约束，以及对应的实验、测试和 benchmark 入口。 + +## 概览 + +StepTronOSS 目前提供三块主要的 Ascend 支持： + +- 通过 `steptronoss.utils.npu_patch.apply_npu_patch()` 手动启用运行时补丁 +- 原生注册 `AttentionCore="npu-flash-attn"` alternative +- 原生注册 `grouped_gemm="npu_gmm"` 和 + `TokenDispatcher="npu_alltoall"` alternatives + +相关工作流还包括： + +- grouped GEMM 和 token dispatcher 的 benchmark +- NPU 专用 SFT 实验入口 +- 用于 Ascend smoke test 的小规模真实数据准备脚本 +- dispatcher 路由逻辑和 grouped GEMM 正确性测试 + +## 测试环境 + +本文档中的 Ascend 使用说明基于以下软件栈验证： + +- CANN `8.3.RC2` +- `torch_npu` `2.8` +- MindSpeed `v2.2.0_core_r0.12.1` + +其余依赖与项目的 `uv` 环境保持一致。 + +## 运行时启用方式 + +Ascend 运行时 patch 需要手动启用。仅仅导入 `steptronoss` 不会自动激活 +NPU patch。 + +如果要启用 Ascend 运行时 patch 层，需要在导入依赖补丁后 NPU 行为的模块 +之前，显式调用 +`apply_npu_patch()`： + +```python +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() +``` + +行为如下： + +- 如果 `torch_npu` 不可用，补丁不会生效，默认 CUDA 路径保持不变 +- 如果 `torch_npu` 可用且调用了 `apply_npu_patch()`，StepTron 会： + - 标记当前运行时为 NPU + - 导入 `torch_npu.contrib.transfer_to_npu` + - 按环境变量决定是否屏蔽 `torch.compile` + - 替换若干 CUDA-only helper 为 NPU 可用实现 + +`apply_npu_patch()` 现在不再负责注册 `npu-flash-attn`、`npu_gmm`、 +`npu_alltoall`。这些 alternative 直接写在对应模块的 +`@optimizable(...)` 定义里，模块导入后即可在 registry 中看到。 + +## 当前补丁覆盖的内容 + +当 NPU patch 激活后，运行时会替换这些路径： + +- CUDA RNG state setter + 改为使用 `torch.npu` 的随机数状态逻辑。 +- grad clipping 和 zero counting helper + 去掉对 CUDA tensor type 的假设，适配 NPU tensor。 +- fp32 <-> fp16/bf16 conversion helper + 放宽 dtype 判断，适配 NPU 张量。 + +## 原生注册的 NPU Alternatives + +NPU 优化后端现在通过 `@optimizable(...)` 原生注册，而不是由 +`apply_npu_patch()` 动态塞进 registry： + +- `AttentionCore` + 注册 `npu-flash-attn`，实现类是 + `steptronoss/model/common/attention_core.py` 里的 + `NpuFlashAttention`，底层走 Hugging Face 的 NPU flash attention 集成。 +- `grouped_gemm` + 注册 `npu_gmm`，底层是 `mindspeed.ops.gmm.npu_gmm_v2`。 +- `TokenDispatcher` + 注册 `npu_alltoall`，实现位于 + `steptronoss/model/ep_dispatcher/npu_alltoall_dispatcher.py`。 + +在实际使用里，大多数 Ascend 训练和 benchmark 入口仍然会很早调用 +`apply_npu_patch()`，因为除了选择这些 alternative 之外，它们还依赖运行时 +helper patch。 + +## `npu_alltoall` Dispatcher + +新的 dispatcher 是一个面向 NPU 的 EP token routing backend。 + +核心行为： + +- 为每个 token 构造唯一的 token-rank 路由映射，同一 token 若 top-k expert 落在同一远端 rank，会先去重再通信 +- 用 `all_to_all_single` 传 hidden states、token ids 和 routing weights +- 通过自定义带 autograd 的 `all_to_all_single` 保留反向传播 +- 在 fast path 下，尝试使用 MindSpeed 的 + `npu_moe_token_permute` / `npu_moe_token_unpermute` +- 如果不满足 fast path 条件，就回退到 `index_select` / `index_add` 逻辑 + +fused fast path 的前提： + +- device 类型必须是 `npu` +- hidden states 必须是 `torch.bfloat16` +- MindSpeed 的 token permute/unpermute 能正常导入 + +dispatcher 还暴露了两个观测字段： + +- `last_dispatch_used_fused_permute` +- `last_combine_used_fused_unpermute` + +benchmark 时可以用它们确认本次是否真的打到了 fused path。 + +## `npu_gmm` Grouped GEMM + +新的 grouped GEMM backend 封装了 `mindspeed.ops.gmm.npu_gmm_v2`。 + +实现上的注意点： + +- 接口保持和默认 `grouped_gemm` 一致 +- `batch_sizes` 会在内部归一化为 `torch.int64` +- 如果 `batch_sizes` 不在 NPU 上，会先搬到输入所在设备 +- 空输入会显式返回 shape 正确的空输出 + +推荐的运行假设： + +- 以 bf16 为主 +- `batch_sizes` 至少要能方便转成 `torch.int64` + +## 环境变量 + +Ascend 运行时会使用或识别这些环境变量： + +- `STEPTRON_ENABLE_TORCH_COMPILE_ON_NPU=1` + 保留真实的 `torch.compile`。不设置时，patch 层会把 `torch.compile` + 替换成 no-op wrapper，避开暂不稳定的 compile 路径。 +- `STEPTRON_SFT_DATALOADER_WORKERS` + 控制新 Qwen3 NPU SFT 配置中的异步 dataloader worker 数量。 +- `PYTORCH_NPU_ALLOC_CONF=expandable_segments:True` + 新的 Muon SFT 资源配置里会设置它。 +- `CUDA_DEVICE_MAX_CONNECTIONS=1` + 同样由该资源配置设置，用来和现有训练假设保持一致。 + +## 如何开启 NPU 优化 + +通过优化 registry 选择 NPU alternative。典型 Ascend 运行里，通常还会在很早 +的位置调用 `apply_npu_patch()` 启用运行时 patch 层： + +```python +from steptronoss.utils.npu_patch import apply_npu_patch +from steptronoss.utils.optimizable import set_optimization + +apply_npu_patch() + +set_optimization( + TokenDispatcher="npu_alltoall", + grouped_gemm="npu_gmm", + AttentionCore="npu-flash-attn", +) +``` + +这也是下面这些 NPU 实验脚本采用的配置方式。 + +## 实验入口 + +可用的实验脚本包括： + +- `playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py` + 单机 toy SFT 调试入口，尽量贴近 Step3.5 Flash runtime。 +- `playground/sft/step3/npu/step3p5_flash_sft_step3_data_muon_npu.py` + 更完整的 Ascend Step3.5 Flash + Muon SFT 配置。 +- `playground/sft/qwen3/npu/qwen3_1p7b_sft_step3_data_npu.py` + 指向编译后真实数据集的 Qwen3-1.7B SFT 配置。 +- `playground/sft/qwen3/npu/qwen3_30a3b_sft_step3_data_npu.py` + 指向同一真实数据集的 Qwen3-30A3B SFT 配置。 + +单机 toy 调试示例： + +```bash +torchrun --standalone --nproc-per-node=8 \ + playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py +``` + +这些 NPU 实验入口都会在文件开头先调用 `apply_npu_patch()`，再导入依赖 +补丁后 NPU 运行时行为的模块。 + +更大的 Step3.5 Muon 配置内部使用了 `TorchrunResourceConfig`， +其中 `replica=4`、`gpu=16`，因此更适合走仓库现有的多机 torchrun +流程，不应把它当成单机 smoke test。 + +## 真实数据准备脚本 + +`playground/sft/qwen3/npu/prepare_ascend_qwen3_real_sft_dataset.py` +用于准备 Ascend 验证用的小规模真实 SFT 数据集。 + +当前脚本会： + +- 从 ModelScope 流式读取 `HuggingFaceH4/ultrachat_200k` +- 归一化为 StepTron 对话格式 +- 只对 assistant turn 打开 loss +- 过滤非法、空内容、过短、过长样本 +- 使用本地 Qwen3 tokenizer 路径做 tokenize +- 输出 `dialogs.json`、`stats.json` 和编译后的 dataset shard + +示例： + +```bash +python playground/sft/qwen3/npu/prepare_ascend_qwen3_real_sft_dataset.py \ + --model-path /oss/model_zoo/Qwen3-1.7B-Base/ \ + --output-root /oss/steptronoss_data/qwen3_1p7b_sft_real_ultrachat +``` + +这个脚本不会调用 `apply_npu_patch()`，因此数据准备流程不依赖 NPU +运行时激活。 + +## 测试 + +可用的测试包括： + +- `tests/test_npu_alltoall_dispatcher.py` + 覆盖 rank bucket、去重、排序契约和 token id dtype 选择逻辑。 +- `tests/test_grouped_gemm_npu.py` + 对 `npu_gmm` 做 forward/backward 与 reference grouped matmul 的一致性校验。 + +推荐命令： + +```bash +pytest tests/test_npu_alltoall_dispatcher.py +pytest tests/test_grouped_gemm_npu.py +``` + +其中 grouped GEMM 测试在没有真实 NPU 运行时的环境下会自动 skip。 + +## Benchmark + +可用的 benchmark 脚本包括： + +- `benchmarks/benchmark_grouped_gemm_npu.py` +- `benchmarks/benchmark_dispatcher_npu.py` + +示例命令： + +```bash +python benchmarks/benchmark_grouped_gemm_npu.py +torchrun --nproc_per_node=8 benchmarks/benchmark_dispatcher_npu.py +``` + +dispatcher benchmark 会： + +- 初始化 HCCL 分布式环境 +- 通过 `PM` 建立 EP mesh +- 通过 `set_optimization` 比较不同 dispatcher alternative +- 统计 forward、backward、correctness 和 memory + +这两个 NPU benchmark 脚本都会先显式调用 `apply_npu_patch()`，再使用 NPU +backend，因为 benchmark 同时覆盖运行时 patch 层和已注册的 NPU +alternatives。 + +## 约束与降级路径 + +需要记住这些运行约束： + +- `npu_gmm` 依赖 MindSpeed 的 `npu_gmm_v2` +- dispatcher fused path 依赖 MindSpeed token permute/unpermute +- 分布式 NPU 运行默认假设 backend 是 HCCL +- fused dispatcher 目前只覆盖 bf16 path +- 如果可选的 MindSpeed op 不可用，dispatcher 仍然能走 fallback gather/scatter 逻辑，但性能预期会下降 + +## 实操检查清单 + +如果要排查 Ascend 回归，先确认： + +- `torch_npu` 能导入，且 `torch.npu.is_available()` 为真 +- 如果本次运行依赖 NPU 运行时 helper patch，`apply_npu_patch()` 调用得足够早 +- `set_optimization(...)` 已选择 `TokenDispatcher="npu_alltoall"` 和 + `grouped_gemm="npu_gmm"`；如果走 NPU flash attention，还应选择 + `AttentionCore="npu-flash-attn"` +- 分布式运行使用 HCCL +- 如果期望命中 fused dispatcher fast path，hidden states 是 bf16 +- 如果期望拿到最高性能，MindSpeed 对应可选 op 可以正常导入 diff --git a/docs/index.md b/docs/index.md index d9ae14b0..94efc871 100644 --- a/docs/index.md +++ b/docs/index.md @@ -47,6 +47,8 @@ cfshow playground/rlvr/qwen3_1p5b_rlvr_math.py - Start here: [LAUNCH_EXPERIMENTS](LAUNCH_EXPERIMENTS.md) - Launch guide (EN): [LAUNCH_EXPERIMENTS](LAUNCH_EXPERIMENTS.md) - Launch guide (ZH): [LAUNCH_EXPERIMENTS_ZH](LAUNCH_EXPERIMENTS_ZH.md) +- Ascend guide (EN): [ASCEND](ASCEND.md) +- Ascend guide (ZH): [ASCEND_ZH](ASCEND_ZH.md) - SFT 数据准备: [SFT_DATA_PREPARATION](SFT_DATA_PREPARATION.md) - SFT Data Preparation (EN): [SFT_DATA_PREPARATION_EN](SFT_DATA_PREPARATION_EN.md) - Memory tuning (ZH): [MEMORY_TUNE](MEMORY_TUNE.md) diff --git a/playground/sft/qwen3/npu/prepare_ascend_qwen3_real_sft_dataset.py b/playground/sft/qwen3/npu/prepare_ascend_qwen3_real_sft_dataset.py new file mode 100644 index 00000000..bda08fdf --- /dev/null +++ b/playground/sft/qwen3/npu/prepare_ascend_qwen3_real_sft_dataset.py @@ -0,0 +1,172 @@ +import argparse +import json +import os +import sys +from dataclasses import dataclass +from pathlib import Path + +from modelscope import MsDataset +from torch.utils.data import Dataset + +from steptronoss.data.chat_templates.text_template import HuggingFaceTemplate +from steptronoss.data.datasets.compile_dataset import compile_dataset +from steptronoss.tokenizer.hf_compat_tokenizer import HFCompatTokenizer + +DEFAULT_DATASET_NAME = "HuggingFaceH4/ultrachat_200k" +DEFAULT_SPLIT = "train_sft" +DEFAULT_MODEL_PATH = "/oss/model_zoo/Qwen3-1.7B-Base/" +DEFAULT_OUTPUT_ROOT = "/oss/steptronoss_data/qwen3_1p7b_sft_real_ultrachat" + + +class LocalQwen3Tokenizer(HFCompatTokenizer): + hf_path = DEFAULT_MODEL_PATH + + +@dataclass +class PreparedStats: + raw_seen: int = 0 + kept_dialogs: int = 0 + filtered_missing_messages: int = 0 + filtered_invalid_role: int = 0 + filtered_no_final_assistant: int = 0 + filtered_empty_content: int = 0 + filtered_too_short: int = 0 + filtered_too_long: int = 0 + + +class TokenizedListDataset(Dataset): + def __init__(self, items): + self.items = items + + def __len__(self): + return len(self.items) + + def __getitem__(self, index): + return self.items[index] + + +def _normalize_messages(sample): + messages = sample.get("messages") + if not isinstance(messages, list) or len(messages) < 2: + return None, "missing_messages" + + dialogs = [] + valid_roles = {"system", "user", "assistant"} + for item in messages: + role = str(item.get("role", "")).lower().strip() + if role not in valid_roles: + return None, "invalid_role" + content = item.get("content", "") + if isinstance(content, list): + content = "\n".join(str(x) for x in content if x is not None) + else: + content = str(content) + content = content.strip() + if not content: + return None, "empty_content" + dialogs.append({ + "role": role, + "content": content, + "name": "", + "loss_mask": 1 if role == "assistant" else 0, + "ground_truth": None, + }) + + if dialogs[-1]["role"] != "assistant": + return None, "no_final_assistant" + + return {"conversations": dialogs, "images": None}, None + + +def build_real_dataset( + dataset_name: str, + split: str, + model_path: str, + output_root: str, + max_samples: int, + max_tokens_per_sample: int, + min_tokens_per_sample: int, +): + class _Tokenizer(LocalQwen3Tokenizer): + hf_path = model_path + + tokenizer = _Tokenizer() + template = HuggingFaceTemplate(tokenizer) + + stats = PreparedStats() + dialogs = [] + tokenized_items = [] + + stream = MsDataset.load(dataset_name, split=split, use_streaming=True) + for sample in stream: + stats.raw_seen += 1 + dialog, err = _normalize_messages(sample) + if err is not None: + setattr(stats, f"filtered_{err}", getattr(stats, f"filtered_{err}") + 1) + continue + + tokenized = template(dialog) + token_count = len(tokenized["tokens"]) + if token_count < min_tokens_per_sample: + stats.filtered_too_short += 1 + continue + if token_count > max_tokens_per_sample: + stats.filtered_too_long += 1 + continue + + dialogs.append(dialog) + tokenized_items.append(tokenized) + stats.kept_dialogs += 1 + if stats.kept_dialogs >= max_samples: + break + + output_root = Path(output_root) + dialogs_path = output_root / "dialogs.json" + stats_path = output_root / "stats.json" + compiled_path = output_root / "compiled" + + os.makedirs(output_root, exist_ok=True) + with open(dialogs_path, "w", encoding="utf-8") as f: + json.dump(dialogs, f, ensure_ascii=False) + with open(stats_path, "w", encoding="utf-8") as f: + json.dump(stats.__dict__, f, ensure_ascii=False, indent=2) + + tokenized_dataset = TokenizedListDataset(tokenized_items) + compile_dataset( + tokenized_dataset, + str(compiled_path), + sample_meta_extractor=lambda x: len(x["tokens"]), + num_workers=min(8, os.cpu_count() or 1), + ) + + print(f"Prepared real dataset: {compiled_path}") + print(json.dumps(stats.__dict__, ensure_ascii=False, indent=2)) + + +def main(): + parser = argparse.ArgumentParser(description="Prepare a small real Qwen3-like SFT dataset for Ascend validation.") + parser.add_argument("--dataset-name", default=DEFAULT_DATASET_NAME) + parser.add_argument("--split", default=DEFAULT_SPLIT) + parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH) + parser.add_argument("--output-root", default=DEFAULT_OUTPUT_ROOT) + parser.add_argument("--max-samples", type=int, default=2048) + parser.add_argument("--max-tokens-per-sample", type=int, default=1024) + parser.add_argument("--min-tokens-per-sample", type=int, default=32) + args = parser.parse_args() + + build_real_dataset( + dataset_name=args.dataset_name, + split=args.split, + model_path=args.model_path, + output_root=args.output_root, + max_samples=args.max_samples, + max_tokens_per_sample=args.max_tokens_per_sample, + min_tokens_per_sample=args.min_tokens_per_sample, + ) + sys.stdout.flush() + sys.stderr.flush() + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/playground/sft/qwen3/npu/qwen3_1p7b_sft_step3_data_npu.py b/playground/sft/qwen3/npu/qwen3_1p7b_sft_step3_data_npu.py new file mode 100644 index 00000000..95606977 --- /dev/null +++ b/playground/sft/qwen3/npu/qwen3_1p7b_sft_step3_data_npu.py @@ -0,0 +1,59 @@ +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() # Ensure the NPU patch is applied before importing any other modules that might use NPU features + +import os + +from playground.data.sft.oss260312.step_sft_data_config0311_qwen_tokenizer import ( + Recipe0311QwenCompiledSFTDataConfig, +) +from playground.sft.qwen3.qwen3_1p7b_sft_step3_data import Exp as BaseExp +from playground.tools.compile_recipe import ( + CompiledDataRecipe, + CompiledDatasetsConfig, +) + + +class AscendRealDatasetsConfig(CompiledDatasetsConfig): + compiled_recipe = CompiledDataRecipe( + domains={ + "real_ultrachat": "/oss/steptronoss_data/qwen3_1p7b_sft_real_ultrachat/compiled", + }, + epochs={ + "real_ultrachat": 1.0, + }, + ) + + +class AscendRealStep3SFTDataQwenTokenizedConfig(Recipe0311QwenCompiledSFTDataConfig): + dataset_cfg = AscendRealDatasetsConfig + + def build_dataloader(self, dp_rank=0, dp_size=1): + from steptronoss.data.dataloader.packed_dataloader import MixedPackedDataloader + from steptronoss.data.nextable import DPMux, async_accelearte_slowfast + + datasets = self.dataset_cfg.build_datasets() + dataloader = MixedPackedDataloader( + datasets=[ds[0] for ds in datasets.values()], + epochs=[ds[1] for ds in datasets.values()], + max_length=self.max_packing_seqlen, + oversize_policy=self.oversize_policy, + transform=self.pack, + ) + dataloader = DPMux(dataloader, dp_size=dp_size, dp_rank=dp_rank) + + num_workers = int(os.getenv("STEPTRON_SFT_DATALOADER_WORKERS", "1")) + return async_accelearte_slowfast(dataloader, num_workers=max(1, num_workers)) + + +class Exp(BaseExp): + data_cfg = AscendRealStep3SFTDataQwenTokenizedConfig + + def __init__(self): + super().__init__() + self.checkpoint_cfg.load_safetensors = "/oss/model_zoo/Qwen3-1.7B-Base/" + self.checkpoint_cfg.save_dir = "/oss/checkpoints/" + + +if __name__ == "__main__": + Exp().train() diff --git a/playground/sft/qwen3/npu/qwen3_30a3b_sft_step3_data_npu.py b/playground/sft/qwen3/npu/qwen3_30a3b_sft_step3_data_npu.py new file mode 100644 index 00000000..45a34e26 --- /dev/null +++ b/playground/sft/qwen3/npu/qwen3_30a3b_sft_step3_data_npu.py @@ -0,0 +1,59 @@ +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() # Ensure the NPU patch is applied before importing any other modules that might use NPU features + +import os + +from playground.data.sft.oss260312.step_sft_data_config0311_qwen_tokenizer import ( + Recipe0311QwenCompiledSFTDataConfig, +) +from playground.sft.qwen3.qwen3_30a3b_sft_step3_data import Exp as BaseExp +from playground.tools.compile_recipe import ( + CompiledDataRecipe, + CompiledDatasetsConfig, +) + + +class AscendRealDatasetsConfig(CompiledDatasetsConfig): + compiled_recipe = CompiledDataRecipe( + domains={ + "real_ultrachat": "/oss/steptronoss_data/qwen3_1p7b_sft_real_ultrachat/compiled", + }, + epochs={ + "real_ultrachat": 1.0, + }, + ) + + +class AscendRealStep3SFTDataQwenTokenizedConfig(Recipe0311QwenCompiledSFTDataConfig): + dataset_cfg = AscendRealDatasetsConfig + + def build_dataloader(self, dp_rank=0, dp_size=1): + from steptronoss.data.dataloader.packed_dataloader import MixedPackedDataloader + from steptronoss.data.nextable import DPMux, async_accelearte_slowfast + + datasets = self.dataset_cfg.build_datasets() + dataloader = MixedPackedDataloader( + datasets=[ds[0] for ds in datasets.values()], + epochs=[ds[1] for ds in datasets.values()], + max_length=self.max_packing_seqlen, + oversize_policy=self.oversize_policy, + transform=self.pack, + ) + dataloader = DPMux(dataloader, dp_size=dp_size, dp_rank=dp_rank) + + num_workers = int(os.getenv("STEPTRON_SFT_DATALOADER_WORKERS", "1")) + return async_accelearte_slowfast(dataloader, num_workers=max(1, num_workers)) + + +class Exp(BaseExp): + data_cfg = AscendRealStep3SFTDataQwenTokenizedConfig + + def __init__(self): + super().__init__() + self.checkpoint_cfg.load_safetensors = "/oss/model_zoo/Qwen/Qwen3-30B-A3B-Base/model" + self.checkpoint_cfg.save_dir = "/oss/checkpoints/" + + +if __name__ == "__main__": + Exp().train() diff --git a/playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py b/playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py new file mode 100644 index 00000000..adca5fa2 --- /dev/null +++ b/playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py @@ -0,0 +1,47 @@ +""" +Single-node SFT debugging experiment for the Step3.5 Flash runtime path. + +Launch from the repo root with 8 GPUs on one machine: + + torchrun --standalone --nproc-per-node=8 playground/sft/step3/npu/step3_toy_sft_step3_data_npu.py + +This experiment is intended to stay aligned with the single-node Step3.5 +Flash setup in this repo while swapping in the toy Step3.5 model for faster +debugging. Aside from the toy model definition and checkpoint/weight handling, +the SFT data pipeline and the key runtime settings are kept consistent with +the Step3.5 Flash single-node configuration. +""" + +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() # Ensure the NPU patch is applied before importing any other modules that might use NPU features + +from playground.sft.step3.step3_toy_sft_step3_data import Exp as BaseExp + + +class Exp(BaseExp): + """Toy-model SFT debug config that mirrors the Step3.5 Flash 1-node path.""" + + def __init__(self): + super().__init__() + self.trainer_cfg.global_seq_length = 8 * 1024 + + self.model_cfg.num_layers = 4 + self.model_cfg.swa_layer_list = [False, True, False, True] + self.model_cfg.parallel_cfg.context_parallel_size = 1 + self.model_cfg.parallel_cfg.tensor_model_parallel_size = 8 + self.model_cfg.tp_cfg.sequence_parallel = True + self.trainer_cfg.offload_optimizer_state = True + + def configure_optimizable(self): + from steptronoss.utils.optimizable import set_optimization + + set_optimization( + TokenDispatcher="npu_alltoall", + grouped_gemm="npu_gmm", + AttentionCore="npu-flash-attn", + ) + + +if __name__ == "__main__": + Exp().train() diff --git a/playground/sft/step3/npu/step3p5_flash_sft_step3_data_muon_npu.py b/playground/sft/step3/npu/step3p5_flash_sft_step3_data_muon_npu.py new file mode 100644 index 00000000..04a286a4 --- /dev/null +++ b/playground/sft/step3/npu/step3p5_flash_sft_step3_data_muon_npu.py @@ -0,0 +1,46 @@ +from steptronoss.utils.npu_patch import apply_npu_patch + +apply_npu_patch() # Ensure the NPU patch is applied before importing any other modules that might use NPU features + + +from playground.sft.step3.step3p5_flash_sft_step3_data_muon import Exp as BaseExp +from steptronoss.exp.resources import TorchrunResourceConfig + + +class Step3F128kSFTResourceConfig(TorchrunResourceConfig): + def __init__(self): + super().__init__() + self.replica = 4 + self.gpu = 16 + self.envs |= { + "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:True", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + } + + +class Exp(BaseExp): + log_dir = "/data/logs/" + + resource_cfg = Step3F128kSFTResourceConfig + + def __init__(self): + super().__init__() + self.trainer_cfg.global_seq_length = 1024 * 8 + + self.model_cfg.parallel_cfg.context_parallel_size = 1 + self.model_cfg.parallel_cfg.tensor_model_parallel_size = 8 + self.model_cfg.tp_cfg.sequence_parallel = True + self.trainer_cfg.offload_optimizer_state = True + + def configure_optimizable(self): + from steptronoss.utils.optimizable import set_optimization + + set_optimization( + TokenDispatcher="npu_alltoall", + grouped_gemm="npu_gmm", + AttentionCore="npu-flash-attn", + ) + + +if __name__ == "__main__": + Exp().train() diff --git a/steptronoss/model/common/attention_core.py b/steptronoss/model/common/attention_core.py index 542bb8ce..7c579b31 100644 --- a/steptronoss/model/common/attention_core.py +++ b/steptronoss/model/common/attention_core.py @@ -32,6 +32,90 @@ def parse_cu_seqlens(cu_seqlens, max_seq_len=None): return cu_seqlens_q, cu_seqlens_k, max_q_len, max_k_len +class NpuFlashAttention(nn.Module): + """Flash Attention implementation wrapper. + + This wraps flash_attn for efficient attention computation. + """ + + def __init__( + self, + causal: bool = True, + attention_dropout: float = 0.0, + sliding_window: int = -1, + **kwargs, + ): + super().__init__() + self.causal = causal + self.attention_dropout = attention_dropout + self.sliding_window = (sliding_window, sliding_window) # for fa3 + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seq_len: int | None = None, + ) -> torch.Tensor: + """Compute flash attention. + + Args: + q: Query tensor of shape [batch, seq, heads, head_dim] + k: Key tensor of shape [batch, seq, kv_heads, head_dim] + v: Value tensor of shape [batch, seq, kv_heads, head_dim] + cu_seqlens: Cumulative sequence lengths for variable length sequences + max_seq_len: Maximum sequence length + cpu_offload_info: CPU offload configuration + + Returns: + Attention output of shape [batch, seq, heads, head_dim] + """ + from transformers.integrations.npu_flash_attention import ( + npu_flash_attn_func as flash_attn_func, + ) + from transformers.integrations.npu_flash_attention import ( + npu_flash_attn_varlen_func as flash_attn_varlen_func, + ) + + batch_size, seq_len, num_heads, head_dim = q.shape + + if cu_seqlens is not None: + # Variable length attention + cu_seqlens_q, cu_seqlens_k, max_q_len, max_k_len = parse_cu_seqlens(cu_seqlens, max_seq_len) + + # Reshape for varlen attention: [batch, seq, heads, dim] -> [total, heads, dim] + q = q.reshape(-1, num_heads, head_dim) + k = k.reshape(-1, k.shape[2], head_dim) + v = v.reshape(-1, v.shape[2], head_dim) + + output = flash_attn_varlen_func( + q.contiguous(), + k.contiguous(), + v.contiguous(), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_q_len, + max_seqlen_k=max_k_len, + dropout_p=self.attention_dropout if self.training else 0.0, + causal=self.causal, + window_size=self.sliding_window, + ) + output = output.reshape(batch_size, seq_len, num_heads, head_dim) + else: + # Standard flash attention + output = flash_attn_func( + q, + k, + v, + dropout_p=self.attention_dropout if self.training else 0.0, + causal=self.causal, + window_size=self.sliding_window, + ) + + return output + + class FlashAttention(nn.Module): """Flash Attention implementation wrapper. @@ -183,7 +267,9 @@ def forward( return output -@optimizable(alternatives={"flash-attn": FlashAttention, "flash-attn-3": FlashAttention3}) +@optimizable( + alternatives={"npu-flash-attn": NpuFlashAttention, "flash-attn": FlashAttention, "flash-attn-3": FlashAttention3} +) class AttentionCore(nn.Module): """Scaled Dot-Product Attention (SDPA) implementation with FlashAttention-compatible API.""" diff --git a/steptronoss/model/ep_dispatcher/npu_alltoall_dispatcher.py b/steptronoss/model/ep_dispatcher/npu_alltoall_dispatcher.py new file mode 100644 index 00000000..d352541b --- /dev/null +++ b/steptronoss/model/ep_dispatcher/npu_alltoall_dispatcher.py @@ -0,0 +1,306 @@ +from __future__ import annotations + +from functools import lru_cache + +import torch +import torch.distributed as dist + +from steptronoss.core.parallel_state import PM + +INT32_MAX = torch.iinfo(torch.int32).max + + +def _all_gather_splits(send_sizes: torch.Tensor, group, world_size: int) -> torch.Tensor: + send_sizes = send_sizes.to(dtype=torch.int64) + if world_size == 1: + return send_sizes.unsqueeze(0) + + gathered = send_sizes.new_empty(world_size * send_sizes.numel()) + dist.all_gather_into_tensor(gathered, send_sizes.contiguous(), group=group) + return gathered.view(world_size, send_sizes.numel()) + + +class _AllToAllSingle(torch.autograd.Function): + @staticmethod + def forward(ctx, input_tensor: torch.Tensor, output_splits: list[int], input_splits: list[int], group): + ctx.output_splits = [int(x) for x in output_splits] + ctx.input_splits = [int(x) for x in input_splits] + ctx.group = group + + total_output = sum(ctx.output_splits) + output = input_tensor.new_empty((total_output,) + input_tensor.shape[1:]) + dist.all_to_all_single( + output, + input_tensor.contiguous(), + output_split_sizes=ctx.output_splits, + input_split_sizes=ctx.input_splits, + group=group, + ) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + total_input = sum(ctx.input_splits) + grad_input = grad_output.new_empty((total_input,) + grad_output.shape[1:]) + dist.all_to_all_single( + grad_input, + grad_output.contiguous(), + output_split_sizes=ctx.input_splits, + input_split_sizes=ctx.output_splits, + group=ctx.group, + ) + return grad_input, None, None, None + + +def _all_to_all_single_tensor(input_tensor: torch.Tensor, output_splits: list[int], input_splits: list[int], group): + return _AllToAllSingle.apply(input_tensor, output_splits, input_splits, group) + + +def _all_to_all_single_no_grad( + input_tensor: torch.Tensor, + output_splits: list[int], + input_splits: list[int], + group, +) -> torch.Tensor: + total_output = sum(int(x) for x in output_splits) + output = input_tensor.new_empty((total_output,) + input_tensor.shape[1:]) + dist.all_to_all_single( + output, + input_tensor.contiguous(), + output_split_sizes=[int(x) for x in output_splits], + input_split_sizes=[int(x) for x in input_splits], + group=group, + ) + return output + + +def _pick_token_id_dtype(num_tokens: int) -> torch.dtype: + if num_tokens <= INT32_MAX: + return torch.int32 + return torch.int64 + + +def _stable_bucket_order_by_rank(dst_rank: torch.Tensor, world_size: int) -> torch.Tensor: + if dst_rank.numel() == 0: + return dst_rank.new_empty((0,), dtype=torch.int64) + + order_parts = [] + for rank in range(world_size): + rank_order = torch.nonzero(dst_rank == rank, as_tuple=True)[0] + if rank_order.numel() > 0: + order_parts.append(rank_order) + + if not order_parts: + return dst_rank.new_empty((0,), dtype=torch.int64) + if len(order_parts) == 1: + return order_parts[0] + return torch.cat(order_parts, dim=0) + + +def _build_rank_permute_map( + token_expert_ranks: torch.Tensor, + world_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if token_expert_ranks.numel() == 0: + empty = token_expert_ranks.new_empty((0,), dtype=torch.int64) + send_sizes = token_expert_ranks.new_zeros((world_size,), dtype=torch.int64) + rank_permute_map = token_expert_ranks.new_full(token_expert_ranks.shape, world_size, dtype=torch.int32) + valid_mask = token_expert_ranks.new_zeros(token_expert_ranks.shape, dtype=torch.bool) + return rank_permute_map, valid_mask, empty, empty, send_sizes + + sorted_ranks, _ = token_expert_ranks.to(dtype=torch.int64).sort(dim=1) + unique_mask = torch.ones_like(sorted_ranks, dtype=torch.bool) + if sorted_ranks.size(1) > 1: + unique_mask[:, 1:] = sorted_ranks[:, 1:] != sorted_ranks[:, :-1] + valid_mask = unique_mask & (sorted_ranks >= 0) & (sorted_ranks < world_size) + invalid_rank = torch.full_like(sorted_ranks, world_size) + rank_permute_map = torch.where(valid_mask, sorted_ranks, invalid_rank).to(torch.int32) + + token_idx, _ = torch.nonzero(valid_mask, as_tuple=True) + if token_idx.numel() == 0: + empty = sorted_ranks.new_empty((0,), dtype=torch.int64) + send_sizes = sorted_ranks.new_zeros((world_size,), dtype=torch.int64) + return rank_permute_map, valid_mask, empty, empty, send_sizes + + dst_rank = sorted_ranks[valid_mask] + send_sizes = torch.bincount(dst_rank, minlength=world_size) + order = _stable_bucket_order_by_rank(dst_rank, world_size) + return rank_permute_map, valid_mask, token_idx.index_select(0, order), dst_rank.index_select(0, order), send_sizes + + +def _build_unique_token_rank_pairs( + token_expert_ranks: torch.Tensor, + world_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + _, _, token_idx, dst_rank, send_sizes = _build_rank_permute_map(token_expert_ranks, world_size) + return token_idx, dst_rank, send_sizes + + +@lru_cache(maxsize=1) +def _load_npu_moe_token_ops(): + try: + from mindspeed.ops.npu_moe_token_permute import npu_moe_token_permute + from mindspeed.ops.npu_moe_token_unpermute import npu_moe_token_unpermute + except Exception: + return None, None + return npu_moe_token_permute, npu_moe_token_unpermute + + +def _permute_hidden_states( + hidden_states: torch.Tensor, + token_idx: torch.Tensor, + rank_permute_map: torch.Tensor, + num_out_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor | None, bool]: + npu_moe_token_permute, _ = _load_npu_moe_token_ops() + if ( + npu_moe_token_permute is None + or hidden_states.device.type != "npu" + or hidden_states.dtype != torch.bfloat16 + or num_out_tokens <= 0 + ): + return hidden_states.index_select(0, token_idx), None, False + + permuted, sorted_indices = npu_moe_token_permute( + hidden_states, + rank_permute_map, + num_out_tokens=int(num_out_tokens), + padded_mode=False, + ) + return permuted, sorted_indices, True + + +def _build_rank_combine_probs(valid_mask: torch.Tensor) -> torch.Tensor: + return valid_mask.to(dtype=torch.bfloat16) + + +def _unpermute_hidden_states( + hidden_states: torch.Tensor, + sorted_indices: torch.Tensor, + combine_probs: torch.Tensor, + restore_shape: tuple[int, int], +) -> torch.Tensor: + _, npu_moe_token_unpermute = _load_npu_moe_token_ops() + return npu_moe_token_unpermute( + hidden_states, + sorted_indices.to(dtype=torch.int32), + probs=combine_probs, + padded_mode=False, + restore_shape=restore_shape, + ) + + +class NPUAllToAllDispatcher: + """Tensorized all-to-all dispatcher for EP token routing on NPU.""" + + def __init__(self, parallel: str, num_experts: int): + self.rank = PM.rank_in(parallel) + self.world_size = PM.size_of(parallel) + self.group = PM.group_of(parallel) + + self.num_local_experts = num_experts // self.world_size + self.comm_record: ( + tuple[list[int], list[int], int, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None] | None + ) = None + self.last_dispatch_used_fused_permute = False + self.last_combine_used_fused_unpermute = False + + def _build_send_layout( + self, + token_expert_ids: torch.Tensor, + token_expert_weights: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[int], list[int], torch.Tensor, torch.Tensor]: + token_expert_ranks = token_expert_ids // self.num_local_experts + rank_permute_map, valid_mask, token_idx, dst_rank, send_sizes_tensor = _build_rank_permute_map( + token_expert_ranks, + self.world_size, + ) + gathered_sizes = _all_gather_splits(send_sizes_tensor, self.group, self.world_size) + recv_sizes_tensor = gathered_sizes[:, self.rank].contiguous() + + send_indices = token_expert_ids.index_select(0, token_idx) + send_probs = token_expert_weights.index_select(0, token_idx) + send_ranks = token_expert_ranks.index_select(0, token_idx) + + rank_offsets = (dst_rank.to(dtype=send_indices.dtype) * self.num_local_experts).unsqueeze(1) + local_mask = send_ranks == dst_rank.unsqueeze(1) + send_indices = send_indices - rank_offsets + send_indices = torch.where(local_mask, send_indices, torch.full_like(send_indices, -1)) + send_probs = torch.where(local_mask, send_probs, torch.zeros_like(send_probs)) + + send_sizes = send_sizes_tensor.cpu().tolist() + recv_sizes = recv_sizes_tensor.cpu().tolist() + return token_idx.to(torch.int64), send_indices, send_probs, send_sizes, recv_sizes, rank_permute_map, valid_mask + + def dispatch( + self, + hidden_states: torch.FloatTensor, + token_expert_ids: torch.IntTensor, + token_expert_weights: torch.FloatTensor, + ): + total_tokens = hidden_states.size(0) + self.last_dispatch_used_fused_permute = False + self.last_combine_used_fused_unpermute = False + + if self.world_size == 1: + local_ids = token_expert_ids.clone() + local_ids -= self.rank * self.num_local_experts + self.comm_record = ([], [], total_tokens, None, None, None) + return hidden_states, local_ids, token_expert_weights + + token_idx, send_indices, send_probs, send_sizes, recv_sizes, rank_permute_map, valid_mask = ( + self._build_send_layout( + token_expert_ids, + token_expert_weights, + ) + ) + send_hidden, sorted_indices, self.last_dispatch_used_fused_permute = _permute_hidden_states( + hidden_states, + token_idx, + rank_permute_map, + token_idx.numel(), + ) + send_token_ids = None + if not self.last_dispatch_used_fused_permute: + send_token_ids = token_idx.to(dtype=_pick_token_id_dtype(total_tokens)) + + recv_hidden = _all_to_all_single_tensor(send_hidden, recv_sizes, send_sizes, self.group) + recv_indices = _all_to_all_single_no_grad(send_indices, recv_sizes, send_sizes, self.group) + recv_probs = _all_to_all_single_tensor(send_probs, recv_sizes, send_sizes, self.group) + recv_token_ids = None + combine_probs = None + if self.last_dispatch_used_fused_permute: + combine_probs = _build_rank_combine_probs(valid_mask) + else: + recv_token_ids = _all_to_all_single_no_grad(send_token_ids, recv_sizes, send_sizes, self.group) + + del send_hidden, send_indices, send_probs, send_token_ids, rank_permute_map + self.comm_record = (send_sizes, recv_sizes, total_tokens, recv_token_ids, sorted_indices, combine_probs) + return recv_hidden, recv_indices, recv_probs + + def combine(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.comm_record is None: + raise RuntimeError("combine called before dispatch") + + send_sizes, recv_sizes, restore_tokens, token_ids, sorted_indices, combine_probs = self.comm_record + if self.world_size == 1: + self.comm_record = None + return hidden_states + + recv_hidden = _all_to_all_single_tensor(hidden_states, send_sizes, recv_sizes, self.group) + + if sorted_indices is not None and combine_probs is not None: + output = _unpermute_hidden_states( + recv_hidden, + sorted_indices, + combine_probs, + (restore_tokens, hidden_states.size(1)), + ) + self.last_combine_used_fused_unpermute = True + else: + recv_token_ids = _all_to_all_single_no_grad(token_ids, send_sizes, recv_sizes, self.group) + output = hidden_states.new_zeros((restore_tokens, hidden_states.size(1))) + if recv_hidden.numel() > 0: + output.index_add_(0, recv_token_ids.to(torch.int64), recv_hidden) + self.comm_record = None + return output diff --git a/steptronoss/model/ep_dispatcher/token_dispatcher.py b/steptronoss/model/ep_dispatcher/token_dispatcher.py index e3c90418..9971677e 100644 --- a/steptronoss/model/ep_dispatcher/token_dispatcher.py +++ b/steptronoss/model/ep_dispatcher/token_dispatcher.py @@ -8,9 +8,10 @@ from steptronoss.utils.optimizable import optimizable from .deepep_dispatcher import DeepEPDispatcher +from .npu_alltoall_dispatcher import NPUAllToAllDispatcher -@optimizable(alternatives={"deep_ep": DeepEPDispatcher}) +@optimizable(alternatives={"npu_alltoall": NPUAllToAllDispatcher, "deep_ep": DeepEPDispatcher}) class TokenDispatcher: """A dispatcher manager using torch.distributed collectives.""" diff --git a/steptronoss/model/optimizations/grouped_gemm/__init__.py b/steptronoss/model/optimizations/grouped_gemm/__init__.py index 5e116ab8..19a5ec20 100644 --- a/steptronoss/model/optimizations/grouped_gemm/__init__.py +++ b/steptronoss/model/optimizations/grouped_gemm/__init__.py @@ -1,4 +1,10 @@ from .function_imple import FunctionImpleGroupedGemm, function_imple_grouped_gemm +from .npu_gmm import mindspeed_npu_grouped_gemm_v2 from .triton import triton_grouped_gemm -__all__ = ["FunctionImpleGroupedGemm", "function_imple_grouped_gemm", "triton_grouped_gemm"] +__all__ = [ + "FunctionImpleGroupedGemm", + "function_imple_grouped_gemm", + "mindspeed_npu_grouped_gemm_v2", + "triton_grouped_gemm", +] diff --git a/steptronoss/model/optimizations/grouped_gemm/npu_gmm.py b/steptronoss/model/optimizations/grouped_gemm/npu_gmm.py new file mode 100644 index 00000000..efdb094c --- /dev/null +++ b/steptronoss/model/optimizations/grouped_gemm/npu_gmm.py @@ -0,0 +1,17 @@ +import torch + + +def mindspeed_npu_grouped_gemm_v2(mat_a_flat, mat_b, batch_sizes, trans_b=False): + try: + from mindspeed.ops.gmm import npu_gmm_v2 + except Exception as exc: + raise ImportError("from mindspeed.ops.gmm import npu_gmm_v2 failed.") from exc + + if mat_a_flat.shape[0] == 0: + return mat_a_flat.new_zeros((0, mat_b.shape[1] if trans_b else mat_b.shape[2])) + + weight = mat_b.transpose(-1, -2) if trans_b else mat_b + if batch_sizes.device.type != "npu": + batch_sizes = batch_sizes.to(device=mat_a_flat.device) + batch_sizes = batch_sizes.to(dtype=torch.int64) + return npu_gmm_v2(mat_a_flat, weight, bias=None, group_list=batch_sizes, group_type=0) diff --git a/steptronoss/model/utils/moe_utils.py b/steptronoss/model/utils/moe_utils.py index 2e880093..de079572 100644 --- a/steptronoss/model/utils/moe_utils.py +++ b/steptronoss/model/utils/moe_utils.py @@ -21,11 +21,6 @@ except: triton_routed_grouped_ffn_fused = None -try: - from steptronoss.model.optimizations.grouped_gemm.triton import triton_grouped_gemm -except: - function_imple_grouped_gemm = None - try: from steptronoss.model.optimizations.moe_routing.triton import ( triton_histogram, @@ -47,6 +42,11 @@ except: triton_grouped_gemm = None +try: + from steptronoss.model.optimizations.grouped_gemm.npu_gmm import mindspeed_npu_grouped_gemm_v2 +except: + mindspeed_npu_grouped_gemm_v2 = None + try: from steptronoss.model.optimizations.moe_gather.triton import triton_moe_weighted_gather except: @@ -436,6 +436,7 @@ def nv_grouped_gemm(mat_a_flat, mat_b, batch_sizes, trans_b=False): alternatives={ "nv_grouped_gemm": nv_grouped_gemm, "triton_grouped_gemm": triton_grouped_gemm, + "npu_gmm": mindspeed_npu_grouped_gemm_v2, "function_imple": function_imple_grouped_gemm, } ) diff --git a/steptronoss/utils/npu_patch.py b/steptronoss/utils/npu_patch.py new file mode 100644 index 00000000..d4b27211 --- /dev/null +++ b/steptronoss/utils/npu_patch.py @@ -0,0 +1,244 @@ +import os + +import torch +from loguru import logger + +from steptronoss.utils.patch_utils import StepTronPatchesManager + +NPU_PATCH_ACTIVE = False +GROUPED_GEMM_REGISTER_NAME = "steptronoss.model.utils.moe_utils.grouped_gemm" +TOKEN_DISPATCHER_REGISTER_NAME = "steptronoss.model.ep_dispatcher.token_dispatcher.TokenDispatcher" + + +def _dummy_compile( + model=None, *, fullgraph=False, dynamic=None, backend="inductor", mode=None, options=None, disable=False, **kwargs +): + if model is not None: + return model + + def identity(fn): + return fn + + return identity + + +def _torch_npu_available() -> bool: + try: + import torch_npu + except ImportError: + return False + return hasattr(torch, "npu") and callable(getattr(torch.npu, "is_available", None)) and torch.npu.is_available() + + +def is_npu_active() -> bool: + return NPU_PATCH_ACTIVE + + +def get_accel_module(): + return torch.npu if NPU_PATCH_ACTIVE else torch.cuda + + +def get_runtime_device_type() -> str: + return "npu" if NPU_PATCH_ACTIVE else "cuda" + + +def get_lazy_call(): + if NPU_PATCH_ACTIVE and hasattr(torch.npu, "_lazy_call"): + return torch.npu._lazy_call + return torch.cuda._lazy_call + + +def get_accel_device(device=-1): + device_type = get_runtime_device_type() + if device == -1: + return torch.device(device_type) + if isinstance(device, str): + return torch.device(device) + if isinstance(device, int): + return torch.device(device_type, device) + return device + + +def current_device() -> int: + return get_accel_module().current_device() + + +def device_count() -> int: + return get_accel_module().device_count() + + +def _set_accel_rng_state(new_state, device=-1): + module = get_accel_module() + lazy_call = get_lazy_call() + accel_device = get_accel_device(device) + + def cb(): + idx = accel_device.index + if idx is None: + idx = module.current_device() + # NOTE: `torch.cuda.default_generators[idx]` raises `IndexError` on Ascend. + default_generator = module.default_generators[idx] + default_generator.set_state(new_state) + + lazy_call(cb) + + +def _clip_grad_norm_fp32_npu(parameters, grads_for_norm, max_norm, norm_type=2, model_parallel_group=None): + from torch import inf + + from steptronoss.core.utils import multi_tensor_applier, multi_tensor_l2_norm, multi_tensor_scale + + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + if isinstance(grads_for_norm, torch.Tensor): + grads_for_norm = [grads_for_norm] + + accel_device = torch.device(get_runtime_device_type(), current_device()) + grads = [] + for param in parameters: + if param.grad is not None: + # NOTE: NPU grads are not `torch.cuda.FloatTensor`. + assert param.grad.dtype == torch.float32, param.grad.dtype + grads.append(param.grad.detach()) + + max_norm = float(max_norm) + norm_type = float(norm_type) + total_norm = 0.0 + if norm_type == inf: + total_norm = max(grad.abs().max() for grad in grads_for_norm) + total_norm_cuda = torch.tensor([float(total_norm)], device=accel_device, dtype=torch.float) + torch.distributed.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=model_parallel_group) + total_norm = total_norm_cuda[0].item() + else: + if norm_type == 2.0: + dummy_overflow_buf = torch.tensor([0], device=accel_device, dtype=torch.int) + if grads_for_norm: + grad_norm, _ = multi_tensor_applier(multi_tensor_l2_norm, dummy_overflow_buf, [grads_for_norm], False) + else: + grad_norm = torch.tensor([0], device=accel_device, dtype=torch.float) + total_norm = grad_norm**norm_type + else: + for grad in grads_for_norm: + grad_norm = torch.norm(grad, norm_type) + total_norm += grad_norm**norm_type + torch.distributed.all_reduce(total_norm, op=torch.distributed.ReduceOp.SUM, group=model_parallel_group) + total_norm = total_norm.item() ** (1.0 / norm_type) + + clip_coeff = max_norm / (total_norm + 1.0e-6) + if clip_coeff < 1.0: + dummy_overflow_buf = torch.tensor([0], device=accel_device, dtype=torch.int) + multi_tensor_applier(multi_tensor_scale, dummy_overflow_buf, [grads, grads], clip_coeff) + return total_norm + + +def _count_zeros_fp32_npu(parameters, model_parallel_group): + from steptronoss.core.tensor_parallel import ( + param_is_not_expert_parallel_duplicate, + param_is_not_tensor_parallel_duplicate, + ) + from steptronoss.model.module import param_is_not_shared + + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + + accel_device = torch.device(get_runtime_device_type(), current_device()) + total_num_zeros = torch.tensor([0], device=accel_device, dtype=torch.float) + for param in parameters: + grad_not_none = param.grad is not None + is_not_shared = param_is_not_shared(param) + is_moe_param = getattr(param, "expert_model_parallel", False) + if is_moe_param: + is_not_duplicate = param_is_not_expert_parallel_duplicate(param) + else: + is_not_duplicate = param_is_not_tensor_parallel_duplicate(param) + if grad_not_none and is_not_shared and is_not_duplicate: + grad = param.grad.detach() + num_zeros = grad.numel() - torch.count_nonzero(grad) + total_num_zeros = num_zeros + total_num_zeros + torch.distributed.all_reduce(total_num_zeros, op=torch.distributed.ReduceOp.SUM, group=model_parallel_group) + return total_num_zeros.item() + + +def _fp32_to_float16_npu(val, dtype=torch.bfloat16): + from torch.autograd import Variable + from torch.nn.parameter import Parameter + + from steptronoss.model.module import conversion_helper + + def half_conversion(item): + val_typecheck = item + if isinstance(val_typecheck, (Parameter, Variable)): + val_typecheck = item.data + # NOTE: NPU fp32 tensors skip the original CUDA-only type check. + if torch.is_tensor(val_typecheck) and val_typecheck.dtype == torch.float32: + item = item.to(dtype) + return item + + return conversion_helper(val, half_conversion) + + +def _float16_to_fp32_npu(val): + from torch.autograd import Variable + from torch.nn.parameter import Parameter + + from steptronoss.model.module import conversion_helper + + def float_conversion(item): + val_typecheck = item + if isinstance(val_typecheck, (Parameter, Variable)): + val_typecheck = item.data + # NOTE: NPU fp16/bf16 tensors skip the original CUDA-only type check. + if torch.is_tensor(val_typecheck) and val_typecheck.dtype in (torch.float16, torch.bfloat16): + item = item.float() + return item + + return conversion_helper(val, float_conversion) + + +def _apply_steptron_ascend_patches(): + logger.info("Applying StepTron Ascend patches.", at=0) + StepTronPatchesManager.register_patch( + "steptronoss.core.tensor_parallel.random._set_cuda_rng_state", + _set_accel_rng_state, + force_patch=True, + ) + StepTronPatchesManager.register_patch( + "steptronoss.core.utils._set_cuda_rng_state", + _set_accel_rng_state, + force_patch=True, + ) + StepTronPatchesManager.register_patch( + "steptronoss.optimizer.clip_grads.clip_grad_norm_fp32", + _clip_grad_norm_fp32_npu, + force_patch=True, + ) + StepTronPatchesManager.register_patch( + "steptronoss.optimizer.clip_grads.count_zeros_fp32", + _count_zeros_fp32_npu, + force_patch=True, + ) + StepTronPatchesManager.register_patch( + "steptronoss.model.module.fp32_to_float16", + _fp32_to_float16_npu, + force_patch=True, + ) + StepTronPatchesManager.register_patch( + "steptronoss.model.module.float16_to_fp32", + _float16_to_fp32_npu, + force_patch=True, + ) + StepTronPatchesManager.apply_patches() + + +def apply_npu_patch(): + global NPU_PATCH_ACTIVE + if not _torch_npu_available(): + raise RuntimeError("NPU patch cannot be applied because torch_npu is not available or NPU is not available.") + # Patch only once, even if `apply_npu_patch` is called multiple times. + if not NPU_PATCH_ACTIVE: + from torch_npu.contrib import transfer_to_npu + + NPU_PATCH_ACTIVE = True + if os.getenv("STEPTRON_ENABLE_TORCH_COMPILE_ON_NPU", "0") != "1": + torch.compile = _dummy_compile + _apply_steptron_ascend_patches() diff --git a/steptronoss/utils/patch_utils.py b/steptronoss/utils/patch_utils.py new file mode 100644 index 00000000..f09ef6d1 --- /dev/null +++ b/steptronoss/utils/patch_utils.py @@ -0,0 +1,122 @@ +import importlib +import inspect +import sys +import types + + +def get_func_name(func): + if isinstance(func, str): + return func + return ".".join((func.__module__, func.__qualname__)) + + +def dummy_function_wrapper(func_name): + def dummy_function(*args, **kwargs): + raise RuntimeError(f"function {func_name} no exist") + + return dummy_function + + +class Patch: + def __init__(self, orig_func_name, new_func, create_dummy): + split_name = orig_func_name.rsplit(".", 1) + if len(split_name) == 1: + self.orig_module_name, self.orig_func_name = orig_func_name, None + else: + self.orig_module_name, self.orig_func_name = split_name + self.orig_module = None + self.orig_func = None + self.patch_func = None + self.final_patch_func = None + self.wrappers = [] + if new_func is None: + new_func = dummy_function_wrapper(orig_func_name) + self.set_patch_func(new_func) + self.is_applied = False + self.create_dummy = create_dummy + + def set_patch_func(self, new_func, force_patch=False): + if hasattr(new_func, "__name__") and new_func.__name__.endswith(("wrapper", "decorator")): + if new_func not in self.wrappers: + self.wrappers.append(new_func) + else: + if self.patch_func and not force_patch: + raise RuntimeError(f"the patch of {self.orig_func_name} exist !") + self.patch_func = new_func + self.is_applied = False + + def apply_patch(self): + if self.is_applied: + return + + current_module, current_func = Patch.parse_path(self.orig_module_name, self.orig_func_name, self.create_dummy) + if self.orig_module is None: + self.orig_module, self.orig_func = current_module, current_func + + final_patch_func = self.orig_func + if self.patch_func is not None: + final_patch_func = self.patch_func + + for wrapper in self.wrappers: + final_patch_func = wrapper(final_patch_func) + + if self.orig_func_name is not None: + setattr(self.orig_module, self.orig_func_name, final_patch_func) + for _, value in sys.modules.copy().items(): + if ( + self.orig_func_name is not None + and hasattr(value, self.orig_func_name) + and getattr(value, self.orig_func_name) is current_func + ): + setattr(value, self.orig_func_name, final_patch_func) + self.is_applied = True + self.final_patch_func = final_patch_func + + @staticmethod + def parse_path(module_path, function_name, create_dummy): + from importlib.machinery import ModuleSpec + + modules = module_path.split(".") + for i in range(1, len(modules) + 1): + parent = ".".join(modules[: i - 1]) + path = ".".join(modules[:i]) + try: + importlib.import_module(path) + except ModuleNotFoundError as e: + if not parent or not hasattr(importlib.import_module(parent), modules[i - 1]): + if not create_dummy: + raise ModuleNotFoundError(e) from e + sys.modules[path] = types.ModuleType(path) + sys.modules[path].__file__ = "steptronoss.dummy_module.py" + sys.modules[path].__spec__ = ModuleSpec(path, None) + if parent: + setattr(importlib.import_module(parent), modules[i - 1], sys.modules[path]) + else: + module = getattr(importlib.import_module(parent), modules[i - 1]) + if hasattr(module, function_name): + return module, getattr(module, function_name) + if create_dummy: + return module, dummy_function_wrapper(function_name) + raise RuntimeError(f"no exist {function_name} of {module}") + + if function_name is not None and not hasattr(sys.modules[module_path], function_name): + setattr(sys.modules[module_path], function_name, None) + return sys.modules[module_path], getattr( + sys.modules[module_path], function_name + ) if function_name is not None else None + + +class StepTronPatchesManager: + patches_info: dict[str, Patch] = {} + + @staticmethod + def register_patch(orig_func_name, new_func=None, force_patch=False, create_dummy=False): + if orig_func_name not in StepTronPatchesManager.patches_info: + StepTronPatchesManager.patches_info[orig_func_name] = Patch(orig_func_name, new_func, create_dummy) + else: + StepTronPatchesManager.patches_info.get(orig_func_name).set_patch_func(new_func, force_patch) + + @staticmethod + def apply_patches(): + for patch in StepTronPatchesManager.patches_info.values(): + patch.apply_patch() diff --git a/tests/test_grouped_gemm_npu.py b/tests/test_grouped_gemm_npu.py new file mode 100644 index 00000000..1d14b4f3 --- /dev/null +++ b/tests/test_grouped_gemm_npu.py @@ -0,0 +1,74 @@ +import pytest +import torch + +pytestmark = pytest.mark.gpu + + +GROUPED_GEMM_REGISTER_NAME = "steptronoss.model.utils.moe_utils.grouped_gemm" + + +def _npu_available() -> bool: + return hasattr(torch, "npu") and callable(getattr(torch.npu, "is_available", None)) and torch.npu.is_available() + + +def _ref_gmm(a: torch.Tensor, b: torch.Tensor, batch_sizes: torch.Tensor, trans_b: bool) -> torch.Tensor: + batch_sizes_list = batch_sizes.tolist() + outputs = [] + start = 0 + for i, size in enumerate(batch_sizes_list): + rhs = b[i].t() if trans_b else b[i] + outputs.append(a[start : start + size] @ rhs) + start += size + if outputs: + return torch.cat(outputs, dim=0) + return a.new_zeros((0, b.shape[1] if trans_b else b.shape[2])) + + +@pytest.mark.skipif(not _npu_available(), reason="grouped_gemm npu_gmm test requires NPU") +@pytest.mark.parametrize("trans_b", [False, True]) +def test_npu_gmm_forward_backward(trans_b: bool): + import steptronoss.utils.npu_patch as npu_patch + + npu_patch.apply_npu_patch() + + from steptronoss.model.utils import grouped_gemm + from steptronoss.utils.optimizable import OPTIMIZABLE_REGISTER, set_optimization + + if not npu_patch.is_npu_active(): + pytest.skip("StepTron NPU patch is not active") + + register = OPTIMIZABLE_REGISTER[GROUPED_GEMM_REGISTER_NAME] + assert "npu_gmm" in register["alternatives"] + print(f"Using npu_gmm optimization: {register['alternatives']['npu_gmm']}") + assert register["alternatives"]["npu_gmm"].__name__ == "mindspeed_npu_grouped_gemm_v2" + + torch.manual_seed(0) + device = torch.device("cuda") + batch_sizes = torch.tensor([3, 1, 4], dtype=torch.int64) + k = 32 + n = 24 + + a = torch.randn(batch_sizes.sum().item(), k, device=device, dtype=torch.bfloat16, requires_grad=True) + if trans_b: + b = torch.randn(batch_sizes.numel(), n, k, device=device, dtype=torch.bfloat16, requires_grad=True) + else: + b = torch.randn(batch_sizes.numel(), k, n, device=device, dtype=torch.bfloat16, requires_grad=True) + + a_ref = a.detach().clone().requires_grad_(True) + b_ref = b.detach().clone().requires_grad_(True) + + try: + set_optimization(grouped_gemm="npu_gmm") + out = grouped_gemm(a, b, batch_sizes, trans_b=trans_b) + expected = _ref_gmm(a_ref, b_ref, batch_sizes, trans_b) + torch.testing.assert_close(out, expected, rtol=1e-2, atol=1e-2) + + loss = out.float().square().mean() + expected_loss = expected.float().square().mean() + loss.backward() + expected_loss.backward() + + torch.testing.assert_close(a.grad, a_ref.grad, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(b.grad, b_ref.grad, rtol=1e-2, atol=1e-2) + finally: + set_optimization(grouped_gemm=None) diff --git a/tests/test_npu_alltoall_dispatcher.py b/tests/test_npu_alltoall_dispatcher.py new file mode 100644 index 00000000..aa74f94c --- /dev/null +++ b/tests/test_npu_alltoall_dispatcher.py @@ -0,0 +1,113 @@ +import torch + +from steptronoss.model.ep_dispatcher.npu_alltoall_dispatcher import ( + _build_rank_permute_map, + _build_unique_token_rank_pairs, + _pick_token_id_dtype, + _stable_bucket_order_by_rank, +) + + +def _baseline_unique_token_rank_pairs(token_expert_ranks: torch.Tensor, world_size: int): + token_idx_parts = [] + dst_rank_parts = [] + send_sizes = [] + + for rank in range(world_size): + this_rank = (token_expert_ranks == rank).any(dim=1) + token_idx = torch.nonzero(this_rank, as_tuple=True)[0] + token_idx_parts.append(token_idx) + dst_rank_parts.append(torch.full_like(token_idx, rank, dtype=torch.int64)) + send_sizes.append(int(token_idx.numel())) + + if token_idx_parts: + token_idx = torch.cat(token_idx_parts, dim=0) + dst_rank = torch.cat(dst_rank_parts, dim=0) + else: + token_idx = torch.empty((0,), dtype=torch.int64) + dst_rank = torch.empty((0,), dtype=torch.int64) + + return token_idx, dst_rank, torch.tensor(send_sizes, dtype=torch.int64) + + +def test_build_unique_token_rank_pairs_matches_baseline(): + generator = torch.Generator().manual_seed(1234) + + for world_size in (1, 2, 4, 8): + for token_num in (0, 1, 7, 17): + for topk in (1, 2, 4): + token_expert_ranks = torch.randint( + -1, + world_size, + (token_num, topk), + generator=generator, + dtype=torch.int64, + ) + got = _build_unique_token_rank_pairs(token_expert_ranks, world_size) + expected = _baseline_unique_token_rank_pairs(token_expert_ranks, world_size) + + assert torch.equal(got[0], expected[0]) + assert torch.equal(got[1], expected[1]) + assert torch.equal(got[2], expected[2]) + + +def test_build_unique_token_rank_pairs_deduplicates_same_rank(): + token_expert_ranks = torch.tensor( + [ + [0, 0, 1, 1], + [2, 2, 2, -1], + [3, 1, 3, 1], + ], + dtype=torch.int64, + ) + + token_idx, dst_rank, send_sizes = _build_unique_token_rank_pairs(token_expert_ranks, world_size=4) + + assert torch.equal(token_idx, torch.tensor([0, 0, 2, 1, 2], dtype=torch.int64)) + assert torch.equal(dst_rank, torch.tensor([0, 1, 1, 2, 3], dtype=torch.int64)) + assert torch.equal(send_sizes, torch.tensor([1, 2, 1, 1], dtype=torch.int64)) + + +def test_build_rank_permute_map_matches_fused_permute_sort_contract(): + token_expert_ranks = torch.tensor( + [ + [0, 2, 2, 0], + [1, 1, 3, 3], + [0, 1, 2, 3], + ], + dtype=torch.int64, + ) + world_size = 4 + + rank_permute_map, valid_mask, token_idx, dst_rank, send_sizes = _build_rank_permute_map( + token_expert_ranks, world_size + ) + + flat_rank_map = rank_permute_map.reshape(-1).to(torch.int64) + flat_sorted_pos = flat_rank_map.argsort(stable=True) + valid_count = int(valid_mask.sum().item()) + selected_pos = flat_sorted_pos[:valid_count] + + expected_token_idx = torch.div(selected_pos, token_expert_ranks.size(1), rounding_mode="floor") + expected_dst_rank = flat_rank_map.index_select(0, selected_pos) + + assert torch.equal(token_idx, expected_token_idx) + assert torch.equal(dst_rank, expected_dst_rank) + assert torch.equal(send_sizes, torch.bincount(expected_dst_rank, minlength=world_size)) + + +def test_stable_bucket_order_matches_stable_argsort(): + generator = torch.Generator().manual_seed(4321) + + for world_size in (1, 2, 4, 8): + for token_num in (0, 1, 7, 31, 127): + dst_rank = torch.randint(0, world_size, (token_num,), generator=generator, dtype=torch.int64) + expected = dst_rank.argsort(stable=True) + got = _stable_bucket_order_by_rank(dst_rank, world_size) + assert torch.equal(got, expected) + + +def test_pick_token_id_dtype_prefers_int32_when_safe(): + assert _pick_token_id_dtype(0) is torch.int32 + assert _pick_token_id_dtype(1024) is torch.int32 + assert _pick_token_id_dtype(torch.iinfo(torch.int32).max + 1) is torch.int64