add tests

test__c_lookup_table/cases.py

@@ -0,0 +1,5 @@
+cases = [
+    (8, "float32", (32032, 12288), (1, 8192), 0, 256256),
+    (8, "float16", (32032, 12288), (1, 8192), 0, 256256),
+    (8, "bfloat16", (32032, 12288), (1, 8192), 0, 256256),
+]

test__c_lookup_table/check.py

@@ -0,0 +1,72 @@
+import sys
+
+sys.path.append("..")
+import numpy as np
+from utils import (
+    TOLERANCE,
+    convert_dtype_to_torch_type,
+    np_assert_accuracy,
+    np_assert_staility,
+)
+import hashlib
+from pathlib import Path
+from cases import cases
+import multiprocess
+
+
+def check(ranks, dtype, table_shape, index_shape, start_index, vacab_size):
+    atol = TOLERANCE[dtype]["atol"]
+    rtol = TOLERANCE[dtype]["rtol"]
+    dataname = hashlib.md5(
+        f"{ranks},{dtype},{table_shape},{index_shape},{start_index},{vacab_size}".encode(
+            "ascii"
+        )
+    ).digest()
+    dataname = dataname.hex()[:6]
+    datadir = Path(__file__).resolve().parent.joinpath("data").joinpath(dataname)
+    paddle_dir = (
+        Path(__file__).resolve().parent.joinpath("paddle_value").joinpath(dataname)
+    )
+    torch_dir = (
+        Path(__file__).resolve().parent.joinpath("torch_value").joinpath(dataname)
+    )
+
+    out_paddle = np.load(paddle_dir.joinpath(f"out.npy"))
+    out_torch = np.load(torch_dir.joinpath(f"out.npy"))
+    np_assert_accuracy(
+        out_paddle,
+        out_torch,
+        atol,
+        rtol,
+        dtype,
+        version_a="paddle_develop",
+        version_b="torch",
+        eager_or_static_mode="eager",
+        fwd_or_bkd="forward",
+        api="_c_lookup_table",
+    )
+    table_grad_paddle = []
+
+    for i in range(0, ranks):
+        table_grad_i_paddle = np.load(paddle_dir.joinpath(f"table_grad{i}.npy"))
+        table_grad_paddle.append(table_grad_i_paddle)
+    table_grad_paddle = np.concatenate(table_grad_paddle, axis=0)
+    table_grad_torch = np.load(torch_dir.joinpath(f"table_grad.npy"))
+    np_assert_accuracy(
+        table_grad_paddle,
+        table_grad_torch,
+        atol,
+        rtol,
+        dtype,
+        version_a="paddle_develop",
+        version_b="torch",
+        eager_or_static_mode="eager",
+        fwd_or_bkd="backward",
+        api="_c_lookup_table",
+    )
+    print("OK")
+
+
+if __name__ == "__main__":
+    args = cases[int(sys.argv[1])]
+    check(*args)

test__c_lookup_table/gen_np_inputs.py

@@ -0,0 +1,34 @@
+import numpy as np
+from pathlib import Path
+import hashlib
+import sys
+from cases import cases
+
+
+def gen_np_inputs(ranks, dtype, table_shape, index_shape, start_index, vacab_size):
+
+    dataname = hashlib.md5(
+        f"{ranks},{dtype},{table_shape},{index_shape},{start_index},{vacab_size}".encode(
+            "ascii"
+        )
+    ).digest()
+    if dtype == "bfloat16":
+        dtype = "float32"
+    dataname = dataname.hex()[:6]
+    datadir = Path(__file__).resolve().parent.joinpath("data").joinpath(dataname)
+    datadir.mkdir(parents=True, exist_ok=True)
+    for i in range(ranks):
+        table = np.random.random(size=table_shape).astype(dtype)
+        np.save(datadir.joinpath(f"table{i}"), table)
+    index = np.random.randint(
+        low=0, high=table_shape[0] * ranks, size=index_shape
+    ).astype("int64")
+    np.save(datadir.joinpath(f"index"), index)
+    out_grad = np.random.random(size=(*index_shape, table_shape[-1])).astype(dtype)
+    np.save(datadir.joinpath(f"out_grad"), out_grad)
+
+
+if __name__ == "__main__":
+    np.random.seed(2024)
+    args = cases[int(sys.argv[1])]
+    gen_np_inputs(*args)

test__c_lookup_table/run__c_lookup_table.py

