wehub-resource-sync
125 lines
5.5 KiB
Python
125 lines
5.5 KiB
Python
# Copyright (c) 2022 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 unittest
|
|
|
|
from paddle.distributed.auto_parallel.static.cluster_v2 import DeviceMesh
|
|
from paddle.framework import core
|
|
|
|
|
|
class TestDeviceMesh(unittest.TestCase):
|
|
def test_device_mesh(self):
|
|
name = "my_device_mesh"
|
|
mesh = [[0, 1, 2], [3, 4, 5]]
|
|
device_mesh = DeviceMesh(name, mesh, dim_names=["x", "y"])
|
|
device_mesh1 = DeviceMesh("another_mesh", [0, 1, 2, 3])
|
|
self.assertEqual(device_mesh.name, "my_device_mesh")
|
|
self.assertEqual(device_mesh.shape, [2, 3])
|
|
self.assertEqual(device_mesh.device_ids, [0, 1, 2, 3, 4, 5])
|
|
self.assertEqual(device_mesh.dim_names, ["x", "y"])
|
|
self.assertEqual(device_mesh.device_type, "UNKNOWN")
|
|
self.assertEqual(device_mesh.size, 6)
|
|
self.assertEqual(device_mesh.ndim, 2)
|
|
self.assertEqual(device_mesh.dim_size(0), 2)
|
|
self.assertEqual(device_mesh.dim_size(-1), 3)
|
|
self.assertEqual(device_mesh.dim_size("x"), 2)
|
|
self.assertEqual(device_mesh.dim_size("y"), 3)
|
|
self.assertEqual(device_mesh.empty(), False)
|
|
self.assertEqual(device_mesh.contains(0), True)
|
|
self.assertEqual(device_mesh.contains(6), False)
|
|
|
|
dev0 = core.Device(global_id=0, local_id=0, machine_id=0, type="GPU")
|
|
dev1 = core.Device(global_id=1, local_id=1, machine_id=0, type="GPU")
|
|
dev2 = core.Device(global_id=2, local_id=2, machine_id=0, type="GPU")
|
|
dev3 = core.Device(global_id=3, local_id=0, machine_id=1, type="GPU")
|
|
dev4 = core.Device(global_id=4, local_id=1, machine_id=1, type="GPU")
|
|
dev5 = core.Device(global_id=5, local_id=2, machine_id=1, type="GPU")
|
|
device_mesh.add_device(dev0)
|
|
device_mesh.add_device(dev1)
|
|
device_mesh.add_device(dev2)
|
|
device_mesh.add_device(dev3)
|
|
device_mesh.add_device(dev4)
|
|
device_mesh.add_device(dev5)
|
|
self.assertEqual(device_mesh.device(0), dev0)
|
|
self.assertEqual(device_mesh.device(1), dev1)
|
|
self.assertEqual(device_mesh.device(2), dev2)
|
|
self.assertEqual(device_mesh.device(3), dev3)
|
|
self.assertEqual(device_mesh.device(4), dev4)
|
|
self.assertEqual(device_mesh.device(5), dev5)
|
|
|
|
link0 = core.Link(source_id=0, target_id=1, type="NVL")
|
|
link1 = core.Link(source_id=1, target_id=0, type="NVL")
|
|
link2 = core.Link(source_id=3, target_id=4, type="NVL")
|
|
link3 = core.Link(source_id=4, target_id=5, type="NVL")
|
|
device_mesh.add_link(link0)
|
|
device_mesh.add_link(link1)
|
|
device_mesh.add_link(link2)
|
|
device_mesh.add_link(link3)
|
|
self.assertEqual(device_mesh.link(0, 1), link0)
|
|
self.assertEqual(device_mesh.link(1, 0), link1)
|
|
self.assertEqual(device_mesh.link(3, 4), link2)
|
|
self.assertEqual(device_mesh.link(4, 5), link3)
|
|
|
|
self.assertEqual(device_mesh.machine(0).id, 0)
|
|
self.assertEqual(device_mesh.machine(0).contains(3), False)
|
|
self.assertEqual(device_mesh.machine(0).device(2), dev2)
|
|
self.assertEqual(device_mesh.machine(1).link(3, 4), link2)
|
|
self.assertEqual(
|
|
device_mesh.machine(0).devices, device_mesh.machine(0).devices
|
|
)
|
|
self.assertEqual(
|
|
device_mesh.machine(0).links, device_mesh.machine(0).links
|
|
)
|
|
|
|
self.assertEqual(device_mesh.device_type, "GPU")
|
|
self.assertEqual(device_mesh.devices, device_mesh.devices)
|
|
self.assertEqual(device_mesh.links, device_mesh.links)
|
|
self.assertEqual(device_mesh.machines, device_mesh.machines)
|
|
self.assertEqual(device_mesh, device_mesh)
|
|
self.assertNotEqual(device_mesh, device_mesh1)
|
|
self.assertEqual(str(device_mesh), str(device_mesh))
|
|
|
|
def test_device(self):
|
|
device = core.Device(global_id=0, local_id=1, machine_id=2, type="GPU")
|
|
device.capability.sflops = 100
|
|
device.capability.dflops = 200
|
|
device.capability.memory = 32
|
|
device.capability.rate = 2
|
|
self.assertEqual(device.global_id, 0)
|
|
self.assertEqual(device.local_id, 1)
|
|
self.assertEqual(device.machine_id, 2)
|
|
self.assertEqual(device.type, "GPU")
|
|
self.assertAlmostEqual(device.capability.sflops, 100)
|
|
self.assertAlmostEqual(device.capability.dflops, 200)
|
|
self.assertAlmostEqual(device.capability.memory, 32)
|
|
self.assertAlmostEqual(device.capability.rate, 2)
|
|
self.assertEqual(device, device)
|
|
self.assertEqual(str(device), str(device))
|
|
|
|
def test_link(self):
|
|
link = core.Link(source_id=0, target_id=1, type="NVL")
|
|
link.capability.bandwidth = 100
|
|
link.capability.latency = 1
|
|
self.assertEqual(link.source_id, 0)
|
|
self.assertEqual(link.target_id, 1)
|
|
self.assertEqual(link.type, "NVL")
|
|
self.assertAlmostEqual(link.capability.bandwidth, 100)
|
|
self.assertAlmostEqual(link.capability.latency, 1)
|
|
self.assertEqual(link, link)
|
|
self.assertEqual(str(link), str(link))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|