Update semi_auto_parallel_sharding_stage_1.py

Author: Xing-lil

test/auto_parallel/semi_auto_parallel_sharding_stage_1.py

@@ -234,6 +234,34 @@ def run_sharding_test(enable_tensor_fusion):
         loss_enable = run_sharding_test(enable_tensor_fusion=True)
         self.check_tensor_eq(loss_disable, loss_enable)
 
+    def test_pure_sharding_multi_mesh_stage_1_with_tensor_fusion_with_chip(
+        self,
+    ):
+        dist.init_parallel_env()
+        os.environ['FLAGS_enable_inplace_master_grad'] = '1'
+        os.environ['FLAGS_enable_tensor_fusion'] = '1'
+        paddle.distributed.auto_parallel.set_mesh(self._multi_dim_mesh)
+        paddle.seed(self._seed)
+        model = paddle.nn.Linear(10, 10)
+        batch = paddle.rand(shape=[10, 10])
+        batch = dist.shard_tensor(batch, self._mesh, [dist.Shard(0)])
+        clip = paddle.nn.ClipGradByGlobalNorm(1.0)
+        opt = paddle.optimizer.AdamW(
+            parameters=model.parameters(), grad_clip=clip
+        )
+        opt = dist.shard_optimizer(
+            opt, dist.ShardingStage1(sharding_mesh_dim="dp")
+        )
+        model, opt = paddle.amp.decorate(
+            model, optimizers=opt, level='O2', master_grad=True
+        )
+        for _ in range(5):
+            with paddle.amp.auto_cast(level='O2'):
+                loss = model(batch)
+                loss.backward()
+                opt.step()
+                opt.clear_grad()
+
     def run_test_case(self):
         if self._backend == "cpu":
             paddle.set_device("cpu")
@@ -251,6 +279,7 @@ def run_test_case(self):
         self.test_sharding_stage_1_overlap_to_static()
         self.test_pure_sharding_multi_mesh_stage_1_with_inplace_master_grad()
         self.test_pure_sharding_multi_mesh_stage_1_with_tensor_fusion()
+        self.test_pure_sharding_multi_mesh_stage_1_with_tensor_fusion_with_chip()
 
 
 if __name__ == '__main__':