@@ -0,0 +1,127 @@
+#   Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import unittest
+
+import numpy as np
+import torch
+
+import paddle
+from paddle.utils import map_structure
+from paddle.distributed.fleet.layers.mpu.mp_ops import _c_lookup_table
+from paddle.distributed.fleet.layers.mpu import mp_ops
+from paddle.distributed.fleet.layers.mpu.mp_layers import VocabParallelEmbedding
+import paddle.distributed as dist
+from paddle.distributed import collective
+from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
+import paddle.distributed.fleet as fleet
+import random
+
+sys.path.append("..")
+from utils import (
+    TOLERANCE,
+    convert_dtype_to_torch_type,
+    np_assert_accuracy,
+    np_assert_staility,
+)
+import hashlib
+from pathlib import Path
+from cases import cases
+
+
+def set_random_seed(seed):
+    """Set random seed for reproducability."""
+    random.seed(seed)
+    np.random.seed(seed)
+    paddle.seed(seed)
+    fleet.meta_parallel.model_parallel_random_seed(seed)
+
+
+def run__c_lookup_table(
+    ranks, dtype, table_shape, index_shape, start_index, vacab_size
+):
+    dataname = hashlib.md5(
+        f"{ranks},{dtype},{table_shape},{index_shape},{start_index},{vacab_size}".encode(
+            "ascii"
+        )
+    ).digest()
+    dataname = dataname.hex()[:6]
+    datadir = Path(__file__).resolve().parent.joinpath("data").joinpath(dataname)
+    resultdir = (
+        Path(__file__).resolve().parent.joinpath("paddle_value").joinpath(dataname)
+    )
+    if dist.get_rank() == 0:
+        resultdir.mkdir(parents=True)
+
+    table = np.load(datadir.joinpath(f"table{dist.get_rank()}.npy"))
+    index = np.load(datadir.joinpath(f"index.npy"))
+    out_grad = np.load(datadir.joinpath(f"out_grad.npy"))
+    index = paddle.to_tensor(index, dtype="int64", stop_gradient=False)
+    if dtype == "float32":
+        table = paddle.to_tensor(table, dtype="float32", stop_gradient=False)
+        out_grad = paddle.to_tensor(out_grad, dtype="float32", stop_gradient=False)
+        table = paddle.cast(table, "uint16")
+        out_grad = paddle.cast(out_grad, "uint16")
+    else:
+        table = paddle.to_tensor(table, dtype=dtype, stop_gradient=False)
+        out_grad = paddle.to_tensor(out_grad, dtype=dtype, stop_gradient=False)
+
+    group = collective._get_default_group()
+
+    # out = _c_lookup_table(
+    #     table,
+    #     index,
+    #     start_index=dist.get_rank() * table.shape[0],
+    #     vocab_size=vacab_size,
+    # )
+    # out = mp_ops._mp_allreduce(
+    #     out,
+    #     group=group,
+    #     use_calc_stream=True,
+    #     use_model_parallel=True,
+    # )
+
+    embedding = VocabParallelEmbedding(vacab_size, table.shape[-1], mp_group=group)
+    embedding.weight.data = table
+    out = embedding(index)
+
+    out.backward(out_grad)
+    table_grad = embedding.weight.grad
+    # table_grad = table.grad.numpy()
+
+    if dtype == "bfloat16":
+        out = paddle.cast(out, "float32")
+        table_grad = paddle.cast(table_grad, "float32")
+
+    out = out.numpy()
+    table_grad = table_grad.numpy()
+
+    np.save(resultdir.joinpath(f"out.npy"), out)
+    np.save(resultdir.joinpath(f"table_grad{dist.get_rank()}.npy"), table_grad)
+
+
+if __name__ == "__main__":
+    dist_strategy = fleet.DistributedStrategy()
+    world_size = dist.get_world_size()
+    dist_strategy.hybrid_configs = {
+        "mp_degree": world_size,
+        "pp_degree": 1,
+        "dp_degree": 1,
+    }
+    dist.fleet.init(is_collective=True, strategy=dist_strategy)
+
+    set_random_seed(1024)
+    args = cases[int(sys.argv[1])]
+    run__c_lookup_table(*args)

test__c_lookup_table/run_torch.py

