用local_map替换LocalLayer

Author: zty-king

test/auto_parallel/hybrid_strategy/single_llama_model.py

@@ -256,35 +256,19 @@ def forward(self, prediction_scores, masked_lm_labels):
                 prediction_scores.astype("float32"),
                 masked_lm_labels.unsqueeze(2),
             )
-        # XPU dose not support allgather mask with bool dtype, so we use LocalLayer here.
         if paddle.device.is_compiled_with_xpu():
 
-            class LocalLossLayer(paddle.distributed.LocalLayer):
-                def __init__(self, out_dist_attrs, grad_dist_attrs):
-                    super().__init__(out_dist_attrs, grad_dist_attrs)
-
-                def forward(self, x, mask):
-                    masked_lm_loss = paddle.masked_select(x, mask).astype(
-                        "float32"
-                    )
-                    loss = paddle.mean(masked_lm_loss).unsqueeze(0)
-                    return loss.unsqueeze(0)
-
-            out_dist_attrs = [
-                (
-                    masked_lm_loss.process_mesh,
-                    [dist.Shard(0), dist.Replicate()],
-                ),
-            ]
-            grad_dist_attrs = [
-                (
-                    masked_lm_loss.process_mesh,
-                    [dist.Shard(0), dist.Replicate()],
-                ),
-                None,
-            ]
-            loss_func = LocalLossLayer(out_dist_attrs, grad_dist_attrs)
+            def LocalLoss(x, mask):
+                masked_lm_loss = paddle.masked_select(x, mask).astype("float32")
+                loss = paddle.mean(masked_lm_loss).unsqueeze(0)
+                return loss.unsqueeze(0)
 
+            loss_func = dist.local_map(
+                LocalLoss,
+                [[dist.Shard(0), dist.Replicate()]],
+                [[dist.Shard(0), dist.Replicate()], None],
+                masked_lm_loss.process_mesh,
+            )
             loss = loss_func(masked_lm_loss, masked_lm_loss > 0)
             loss = loss.mean()
             return loss

test/auto_parallel/local_view_compute.py

@@ -12,15 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import hashlib
 import random
 
 import numpy as np
 
 import paddle
 import paddle.distributed as dist
 from paddle.distributed import ProcessMesh, fleet, get_rank, shard_dataloader
-from paddle.distributed.auto_parallel.local_layer import LocalLayer
 from paddle.io import BatchSampler, DataLoader, DistributedBatchSampler
 
 base_lr = 0.01  # Learning rate
@@ -30,6 +28,7 @@
 batch_num = 100  # Number of batches per epoch
 batch_size = 32  # Batch size for training
 class_dim = 10
+global_local_loss_list = []
 
 
 class RandomDataset(paddle.io.Dataset):
@@ -65,7 +64,8 @@ def forward(self, x):
         return x
 
 
-def masked_lm_loss_func(pred, label):
+def masked_lm_loss_func(pred, label, global_local_loss_list_item=None):
+    """自定义损失函数，基于rank进行掩码"""
     lossmask = paddle.zeros_like(label).astype('float32')
     if dist.get_rank() == 0:
         lossmask[:3] = 1
@@ -74,35 +74,23 @@ def masked_lm_loss_func(pred, label):
 
     pred_sub = pred[:, 0:1]  # shape [B,1]
     label_float = paddle.cast(label, 'float32')  # shape [B,1]
-
     raw_loss = paddle.abs(pred_sub - label_float)
-
     lossmask_ = lossmask.reshape([-1]).cast('float32')
     raw_loss_flat = raw_loss.reshape([-1]).cast('float32')
 
     masked_lm_loss_sum = paddle.sum(raw_loss_flat * lossmask_)
     valid_count = paddle.sum(lossmask_)
 
     loss = masked_lm_loss_sum / (valid_count + 1e-8)
+    if global_local_loss_list_item is not None:
+        np.testing.assert_allclose(
+            global_local_loss_list_item,
+            loss,
+            rtol=1e-8,
+        )
     return loss
 
 
-class LocalViewMaskLoss(LocalLayer):
-    def __init__(self, out_dist_attrs, grad_dist_attrs):
-        super().__init__(out_dist_attrs, grad_dist_attrs)
-        self.local_loss = None
-
-    def forward(self, pred, label):
-        loss = masked_lm_loss_func(pred, label)
-        self.local_loss = loss
-        return loss
-
-
-def get_md5(tensor):
-    tensor_numpy = tensor.cpu().numpy()
-    return hashlib.md5(tensor_numpy.tobytes()).hexdigest()
-
-
 class TestLocalViewCompute:
     def __init__(self):
         self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
@@ -125,23 +113,8 @@ def create_dataset(self):
     def run_test_cases(self):
         self.set_random_seed()
         dataset = self.create_dataset()
-        dy_hand_loss_list = self.run_dy_hand(dataset)
-        self.set_random_seed()
-        dataset = self.create_dataset()
-        dy_semi_auto_local_loss_list = self.run_dy_semi_auto(dataset)
-        self.set_random_seed()
-        dy2s_semi_auto_local_loss_list = self.run_dy2s_semi_auto(dataset)
-
-        np.testing.assert_allclose(
-            dy_hand_loss_list[-1], dy_semi_auto_local_loss_list[-1], rtol=1e-8
-        )
-        np.testing.assert_allclose(
-            dy_semi_auto_local_loss_list[-1],
-            dy2s_semi_auto_local_loss_list[-1],
-            rtol=1e-8,
-        )
 
