test for distributionlib py import functools import os from unittest import mock import jax import numpy as np import pytest from keras import backend from keras import layers from keras import models from keras import testing from keras backend import distributionlib as backenddlib from keras distribution import distributionlib if backend backend jax due to https github comgooglejaxissues17188 we can t override the xla flag after the jax back init we have to run this at top level to let jax pick the flag value xlaflags os getenvxlaflags or don t override userspecified device count or other xla flags if xlaforcehostplatformdevicecount not in xlaflags os environxlaflags xlaflags xlaforcehostplatformdevicecount8 pytest mark skipif backend backend jax reasonbackend specific test class jaxdistributionlibtesttesting testcase def testlistdevicesself self assertequallendistributionlib listdevices 8 self assertequallendistributionlib listdevicescpu 8 self assertequallendistributionlib listdevicescpu 8 def testdeviceconversionself devices distributionlib listdevicescpu jaxdevices jax devicescpu for d jaxd in zipdevices jaxdevices convertedjaxdevice backenddlib tojaxdeviced self assertisinstanceconvertedjaxdevice jax device self assertequaljaxd convertedjaxdevice mock patch objectjax distributed initialize returnvaluenone def testinitializewithalljobaddressesself mockjaxinitialze backenddlib initialize10 0 0 1 1234 10 0 0 2 2345 2 0 mockjaxinitialze assertcalledoncewith corrdinatoraddress10 0 0 1 1234 numprocesses2 processid0 def testinitializevalidatejobandprocessself with self assertraisesregex valueerror has 2 jobs but numprocesses is 3 backenddlib initialize10 0 0 1 1234 10 0 0 2 2345 3 0 mock patch objectjax distributed initialize returnvaluenone def testinitializewithcoordinateraddressself mockjaxinitialze backenddlib initialize10 0 0 1 1234 2 0 mockjaxinitialze assertcalledoncewith corrdinatoraddress10 0 0 1 1234 numprocesses2 processid0 def testdistributetensorself jaxmesh jax sharding mesh np arrayjax devices reshape2 4 batch model inputs jax numpy arraynp random normalsize16 8 targetlayout jax sharding namedsharding jaxmesh jax sharding partitionspecbatch none functools partialjax jit staticargnamestargetlayout def testfunctioninputs targetlayout return distributionlib distributetensorinputs targetlayout result testfunctioninputs targetlayout note that the returned tensor has a different sharding implementation which is gspmdsharding but it should be equivalent as the target layout specified self asserttrueresult sharding isequivalenttotargetlayout ndim2 test without jit result distributionlib distributetensorinputs targetlayout self asserttrueresult sharding isequivalenttotargetlayout ndim2 def testdistributevariableself this test only verify the single workerprocess behavior the multiprocess test lives in g3 jaxmesh jax sharding mesh np arrayjax devices reshape2 4 batch model variable jax numpy arraynp random normalsize16 8 targetlayout jax sharding namedsharding jaxmesh jax sharding partitionspecmodel none result backenddlib distributevariablevariable targetlayout note that the returned tensor has a different sharding implementation which is gspmdsharding but it should be equivalent as the target layout specified self asserttrueresult sharding isequivalenttotargetlayout ndim2 def testdistributeinputdataself this test only verify the single workerprocess behavior the multiprocess test lives in g3 jaxmesh jax sharding mesh np arrayjax devices reshape2 4 batch model inputdata jax numpy arraynp random normalsize16 8 targetlayout jax sharding namedsharding jaxmesh jax sharding partitionspecbatch none result backenddlib distributevariableinputdata targetlayout note that the returned tensor has a different sharding implementation which is gspmdsharding but it should be equivalent as the target layout specified self asserttrueresult sharding isequivalenttotargetlayout ndim2 def testprocessesself self assertequalbackenddlib processid 0 self assertequalbackenddlib numprocesses 1 def testtojaxmeshself devices fcpu i for i in range8 shape 4 2 axisnames batch model mesh distributionlib devicemeshshape axisnames devices jaxmesh backenddlib tojaxmeshmesh self assertisinstancejaxmesh jax sharding mesh self assertequaljaxmesh devices shape shape self assertequaljaxmesh axisnames batch model def testtojaxlayoutself axes data none mesh distributionlib devicemesh 4 2 data model fcpu i for i in range8 layout distributionlib tensorlayoutaxes mesh jaxsharding backenddlib tojaxlayoutlayout jaxmesh backenddlib tojaxmeshmesh self assertequal jaxsharding jax sharding namedsharding jaxmesh jax sharding partitionspecdata none def testvalidationfordevicemeshself axes data none layout distributionlib tensorlayoutaxes devicemeshnone with self assertraisesregex valueerror cannot create sharding when device mesh is not set backenddlib tojaxlayoutlayout def testvariableassignmentreuselayoutself shape 4 2 axisnames batch model devicemesh distributionlib devicemesh shape axisnames backenddlib listdevices layoutmap distributionlib layoutmapdevicemesh layoutmap dense kernel distributionlib tensorlayout none model layoutmap dense bias distributionlib tensorlayoutmodel distribution distributionlib modelparallel devicemesh layoutmap batchdimnamebatch with distribution scope denselayer layers dense8 denselayer build16 16 self assertequal denselayer kernel value sharding spec none model self assertequaldenselayer bias value sharding spec model assign a numpy value to dense layer to mimic the model weight loading newkernel np random normalsize16 8 newbias np random normalsize8 denselayer kernel assignnewkernel denselayer bias assignnewbias make sure the loaded value still use the layout when it is initialized even outside of the distribution scope self assertequal denselayer kernel value sharding spec none model self assertequaldenselayer bias value sharding spec model def teste2edataparallelmodelself distribution distributionlib dataparallel devicesbackenddlib listdevices with distribution scope inputs layers inputshape28 28 1 y layers flatteninputs y layers denseunits200 usebiasfalse activationreluy y layers dropout0 4y y layers denseunits10 activationsoftmaxy model models modelinputsinputs outputsy make sure all the weights are properly sharded for weight in model weights self asserttrueweight value sharding isfullyreplicated inputs np random normalsize32 28 28 1 labels np random normalsize32 10 with distribution scope model compilelossmse model fitinputs labels def teste2emodelparallelmodelself shape 4 2 axisnames batch model devicemesh distributionlib devicemesh shape axisnames backenddlib listdevices layoutmap distributionlib layoutmapdevicemesh layoutmap dense kernel distributionlib tensorlayout none model layoutmap dense bias distributionlib tensorlayoutmodel distribution distributionlib modelparallel devicemesh layoutmap batchdimnamebatch with distribution scope inputs layers inputshape28 28 1 y layers flatteninputs y layers denseunits200 usebiasfalse activationreluy y layers dropout0 4y y layers denseunits10 activationsoftmaxy model models modelinputsinputs outputsy for weight in model weights if kernel in weight name self assertequalweight value sharding spec none model elif bias in weight name self assertequalweight value sharding spec model else self asserttrueweight value sharding isfullyreplicated inputs np random normalsize32 28 28 1 labels np random normalsize32 10 with distribution scope model compilelossmse model fitinputs labels def teste2emodelparallelwithoutputshardingself shape 4 2 axisnames batch model devicemesh distributionlib devicemesh shape axisnames backenddlib listdevices layoutmap distributionlib layoutmapdevicemesh layoutmap dense kernel distributionlib tensorlayout none model layoutmap dense bias distributionlib tensorlayoutmodel force the dense layer output to be batch parallel only and not sharded on model dimension layoutmap dense output batch none distribution distributionlib modelparallel devicemesh layoutmap batchdimnamebatch shardingcapture shardingcapturelayer with distribution scope inputs layers inputshape28 28 1 y layers flatteninputs y layers denseunits200 usebiasfalse activationreluy y shardingcapturey y layers dropout0 4y y layers denseunits10 activationsoftmaxy model models modelinputsinputs outputsy for weight in model weights if kernel in weight name self assertequalweight value sharding spec none model elif bias in weight name self assertequalweight value sharding spec model else self asserttrueweight value sharding isfullyreplicated inputs np random normalsize32 28 28 1 labels np random normalsize32 10 with distribution scope model compilelossmse model fitinputs labels note that the intermediatetensorlayout is only captured during the actual training and not at the model building time intermediatetensorlayout jax sharding namedsharding backenddlib tojaxmeshdistribution devicemesh jax sharding partitionspecbatch none self asserttrue shardingcapture capturedinputsharding isequivalentto intermediatetensorlayout ndim2 class shardingcapturelayerlayers layer def initself kwargs super initkwargs self capturedinputsharding none self supportsmasking true def callself inputs jax debug inspectarraysharding inputs callbacklambda x self captureinputshardingx return inputs def captureinputshardingself sharding self capturedinputsharding sharding due to https github com google jax issues 17188 we can t override the xla flag after the jax back init we have to run this at top level to let jax pick the flag value don t override user specified device count or other xla flags note that the returned tensor has a different sharding implementation which is gspmdsharding but it should be equivalent as the target layout specified test without jit this test only verify the single worker process behavior the multi process test lives in g3 note that the returned tensor has a different sharding implementation which is gspmdsharding but it should be equivalent as the target layout specified this test only verify the single worker process behavior the multi process test lives in g3 note that the returned tensor has a different sharding implementation which is gspmdsharding but it should be equivalent as the target layout specified assign a numpy value to dense layer to mimic the model weight loading make sure the loaded value still use the layout when it is initialized even outside of the distribution scope make sure all the weights are properly sharded force the dense layer output to be batch parallel only and not sharded on model dimension note that the intermediate_tensor_layout is only captured during the actual training and not at the model building time
import functools
import os
from unittest import mock

