diff --git a/src/code/issue9/README.md b/src/code/issue9/README.md new file mode 100644 index 0000000..6090693 --- /dev/null +++ b/src/code/issue9/README.md @@ -0,0 +1,22 @@ +# Issue 9 单机可完成部分:多表 AlltoAllv 聚合参考实现 + +`fused_alltoallv_reference.py` 把各表数据按“destination-major、table-minor”顺序 +打成一个字节 buffer,使每个目标 rank 只需要一个连续 segment。接收端根据每表 +dtype、尾部 shape 和 split 元数据恢复结果。 + +本机模拟覆盖: + +- 多张表 first-dim 不同; +- 表间 hidden shape 和 dtype 不同; +- split 中存在零长度消息; +- 接收顺序与逐表 AlltoAllv 一致,并逐位相等。 + +```bash +python -m pytest src/code/issue9/test_fused_alltoallv_reference.py +python src/code/issue9/fused_alltoallv_reference.py +``` + +默认演示把 8 次逐表 launch 合成 1 次,模型控制面开销下降 87.5%。这只是 CPU +参考实现与启动开销模型;下一阶段需要把同一布局接到 NCCL/DeepEP 数据面,在多机多卡 +实测 launch 数、协商开销与端到端时间。 + diff --git a/src/code/issue9/fused_alltoallv_reference.py b/src/code/issue9/fused_alltoallv_reference.py new file mode 100644 index 0000000..b29e85a --- /dev/null +++ b/src/code/issue9/fused_alltoallv_reference.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""Byte-exact reference for fusing multiple table-wise AlltoAllv calls.""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Sequence + +import numpy as np + + +@dataclass(frozen=True) +class TableDescriptor: + dtype: str + trailing_shape: tuple[int, ...] + splits: tuple[int, ...] + row_bytes: int + + +@dataclass +class PackedRank: + buffer: np.ndarray + destination_offsets: tuple[int, ...] + tables: tuple[TableDescriptor, ...] + + @property + def world_size(self) -> int: + return len(self.destination_offsets) - 1 + + @property + def destination_sizes(self) -> tuple[int, ...]: + return tuple( + self.destination_offsets[index + 1] - self.destination_offsets[index] + for index in range(self.world_size) + ) + + +def _descriptor(table: np.ndarray, splits: Sequence[int], world_size: int) -> TableDescriptor: + if table.ndim < 1: + raise ValueError("each table must have at least one dimension") + if len(splits) != world_size or any(split < 0 for split in splits): + raise ValueError("each split vector must contain world_size non-negative entries") + if sum(splits) != table.shape[0]: + raise ValueError("split sizes must sum to the table's first dimension") + row_elements = int(np.prod(table.shape[1:], dtype=np.int64)) if table.ndim > 1 else 1 + return TableDescriptor( + dtype=table.dtype.str, + trailing_shape=tuple(table.shape[1:]), + splits=tuple(int(split) for split in splits), + row_bytes=row_elements * table.dtype.itemsize, + ) + + +def pack_rank( + tables: Sequence[np.ndarray], splits_per_table: Sequence[Sequence[int]], world_size: int +) -> PackedRank: + """Pack one rank's tables in destination-major, table-minor order.""" + + if world_size <= 0 or not tables or len(tables) != len(splits_per_table): + raise ValueError("world_size must be positive and every table needs a split vector") + contiguous = [np.ascontiguousarray(table) for table in tables] + descriptors = [ + _descriptor(table, splits, world_size) + for table, splits in zip(contiguous, splits_per_table, strict=True) + ] + table_starts = [np.cumsum((0,) + descriptor.splits) for descriptor in descriptors] + + chunks: list[np.ndarray] = [] + destination_offsets = [0] + for destination in range(world_size): + for table, descriptor, starts in zip( + contiguous, descriptors, table_starts, strict=True + ): + row_start = int(starts[destination]) + row_end = int(starts[destination + 1]) + chunk = table[row_start:row_end].reshape(-1).view(np.uint8) + chunks.append(chunk) + destination_offsets.append(sum(chunk.nbytes for chunk in chunks)) + + buffer = np.concatenate(chunks) if chunks else np.empty(0, dtype=np.uint8) + return PackedRank(buffer, tuple(destination_offsets), tuple(descriptors)) + + +def _validate_compatible(packed_ranks: Sequence[PackedRank]) -> tuple[int, int]: + if not packed_ranks: + raise ValueError("at least one packed rank is required") + world_size = len(packed_ranks) + if any(packed.world_size != world_size for packed in packed_ranks): + raise ValueError("packed rank count must equal every split vector's world size") + num_tables = len(packed_ranks[0].tables) + for packed in packed_ranks[1:]: + if len(packed.tables) != num_tables: + raise ValueError("all ranks must contain the same number of tables") + for reference, candidate in zip(packed_ranks[0].tables, packed.tables, strict=True): + if (reference.dtype, reference.trailing_shape) != ( + candidate.dtype, + candidate.trailing_shape, + ): + raise ValueError("matching tables must use the same dtype and trailing shape") + return world_size, num_tables + + +def simulate_fused_alltoallv(packed_ranks: Sequence[PackedRank]) -> list[list[np.ndarray]]: + """Simulate one fused byte-buffer AlltoAllv and unpack its table results.""" + + world_size, num_tables = _validate_compatible(packed_ranks) + outputs: list[list[np.ndarray]] = [] + for destination in range(world_size): + destination_tables: list[np.ndarray] = [] + for table_index in range(num_tables): + chunks = [] + for source in packed_ranks: + descriptor = source.tables[table_index] + segment_start = source.destination_offsets[destination] + table_offset = sum( + source.tables[index].splits[destination] + * source.tables[index].row_bytes + for index in range(table_index) + ) + rows = descriptor.splits[destination] + byte_count = rows * descriptor.row_bytes + raw = source.buffer[ + segment_start + table_offset : segment_start + table_offset + byte_count + ] + shape = (rows,) + descriptor.trailing_shape + chunks.append(np.frombuffer(raw, dtype=np.dtype(descriptor.dtype)).reshape(shape).copy()) + descriptor = packed_ranks[0].tables[table_index] + if chunks: + destination_tables.append(np.concatenate(chunks, axis=0)) + else: + destination_tables.append( + np.empty((0,) + descriptor.trailing_shape, dtype=np.dtype(descriptor.dtype)) + ) + outputs.append(destination_tables) + return outputs + + +def simulate_separate_alltoallv( + rank_tables: Sequence[Sequence[np.ndarray]], + rank_splits: Sequence[Sequence[Sequence[int]]], +) -> list[list[np.ndarray]]: + """Reference semantics for launching one AlltoAllv per table.""" + + world_size = len(rank_tables) + packed = [ + pack_rank(tables, splits, world_size) + for tables, splits in zip(rank_tables, rank_splits, strict=True) + ] + _, num_tables = _validate_compatible(packed) + outputs: list[list[np.ndarray]] = [[] for _ in range(world_size)] + for destination in range(world_size): + for table_index in range(num_tables): + chunks = [] + for source_rank in range(world_size): + table = np.ascontiguousarray(rank_tables[source_rank][table_index]) + splits = rank_splits[source_rank][table_index] + start = sum(splits[:destination]) + end = start + splits[destination] + chunks.append(table[start:end]) + outputs[destination].append(np.concatenate(chunks, axis=0)) + return outputs + + +def demo_report(world_size: int, num_tables: int, launch_overhead_us: float) -> dict[str, Any]: + if world_size <= 0 or num_tables <= 0 or launch_overhead_us < 0: + raise ValueError("invalid demo parameters") + rng = np.random.default_rng(20260726) + rank_tables = [] + rank_splits = [] + for rank in range(world_size): + tables = [] + splits_for_rank = [] + for table_index in range(num_tables): + splits = rng.integers(0, 6, size=world_size, dtype=np.int64) + if splits.sum() == 0: + splits[rank] = 1 + rows = int(splits.sum()) + hidden = 4 * (table_index + 1) + dtype = np.float16 if table_index % 2 == 0 else np.float32 + values = np.arange(rows * hidden, dtype=dtype).reshape(rows, hidden) + rank * 100 + tables.append(values) + splits_for_rank.append(splits.tolist()) + rank_tables.append(tables) + rank_splits.append(splits_for_rank) + + packed = [ + pack_rank(tables, splits, world_size) + for tables, splits in zip(rank_tables, rank_splits, strict=True) + ] + fused = simulate_fused_alltoallv(packed) + separate = simulate_separate_alltoallv(rank_tables, rank_splits) + bitwise_equal = all( + np.array_equal(fused[destination][table], separate[destination][table]) + for destination in range(world_size) + for table in range(num_tables) + ) + payload_bytes = sum(item.buffer.nbytes for item in packed) + separate_control_us = num_tables * launch_overhead_us + fused_control_us = launch_overhead_us + return { + "schema_version": 1, + "world_size": world_size, + "num_tables": num_tables, + "bitwise_equal_to_separate": bitwise_equal, + "payload_bytes": payload_bytes, + "per_rank_destination_bytes": [list(item.destination_sizes) for item in packed], + "launches": {"separate": num_tables, "fused": 1}, + "modeled_control_us": {"separate": separate_control_us, "fused": fused_control_us}, + "modeled_control_saving_ratio": ( + 0.0 if separate_control_us == 0 else 1.0 - fused_control_us / separate_control_us + ), + "warning": "control-time saving is a launch-overhead model, not a multi-rank measurement", + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Validate multi-table AlltoAllv packing") + parser.add_argument("--world-size", type=int, default=4) + parser.add_argument("--num-tables", type=int, default=8) + parser.add_argument("--launch-overhead-us", type=float, default=8.0) + parser.add_argument( + "--output", type=Path, default=Path("src/code/issue9/results/fusion_reference.json") + ) + return parser + + +def main() -> int: + args = build_parser().parse_args() + report = demo_report(args.world_size, args.num_tables, args.launch_overhead_us) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + print(json.dumps(report, ensure_ascii=False, indent=2)) + print(f"Wrote {args.output}") + return 0 if report["bitwise_equal_to_separate"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/code/issue9/results/fusion_reference.json b/src/code/issue9/results/fusion_reference.json new file mode 100644 index 0000000..a5035d7 --- /dev/null +++ b/src/code/issue9/results/fusion_reference.json @@ -0,0 +1,43 @@ +{ + "schema_version": 1, + "world_size": 4, + "num_tables": 8, + "bitwise_equal_to_separate": true, + "payload_bytes": 19160, + "per_rank_destination_bytes": [ + [ + 808, + 1200, + 1264, + 1368 + ], + [ + 1072, + 1296, + 1392, + 912 + ], + [ + 1232, + 1304, + 1032, + 1192 + ], + [ + 1040, + 1560, + 1248, + 1240 + ] + ], + "launches": { + "separate": 8, + "fused": 1 + }, + "modeled_control_us": { + "separate": 64.0, + "fused": 8.0 + }, + "modeled_control_saving_ratio": 0.875, + "warning": "control-time saving is a launch-overhead model, not a multi-rank measurement" +} diff --git a/src/code/issue9/test_fused_alltoallv_reference.py b/src/code/issue9/test_fused_alltoallv_reference.py new file mode 100644 index 0000000..08f774a --- /dev/null +++ b/src/code/issue9/test_fused_alltoallv_reference.py @@ -0,0 +1,41 @@ +import numpy as np + +from fused_alltoallv_reference import ( + demo_report, + pack_rank, + simulate_fused_alltoallv, + simulate_separate_alltoallv, +) + + +def test_fused_matches_separate_with_mixed_shapes_dtypes_and_zero_splits(): + rank_tables = [ + [ + np.arange(12, dtype=np.float32).reshape(3, 4), + np.arange(10, dtype=np.int16).reshape(5, 2), + ], + [ + (np.arange(16, dtype=np.float32) + 100).reshape(4, 4), + (np.arange(4, dtype=np.int16) + 100).reshape(2, 2), + ], + ] + rank_splits = [ + [[1, 2], [0, 5]], + [[3, 1], [2, 0]], + ] + packed = [pack_rank(rank_tables[index], rank_splits[index], 2) for index in range(2)] + + fused = simulate_fused_alltoallv(packed) + separate = simulate_separate_alltoallv(rank_tables, rank_splits) + + for destination in range(2): + for table in range(2): + np.testing.assert_array_equal(fused[destination][table], separate[destination][table]) + + +def test_demo_reduces_launch_count_and_preserves_bits(): + report = demo_report(world_size=4, num_tables=8, launch_overhead_us=8) + assert report["bitwise_equal_to_separate"] is True + assert report["launches"] == {"separate": 8, "fused": 1} + assert report["modeled_control_saving_ratio"] == 0.875 +