-    def run_dy_hand(self, dataset):
+        # run_dy_hand_get_local_loss
         dist_strategy = fleet.DistributedStrategy()
         dist_strategy.hybrid_configs = {
             "dp_degree": 2,
@@ -182,21 +155,14 @@ def run_dy_hand(self, dataset):
             img, label = data
 
             out = model(img)
-            lossmask = paddle.zeros_like(label).astype('float32')
-            if dist.get_rank() == 0:
-                lossmask[:3] = 1
-            else:
-                lossmask[4:9] = 1
 
             avg_loss = masked_lm_loss_func(out, label)
             avg_loss.backward()
             optimizer.step()
             model.clear_gradients()
+            global_local_loss_list.append(avg_loss.numpy())
 
-            loss_list.append(avg_loss.numpy())
-        return loss_list
-
-    def run_dy_semi_auto(self, dataset):
+        # run_dy_semi_auto
         world_process_mesh = ProcessMesh([0, 1], dim_names=["dp"])
         model = SimpleNet(
             input_size=256, inner_size=102400, output_size=class_dim
@@ -219,73 +185,30 @@ def run_dy_semi_auto(self, dataset):
         )
 
         model.train()
-        out_process_mesh = ProcessMesh([0, 1], dim_names=["dp"])
+        process_mesh = ProcessMesh([0, 1], dim_names=["dp"])
         out_placements = [dist.Partial(dist.ReduceType.kRedAvg)]
 
-        local_loss_list = []
-
         for batch_id, data in enumerate(dist_dataloader()):
             if batch_id > 10:
                 break
 
             img, label = data
 
             out = model(img)
-            loss_func = LocalViewMaskLoss(
-                out_dist_attrs=[(out_process_mesh, out_placements)],
-                grad_dist_attrs=[None, None],
+            loss_func = dist.local_map(
+                masked_lm_loss_func,
+                out_placements=out_placements,
+                in_placements=[None, None],
+                process_mesh=process_mesh,
+            )
+            avg_loss = loss_func(
+                out,
+                label,
+                global_local_loss_list_item=global_local_loss_list[batch_id],
             )
-            avg_loss = loss_func(out, label)
             avg_loss.backward()
-            local_loss_list.append(loss_func.local_loss)
             optimizer.step()
             model.clear_gradients()
-        return local_loss_list
-
-    def run_dy2s_semi_auto(self, dataset):
-        world_process_mesh = ProcessMesh([0, 1], dim_names=["dp"])
-        model = SimpleNet(
-            input_size=256, inner_size=102400, output_size=class_dim
-        )
-        optimizer = paddle.optimizer.AdamW(
-            learning_rate=base_lr,
-            weight_decay=l2_decay,
-            parameters=model.parameters(),
-        )
-
-        sampler = BatchSampler(
-            dataset, batch_size=batch_size, shuffle=False, drop_last=True
-        )
-        train_loader = DataLoader(
-            dataset, batch_sampler=sampler, num_workers=1, shuffle=False
-        )
-
-        dist_dataloader = shard_dataloader(
-            dataloader=train_loader, meshes=world_process_mesh, shard_dims="dp"
-        )
-
-        process_mesh = ProcessMesh([0, 1], dim_names=["dp"])
-        out_placements = [dist.Partial(dist.ReduceType.kRedAvg)]
-        in_grad_placements = [dist.Shard(0)]
-        loss_func = LocalViewMaskLoss(
-            out_dist_attrs=[(process_mesh, out_placements)],
-            grad_dist_attrs=[(process_mesh, in_grad_placements), None],
-        )
-        dist_model = dist.to_static(
-            model, dist_dataloader, loss_func, optimizer
-        )
-        dist_model.train()
-
-        local_loss_list = []
-        for batch_id, data in enumerate(dist_dataloader()):
-            if batch_id > 10:
-                break
-
-            img, label = data
-            loss = dist_model(img, label)
-            local_loss_list.append(loss)
-
-        return local_loss_list
 
 
 if __name__ == '__main__':

test/auto_parallel/pir/vpp_pass_unittest_local_view_pir.py

@@ -53,15 +53,11 @@ def is_optimize_op(op):
     return False
 
 
-class CustomLayer(dist.LocalLayer):
-    def __init__(self, out_dist_attrs, grad_dist_attrs):
-        super().__init__(out_dist_attrs, grad_dist_attrs)
-
-    def forward(self, input):
-        input += 0.1
-        input -= 0.3
-        input *= 0.5
-        return input
+def customFunction(input):
+    input += 0.1
+    input -= 0.3
+    input *= 0.5
+    return input
 
 
 class MyLinear(nn.Layer):
@@ -95,8 +91,8 @@ def __init__(
             [dist.Replicate()],
             stop_gradient=False,
         )
-        self.custom_local_layer = CustomLayer(
-            [(mesh, [dist.Replicate()])], [(mesh, [dist.Replicate()])]
+        self.custom_local_layer = dist.local_map(
+            customFunction, [[dist.Replicate()]], [[dist.Replicate()]], mesh
         )
 
     def forward(self, input):