import jax
import numpy as np
import pytest

from keras import backend
from keras import layers
from keras import models
from keras import testing
from keras.backend import distribution_lib as backend_dlib
from keras.distribution import distribution_lib

if backend.backend() == "jax":
    
    
    
    xla_flags = os.getenv("XLA_FLAGS") or ""
    
    if "xla_force_host_platform_device_count" not in xla_flags:
        os.environ["XLA_FLAGS"] = (
            xla_flags + " --xla_force_host_platform_device_count=8"
        )


@pytest.mark.skipif(
    backend.backend() != "jax",
    reason="Backend specific test",
)
class JaxDistributionLibTest(testing.TestCase):
    def test_list_devices(self):
        self.assertEqual(len(distribution_lib.list_devices()), 8)
        self.assertEqual(len(distribution_lib.list_devices("cpu")), 8)
        self.assertEqual(len(distribution_lib.list_devices("cpu")), 8)

    def test_device_conversion(self):
        devices = distribution_lib.list_devices("cpu")
        jax_devices = jax.devices("cpu")

        for d, jax_d in zip(devices, jax_devices):
            converted_jax_device = backend_dlib._to_jax_device(d)
            self.assertIsInstance(converted_jax_device, jax.Device)
            self.assertEqual(jax_d, converted_jax_device)

    @mock.patch.object(jax.distributed, "initialize", return_value=None)
    def test_initialize_with_all_job_addresses(self, mock_jax_initialze):
        backend_dlib.initialize("10.0.0.1:1234,10.0.0.2:2345", 2, 0)
        mock_jax_initialze.assert_called_once_with(
            corrdinator_address="10.0.0.1:1234", num_processes=2, process_id=0
        )

    def test_initialize_validate_job_and_process(self):
        with self.assertRaisesRegex(
            ValueError, "has 2 jobs, but num_processes is 3"
        ):
            backend_dlib.initialize("10.0.0.1:1234,10.0.0.2:2345", 3, 0)

    @mock.patch.object(jax.distributed, "initialize", return_value=None)
    def test_initialize_with_coordinater_address(self, mock_jax_initialze):
        backend_dlib.initialize("10.0.0.1:1234", 2, 0)
        mock_jax_initialze.assert_called_once_with(
            corrdinator_address="10.0.0.1:1234", num_processes=2, process_id=0
        )

    def test_distribute_tensor(self):
        jax_mesh = jax.sharding.Mesh(
            np.array(jax.devices()).reshape(2, 4), ("batch", "model")
        )

        inputs = jax.numpy.array(np.random.normal(size=(16, 8)))
        target_layout = jax.sharding.NamedSharding(
            jax_mesh, jax.sharding.PartitionSpec("batch", None)
        )

        @functools.partial(jax.jit, static_argnames="target_layout")
        def test_function(inputs, target_layout):
            return distribution_lib.distribute_tensor(inputs, target_layout)

        result = test_function(inputs, target_layout)
        
        
        
        self.assertTrue(result.sharding.is_equivalent_to(target_layout, ndim=2))

        
        result = distribution_lib.distribute_tensor(inputs, target_layout)
        self.assertTrue(result.sharding.is_equivalent_to(target_layout, ndim=2))

    def test_distribute_variable(self):
        
        
        jax_mesh = jax.sharding.Mesh(
            np.array(jax.devices()).reshape(2, 4), ("batch", "model")
        )

        variable = jax.numpy.array(np.random.normal(size=(16, 8)))
        target_layout = jax.sharding.NamedSharding(
            jax_mesh, jax.sharding.PartitionSpec("model", None)
        )

        result = backend_dlib.distribute_variable(variable, target_layout)
        
        
        
        self.assertTrue(result.sharding.is_equivalent_to(target_layout, ndim=2))

    def test_distribute_input_data(self):
        
        
        jax_mesh = jax.sharding.Mesh(
            np.array(jax.devices()).reshape(2, 4), ("batch", "model")
        )

        input_data = jax.numpy.array(np.random.normal(size=(16, 8)))
        target_layout = jax.sharding.NamedSharding(
            jax_mesh, jax.sharding.PartitionSpec("batch", None)
        )

        result = backend_dlib.distribute_variable(input_data, target_layout)
        
        
        
        self.assertTrue(result.sharding.is_equivalent_to(target_layout, ndim=2))

    def test_processes(self):
        self.assertEqual(backend_dlib.process_id(), 0)
        self.assertEqual(backend_dlib.num_processes(), 1)

    def test_to_jax_mesh(self):
        devices = [f"cpu:{i}" for i in range(8)]
        shape = (4, 2)
        axis_names = ["batch", "model"]

        mesh = distribution_lib.DeviceMesh(shape, axis_names, devices)
        jax_mesh = backend_dlib._to_jax_mesh(mesh)

        self.assertIsInstance(jax_mesh, jax.sharding.Mesh)
        self.assertEqual(jax_mesh.devices.shape, shape)
        self.assertEqual(jax_mesh.axis_names, ("batch", "model"))

    def test_to_jax_layout(self):
        axes = ["data", None]
        mesh = distribution_lib.DeviceMesh(
            (4, 2), ["data", "model"], [f"cpu:{i}" for i in range(8)]
        )
        layout = distribution_lib.TensorLayout(axes, mesh)
        jax_sharding = backend_dlib._to_jax_layout(layout)
        jax_mesh = backend_dlib._to_jax_mesh(mesh)
        self.assertEqual(
            jax_sharding,
            jax.sharding.NamedSharding(
                jax_mesh, jax.sharding.PartitionSpec("data", None)
            ),
        )

    def test_validation_for_device_mesh(self):
        axes = ["data", None]
        layout = distribution_lib.TensorLayout(axes, device_mesh=None)

        with self.assertRaisesRegex(
            ValueError, "Cannot create sharding when device mesh is not set"
        ):
            backend_dlib._to_jax_layout(layout)

    def test_variable_assignment_reuse_layout(self):
        shape = (4, 2)
        axis_names = ["batch", "model"]
        device_mesh = distribution_lib.DeviceMesh(
            shape, axis_names, backend_dlib.list_devices()
        )
        layout_map = distribution_lib.LayoutMap(device_mesh)
        layout_map[".*dense.*kernel"] = distribution_lib.TensorLayout(
            [None, "model"]
        )
        layout_map[".*dense.*bias"] = distribution_lib.TensorLayout(["model"])

        distribution = distribution_lib.ModelParallel(
            device_mesh, layout_map, batch_dim_name="batch"
        )

        with distribution.scope():
            dense_layer = layers.Dense(8)
            dense_layer.build((16, 16))

        self.assertEqual(
            dense_layer.kernel._value.sharding.spec, (None, "model")
        )
        self.assertEqual(dense_layer.bias._value.sharding.spec, ("model",))

        
        new_kernel = np.random.normal(size=(16, 8))
        new_bias = np.random.normal(size=(8))
        dense_layer.kernel.assign(new_kernel)
        dense_layer.bias.assign(new_bias)

        
        
        self.assertEqual(
            dense_layer.kernel._value.sharding.spec, (None, "model")
        )
        self.assertEqual(dense_layer.bias._value.sharding.spec, ("model",))

    def test_e2e_data_parallel_model(self):
        distribution = distribution_lib.DataParallel(
            devices=backend_dlib.list_devices()
        )

        with distribution.scope():
            inputs = layers.Input(shape=[28, 28, 1])
            y = layers.Flatten()(inputs)
            y = layers.Dense(units=200, use_bias=False, activation="relu")(y)
            y = layers.Dropout(0.4)(y)
            y = layers.Dense(units=10, activation="softmax")(y)
            model = models.Model(inputs=inputs, outputs=y)

        
        for weight in model.weights:
            self.assertTrue(weight._value.sharding.is_fully_replicated)

        inputs = np.random.normal(size=(32, 28, 28, 1))
        labels = np.random.normal(size=(32, 10))

        with distribution.scope():
            model.compile(loss="mse")
            model.fit(inputs, labels)

    def test_e2e_model_parallel_model(self):
        shape = (4, 2)
        axis_names = ["batch", "model"]
        device_mesh = distribution_lib.DeviceMesh(
            shape, axis_names, backend_dlib.list_devices()
        )

        layout_map = distribution_lib.LayoutMap(device_mesh)
        layout_map[".*dense.*kernel"] = distribution_lib.TensorLayout(
            [None, "model"]
        )
        layout_map[".*dense.*bias"] = distribution_lib.TensorLayout(["model"])

        distribution = distribution_lib.ModelParallel(
            device_mesh, layout_map, batch_dim_name="batch"
        )
        with distribution.scope():
            inputs = layers.Input(shape=[28, 28, 1])
            y = layers.Flatten()(inputs)
            y = layers.Dense(units=200, use_bias=False, activation="relu")(y)
            y = layers.Dropout(0.4)(y)
            y = layers.Dense(units=10, activation="softmax")(y)
            model = models.Model(inputs=inputs, outputs=y)

        for weight in model.weights:
            if "kernel" in weight.name:
                self.assertEqual(weight._value.sharding.spec, (None, "model"))
            elif "bias" in weight.name:
                self.assertEqual(weight._value.sharding.spec, ("model",))
            else:
                self.assertTrue(weight._value.sharding.is_fully_replicated)

        inputs = np.random.normal(size=(32, 28, 28, 1))
        labels = np.random.normal(size=(32, 10))

        with distribution.scope():
            model.compile(loss="mse")
            model.fit(inputs, labels)

    def test_e2e_model_parallel_with_output_sharding(self):
        shape = (4, 2)
        axis_names = ["batch", "model"]
        device_mesh = distribution_lib.DeviceMesh(
            shape, axis_names, backend_dlib.list_devices()
        )

        layout_map = distribution_lib.LayoutMap(device_mesh)
        layout_map[".*dense.*kernel"] = distribution_lib.TensorLayout(
            [None, "model"]
        )
        layout_map[".*dense.*bias"] = distribution_lib.TensorLayout(["model"])
        
        
        layout_map[".*dense.*output"] = ("batch", None)

        distribution = distribution_lib.ModelParallel(
            device_mesh, layout_map, batch_dim_name="batch"
        )
        sharding_capture = ShardingCaptureLayer()
        with distribution.scope():
            inputs = layers.Input(shape=[28, 28, 1])
            y = layers.Flatten()(inputs)
            y = layers.Dense(units=200, use_bias=False, activation="relu")(y)
            y = sharding_capture(y)
            y = layers.Dropout(0.4)(y)
            y = layers.Dense(units=10, activation="softmax")(y)
            model = models.Model(inputs=inputs, outputs=y)

        for weight in model.weights:
            if "kernel" in weight.name:
                self.assertEqual(weight._value.sharding.spec, (None, "model"))
            elif "bias" in weight.name:
                self.assertEqual(weight._value.sharding.spec, ("model",))
            else:
                self.assertTrue(weight._value.sharding.is_fully_replicated)

        inputs = np.random.normal(size=(32, 28, 28, 1))
        labels = np.random.normal(size=(32, 10))

        with distribution.scope():
            model.compile(loss="mse")
            model.fit(inputs, labels)

        
        
        intermediate_tensor_layout = jax.sharding.NamedSharding(
            backend_dlib._to_jax_mesh(distribution.device_mesh),
            jax.sharding.PartitionSpec("batch", None),
        )
        self.assertTrue(
            sharding_capture.captured_input_sharding.is_equivalent_to(
                intermediate_tensor_layout, ndim=2
            )
        )


class ShardingCaptureLayer(layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.captured_input_sharding = None
        self.supports_masking = True

    def call(self, inputs):
        jax.debug.inspect_array_sharding(
            inputs, callback=lambda x: self.capture_input_sharding(x)
        )
        return inputs

    def capture_input_sharding(self, sharding):
        self.captured_input_sharding = sharding