@@ -0,0 +1,62 @@
+import torch
+import hashlib
+from pathlib import Path
+import numpy as np
+from torch.distributed.tensor.parallel import loss_parallel
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed._tensor import Shard, distribute_tensor, Replicate, DTensor
+import os
+import sys
+from cases import cases
+from torch import nn
+
+
+def run_torch(ranks, dtype, table_shape, index_shape, start_index, vacab_size):
+    dataname = hashlib.md5(
+        f"{ranks},{dtype},{table_shape},{index_shape},{start_index},{vacab_size}".encode(
+            "ascii"
+        )
+    ).digest()
+    dataname = dataname.hex()[:6]
+    datadir = Path(__file__).resolve().parent.joinpath("data").joinpath(dataname)
+    resultdir = (
+        Path(__file__).resolve().parent.joinpath("torch_value").joinpath(dataname)
+    )
+    os.makedirs(resultdir)
+
+    table = []
+    for i in range(ranks):
+        table.append(np.load(datadir.joinpath(f"table{i}.npy")))
+    table = np.concatenate(table, axis=0)
+    index = np.load(datadir.joinpath(f"index.npy"))
+
+    table = torch.tensor(table, device="cuda", requires_grad=True)
+    index = torch.tensor(index, device="cuda")
+    out_grad = np.load(datadir.joinpath(f"out_grad.npy"))
+    out_grad = torch.tensor(out_grad, device="cuda", requires_grad=True)
+    if dtype == "bfloat16":
+        table = table.bfloat16()
+        out_grad = out_grad.bfloat16()
+
+    embedding = nn.Embedding(vacab_size, table_shape[1])
+    embedding.weight = torch.nn.Parameter(table, requires_grad=True)
+    embedding.zero_grad()
+    out = embedding(index)
+    out.backward(out_grad)
+
+    out = out.detach()
+    table_grad = embedding.weight.grad.detach()
+    if dtype == "bfloat16":
+        out = out.cpu().to(torch.float32)
+        table_grad = table_grad.cpu().to(torch.float32)
+
+    np.save(resultdir.joinpath(f"out.npy"), out.numpy())
+    np.save(
+        resultdir.joinpath(f"table_grad.npy"),
+        table_grad.numpy(),
+    )
+
+
+if __name__ == "__main__":
+    args = cases[int(sys.argv[1])]
+    run_torch(*args)

test__c_lookup_table/test.sh

@@ -0,0 +1,14 @@
+rm -rf ./data
+rm -rf ./log
+rm -rf ./torch_value
+rm -rf ./paddle_value
+rm -f ./out_torch.log ./out_paddle.log
+cases=(0 1 2)
+for case in ${cases[@]}
+do
+    echo "Testing Case $case"
+    python gen_np_inputs.py $case
+    python run_torch.py $case 1>>out_torch.log 2>>out_torch.log
+    python -m paddle.distributed.launch run__c_lookup_table.py $case 1>>out_paddle.log 2>>out_paddle.log
+    python check.py $case
+done

test__c_softmax_with_cross_entropy/cases.py

@@ -0,0 +1 @@
+cases = [(8, "float32", (1, 8192, 32032), (1, 8192, 1))]

test__c_softmax_with_cross_entropy/check.py

@@ -0,0 +1,63 @@
+import sys
+sys.path.append("..")
+import numpy as np
+from utils import (
+    TOLERANCE,
+    convert_dtype_to_torch_type,
+    np_assert_accuracy,
+    np_assert_staility,
+)
+import hashlib
+from pathlib import Path
+from cases import cases
+
+
+def check(ranks, dtype, logits_shape, label_shape):
+    atol = TOLERANCE[dtype]["atol"]
+    rtol = TOLERANCE[dtype]["rtol"]
+    dataname = hashlib.md5(
+        f"{ranks},{dtype},{logits_shape},{label_shape}".encode("ascii")
+    ).digest()
+    dataname = dataname.hex()[:6]
+    datadir = Path(__file__).resolve().parent.joinpath("data").joinpath(dataname)
+    paddle_dir = (
+        Path(__file__).resolve().parent.joinpath("paddle_value").joinpath(dataname)
+    )
+    torch_dir = (
+        Path(__file__).resolve().parent.joinpath("torch_value").joinpath(dataname)
+    )
+    loss_paddle = np.load(paddle_dir.joinpath(f"loss.npy"))
+    loss_torch = np.load(torch_dir.joinpath(f"loss.npy"))
+    np_assert_accuracy(
+        loss_paddle,
+        loss_torch,
+        atol,
+        rtol,
+        dtype,
+        version_a="paddle_develop",
+        version_b="torch",
+        eager_or_static_mode="eager",
+        fwd_or_bkd="forward",
+        api="_c_softmax_with_cross_entropy",
+    )
+    for i in range(0, ranks):
+        logits_grad_paddle = np.load(paddle_dir.joinpath(f"logits_grad{i}.npy"))
+        logits_grad_torch = np.load(torch_dir.joinpath(f"logits_grad{i}.npy"))
+        np_assert_accuracy(
+            logits_grad_paddle,
+            logits_grad_torch,
+            atol,
+            rtol,
+            dtype,
+            version_a="paddle_develop",
+            version_b="torch",
+            eager_or_static_mode="eager",
+            fwd_or_bkd="backward",
+            api="_c_softmax_with_cross_entropy",
+        )
+    print("OK")
+
+
+if __name__ == "__main__":
+    args = cases[int(sys.argv[1])]
+    check(*args)