[AutoParallel][PIR] support using custom spec for dist.to_static() (#69183)

AndSonder · web-flow · commit 95341a83e5ac · 2024-11-12T17:03:45.000+08:00
* support custom input spec

* change name

* change note

* Update custom_spec_pir.py
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -94,6 +94,9 @@
     from paddle.amp import GradScaler
     from paddle.base.framework import Program
     from paddle.distributed import Placement
+    from paddle.distributed.auto_parallel.static.dist_input_spec import (
+        DistributedInputSpec,
+    )
     from paddle.io import DataLoader
     from paddle.metric import Metric
     from paddle.nn import Layer
@@ -2166,6 +2169,13 @@ class DistModel:
         strategy(paddle.distributed.Strategy|None, optional): Configs for
             parallel strategies and optimization settings (e.g. sharding,
             pipeline parallelism). Default: None.
+        input_spec(list[list[paddle.distributed.DistributedInputSpec]]|None, optional):
+            The custom input specs specify the shape, dtype, and name information
+            of model inputs and labels. If it is not None, the input specs and
+            label specs will be inferred from the custom input specs. The custom
+            input specs should be a list containing two sublists: the first
+            sublist represents theinput specs, and the second sublist represents
+            the label specs. Default: None.
     """
 
     def __init__(
@@ -2176,6 +2186,7 @@ def __init__(
         optimizer: Optimizer | None = None,
         strategy: Strategy | None = None,
         metrics: list[Metric] | None = None,
+        input_spec: list[list[DistributedInputSpec]] | None = None,
     ) -> None:
         self._feed_name_list = []
         self._inner_strategy = self.__convert_strategy(strategy)
@@ -2206,7 +2217,10 @@ def __init__(
         self._feed_name_list = {}
 
         # convert dygraph model to static model
-        if isinstance(loader, ShardDataloader):
+        if input_spec is not None:
+            self._engine._inputs_spec = input_spec[0]
+            self._engine._labels_spec = input_spec[1]
+        elif isinstance(loader, ShardDataloader):
             (
                 self._engine._inputs_spec,
                 self._engine._labels_spec,
@@ -2734,6 +2748,7 @@ def to_static(
     loss: Layer | Callable[..., Any] | None = None,
     optimizer: Optimizer | None = None,
     strategy: Strategy | None = None,
+    input_spec: list[list[DistributedInputSpec]] | None = None,
 ) -> DistModel:
     """
     Converts the ``layer`` with distributed tensor (constructed from
@@ -2755,6 +2770,13 @@ def to_static(
         strategy(paddle.distributed.Strategy|None, optional): Configs for
             parallel strategies and optimization settings (e.g. sharding,
             pipeline parallelism). Default: None.
+        input_spec(list[list[paddle.distributed.DistributedInputSpec]]|None, optional):
+            The custom input specs specify the shape, dtype, and name information
+            of model inputs and labels. If it is not None, the input specs and
+            label specs will be inferred from the custom input specs. The custom
+            input specs should be a list containing two sublists: the first
+            sublist represents theinput specs, and the second sublist represents
+            the label specs. Default: None.
 
     Returns:
         DistModel: A ``DistModel`` instance converted the input ``layer``.
@@ -2884,7 +2906,10 @@ def to_static(
                 raise NotImplementedError(
                     "Only sharding stage 1, 2 and 3 can to_static for now. User-defined shard_fn will be supported later."
                 )
-    dist_model = DistModel(layer, loader, loss, optimizer, strategy)
+
+    dist_model = DistModel(
+        layer, loader, loss, optimizer, strategy, input_spec=input_spec
+    )
     return dist_model
 
 
diff --git a/test/auto_parallel/pir/CMakeLists.txt b/test/auto_parallel/pir/CMakeLists.txt
@@ -66,6 +66,8 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
     test_semi_auto_parallel_simple_net_ep ENVS FLAGS_enable_pir_api=1)
   set_tests_properties(test_semi_auto_parallel_simple_net_ep
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 60)
+  py_test_modules(test_custom_spec_pir MODULES test_custom_spec_pir ENVS
+                  FLAGS_enable_pir_api=1)
 endif()
 py_test_modules(test_pir_1f1b_plan MODULES test_pir_1f1b_plan ENVS
                 FLAGS_enable_pir_api=1)
diff --git a/test/auto_parallel/pir/custom_spec_pir.py b/test/auto_parallel/pir/custom_spec_pir.py
@@ -0,0 +1,145 @@
+# Copyright (c) 2024 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
+#
+#     https://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 copy
+import random
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.static.dist_input_spec import (
+    DistributedInputSpec,
+)
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+np.random.seed(1127)
+paddle.seed(1127)
+random.seed(1127)
+
+mesh = dist.ProcessMesh([0, 1], dim_names=["dp"])
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, hidden, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.hidden = hidden
+        self.num_samples = num_samples
+        self.mode = "A"
+
+    def __getitem__(self, index):
+        if self.mode == "A":
+            input = np.random.uniform(size=[self.seq_len, self.hidden]).astype(
+                "float32"
+            )
+        else:
+            input = np.random.normal(
+                size=[self.seq_len * 2, self.hidden]
+            ).astype("float32")
+        label = np.random.randint(0, 2, size=[128]).astype("int64")
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+class MlpModel(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.w0 = self.create_parameter(shape=[1024, 4096])
+        self.w1 = self.create_parameter(shape=[1024, 4096])
+
+    def forward(self, x):
+        y1 = paddle.matmul(x, self.w0)
+        y2 = paddle.matmul(x, self.w1)
+        z = y1 + y2
+        return z
+
+
+class TestCustomSpec:
+    def get_input_spec(self, dist_dataloader):
+        dist_dataloader._dataloader.mode = "A"
+        input1, label1 = next(dist_dataloader())
+        dist_dataloader._dataloader.mode = "B"
+        input2, label2 = next(dist_dataloader())
+
+        inputs_spec1 = [DistributedInputSpec.from_dtensor(input1, "input0")]
+        inputs_spec2 = [DistributedInputSpec.from_dtensor(input2, "input0")]
+        labels_spec = [DistributedInputSpec.from_dtensor(label1, "label0")]
+
+        return [inputs_spec1, labels_spec], [inputs_spec2, labels_spec]
+
+    def run_test(self):
+        model = MlpModel()
+        loss_func = paddle.nn.CrossEntropyLoss()
+
+        dataset = RandomDataset(128, 1024, 40)
+        sampler = BatchSampler(
+            dataset,
+            batch_size=4,
+        )
+        dataloader = DataLoader(
+            dataset,
+            batch_sampler=sampler,
+        )
+        dist_dataloader = dist.shard_dataloader(
+            dataloader=dataloader, meshes=mesh, shard_dims="dp"
+        )
+
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+        opt2 = copy.deepcopy(opt)
+
+        input_spec1, input_spec2 = self.get_input_spec(dist_dataloader)
+
+        model.w0.stop_gradient = True
+        model.w1.stop_gradient = False
+        dist_model1 = dist.to_static(
+            model, dist_dataloader, loss_func, opt, input_spec=input_spec1
+        )
+        dist_model1.train()
+
+        model.w0.stop_gradient = False
+        model.w1.stop_gradient = True
+        dist_model2 = dist.to_static(
+            model, dist_dataloader, loss_func, opt2, input_spec=input_spec2
+        )
+        dist_model2.train()
+
+        datasets_modes = ["A", "A", "B", "A", "B"]
+        for mode in datasets_modes:
+            dist_dataloader._dataloader.mode = mode
+            input, label = next(iter(dist_dataloader))
+            if mode == "A":
+                before_w0 = dist_model1.state_dict("param")['w0'].mean().numpy()
+                before_w1 = dist_model1.state_dict("param")['w1'].mean().numpy()
+                loss = dist_model1(input, label)
+                after_w0 = dist_model1.state_dict("param")['w0'].mean().numpy()
+                after_w1 = dist_model1.state_dict("param")['w1'].mean().numpy()
+                assert np.equal(before_w0, after_w0).all()
+                assert not np.equal(before_w1, after_w1).all()
+            else:
+                before_w0 = dist_model2.state_dict("param")['w0'].mean().numpy()
+                before_w1 = dist_model2.state_dict("param")['w1'].mean().numpy()
+                loss = dist_model2(input, label)
+                after_w0 = dist_model2.state_dict("param")['w0'].mean().numpy()
+                after_w1 = dist_model2.state_dict("param")['w1'].mean().numpy()
+                assert not np.equal(before_w0, after_w0).all()
+                assert np.equal(before_w1, after_w1).all()
+
+
+if __name__ == '__main__':
+    TestCustomSpec().run_test()
diff --git a/test/auto_parallel/pir/test_custom_spec_pir.py b/test/auto_parallel/pir/test_custom_spec_pir.py
@@ -0,0 +1,47 @@
+# Copyright (c) 2024 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
+#
+#     https://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
+
+import collective.test_communication_api_base as test_base
+
+
+class TestCustomSpec(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=50,
+        )
+        self._default_envs = {
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+        }
+
+    def test_custom_spec(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "custom_spec_pir.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
