[AutoParallel] optimize sharding stage1 tensor fusion save&load strategy (#70309)

AndSonder · web-flow · commit 8d97458cbbab · 2024-12-26T20:19:35.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -2330,14 +2330,21 @@ def __init__(
             )
             dist.fleet.init(is_collective=True)
 
-        if int(os.environ.get('FLAGS_enable_sharding_stage1_tensor_fusion', 0)):
-            if isinstance(optimizer, _ShardOptimizer) and use_pir_api():
-                shard_fn = optimizer._shard_fn
-                inner_opt = optimizer._inner_opt
-                if isinstance(optimizer._shard_fn, ShardingStage1):
-                    optimizer = ShardingOptimizerStage1(
-                        inner_opt, shard_fn, self._inner_strategy
-                    )
+        if (
+            strategy
+            and strategy.sharding.enable_stage1_tensor_fusion
+            and isinstance(optimizer, _ShardOptimizer)
+            and use_pir_api()
+        ):
+            assert isinstance(optimizer._shard_fn, ShardingStage1), (
+                "The shard_fn should be ShardingStage1 "
+                "when stage1 tensor fusion is enabled."
+            )
+            shard_fn = optimizer._shard_fn
+            inner_opt = optimizer._inner_opt
+            optimizer = ShardingOptimizerStage1(
+                inner_opt, shard_fn, self._inner_strategy
+            )
 
         self._engine = Engine(
             layer, loss, optimizer, metrics, strategy=self._inner_strategy
@@ -2724,17 +2731,35 @@ def state_dict(
                                     ] = dist_tensor
                                 dist_state_dict.pop(param)
 
-        # when tensor-fusion is enabled, the optimizer parameters are unbalanced
-        # in their sharding. We need to process the optimizer parameters to make
-        # them evenly balanced
-        if self._engine._optimizer is not None and load_sharded_model:
+        # When tensor fusion is enabled, optimizer parameters can become unbalanced in
+        # sharding. We need to either balance them or rename unbalanced parameters for each
+        # rank and directly load them.
+        enable_stage1_tensor_fusion = (
+            self._inner_strategy.sharding.enable_stage1_tensor_fusion
+            if self._inner_strategy
+            else False
+        )
+        if (
+            self._engine._optimizer is not None
+            and load_sharded_model
+            and enable_stage1_tensor_fusion
+        ):
             optimizer = self._engine._optimizer
             if isinstance(
                 optimizer,
                 paddle.static.amp.decorator.OptimizerWithMixedPrecision,
             ):
                 optimizer = optimizer._optimizer
-            if isinstance(optimizer, ShardingOptimizerStage1):
+
+            assert isinstance(
+                optimizer, ShardingOptimizerStage1
+            ), "The optimizer should be ShardingOptimizerStage1 when stage1 tensor fusion is enabled."
+
+            if self._inner_strategy.sharding.save_unbalanced_param:
+                optimizer.convert_state_dict_with_rank_unique_name(
+                    dist_state_dict
+                )
+            else:
                 optimizer.convert_state_dict_without_tensor_fusion_param(
                     dist_state_dict
                 )
@@ -2815,22 +2840,36 @@ def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
         # For sharding with tensor-fusion, we need to convert the state_dict
         # to include tensor-fusion parameters before calling set_state_dict,
         # as stored parameters are processed as if tensor-fusion is not applied
-        if self._engine._optimizer is not None:
+        # or we can choose to load the unblanced parameters directly.
+        enable_stage1_tensor_fusion = (
+            self._inner_strategy.sharding.enable_stage1_tensor_fusion
+            if self._inner_strategy
+            else False
+        )
+        if self._engine._optimizer is not None and enable_stage1_tensor_fusion:
             optimizer = self._engine._optimizer
             if isinstance(
                 optimizer,
                 paddle.static.amp.decorator.OptimizerWithMixedPrecision,
             ):
                 optimizer = optimizer._optimizer
-            if isinstance(optimizer, ShardingOptimizerStage1):
+
+            assert isinstance(
+                optimizer, ShardingOptimizerStage1
+            ), "The optimizer should be ShardingOptimizerStage1 when stage1 tensor fusion is enabled."
+
+            if self._inner_strategy.sharding.save_unbalanced_param:
+                optimizer.convert_state_dict_with_origin_name(state_dict)
+            else:
                 optimizer.convert_state_dict_with_tensor_fusion_param(
                     state_dict
                 )
-                # When using the tensor-fusion strategy, model parameters are shared with
-                # slice@ parameters. When setting the state_dict, we must copy the tensor
-                # instead of changing the handle directly, as this could cause errors in
-                # the slice@ parameters and increase memory usage.
-                copy_tensor = True
+
+            # When using the tensor-fusion strategy, model parameters are shared with
+            # slice@ parameters. When setting the state_dict, we must copy the tensor
+            # instead of changing the handle directly, as this could cause errors in
+            # the slice@ parameters and increase memory usage.
+            copy_tensor = True
 
         for k, v in state_dict.items():
             assert v.is_dist(), f"key {k} value:{v} is not a dist tensor."
@@ -2906,9 +2945,14 @@ def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
                                 for ori_p in ori_params_meta:
                                     local_state_dict.pop(ori_p + suffix)
 
-        dist_main_program.set_state_dict(
-            local_state_dict, paddle.static.global_scope()
-        )
+        if use_pir_api():
+            dist_main_program.set_state_dict(
+                local_state_dict, paddle.static.global_scope(), copy_tensor
+            )
+        else:
+            dist_main_program.set_state_dict(
+                local_state_dict, paddle.static.global_scope()
+            )
 
 
 def to_static(
diff --git a/python/paddle/distributed/auto_parallel/constants.py b/python/paddle/distributed/auto_parallel/constants.py
@@ -159,6 +159,9 @@ class _AMPConfig(TypedDict, total=False):  # noqa: PYI049
 set_field_default_config(SHARDING, "tuning_range", [])
 set_field_default_config(SHARDING, "release_gradients", False)
 set_field_default_config(SHARDING, "comm_buffer_size_MB", 256)
+set_field_default_config(SHARDING, "enable_stage1_tensor_fusion", False)
+set_field_default_config(SHARDING, "save_unbalanced_param", True)
+
 
 if TYPE_CHECKING:
 
diff --git a/python/paddle/distributed/auto_parallel/sharding.py b/python/paddle/distributed/auto_parallel/sharding.py
@@ -768,6 +768,7 @@ def convert_state_dict_without_tensor_fusion_param(self, state_dict):
         # when the model is saved, we no need to save the slice@ parameters
         for name in slice_param_names:
             del state_dict[name]
+        paddle.device.cuda.empty_cache()
 
         if self._dy_shard_group is None:
             self._create_dy_sharding_group()
@@ -784,6 +785,7 @@ def convert_state_dict_without_tensor_fusion_param(self, state_dict):
             self._broadcast_pow_acc_opt_params(
                 state_dict, group_info, pow_acc_opt_param_names
             )
+            paddle.device.cuda.empty_cache()
 
     def _create_dy_sharding_group(self):
         mesh = self._shard_fn._mesh
@@ -903,6 +905,7 @@ def _re_slicing_opt_param(
                 if cur_rank in group_rank_mapping[idx]:
                     # param tensor may be sliced into multiple devices
                     # we need calculate the start index of the current rank
+                    opt_param = state_dict[param_name + param_suffix]
                     cur_rank_start_index = param_index
                     for i, rank_id in enumerate(group_rank_mapping[idx]):
                         if rank_id == cur_rank:
@@ -947,7 +950,9 @@ def _all_gather_opt_params(
                 continue
             if opt_param_name not in opt_param_names:
                 continue
-            opt_param_list.append(state_dict[opt_param_name]._local_value())
+            opt_param_list.append(
+                state_dict[opt_param_name]._local_value().clone()
+            )
 
         if len(opt_param_list) == 0:
             return
@@ -958,11 +963,20 @@ def _all_gather_opt_params(
         dist.all_gather(
             fused_opt_param_list, fused_opt_param, group=self._dy_shard_group
         )
-        del fused_opt_param
-        paddle.device.cuda.empty_cache()
 
+        fused_opt_param_list = [item.cpu() for item in fused_opt_param_list]
         fused_opt_param = paddle.concat(fused_opt_param_list, axis=0)
-        del fused_opt_param_list
+
+        for param_name, param_info in group_info.items():
+            opt_param_name = "slice@" + param_name + opt_suffix
+            if opt_param_name not in state_dict:
+                continue
+            if opt_param_name not in opt_param_names:
+                continue
+
+            local_tensor = state_dict[opt_param_name]._local_value()
+            del state_dict[opt_param_name]
+
         paddle.device.cuda.empty_cache()
 
         param_index = 0
@@ -1015,20 +1029,18 @@ def _all_gather_opt_params(
             shard_opt_param = shard_opt_param[tuple(shard_index)]
 
             shard_opt_param = _dtensor_from_local(
-                shard_opt_param,
+                shard_opt_param.cuda(),
                 opt_param_mesh,
                 opt_param_placements,
                 shard_opt_param.shape,
             )
 
             state_dict[param_name + opt_suffix] = shard_opt_param
-            # remove the slice@ parameter
-            if opt_param_name in state_dict:
-                del state_dict[opt_param_name]
-
             padded_size = param_info["padded_size"]
             param_index += padded_size
 
+        paddle.device.cuda.empty_cache()
+
     def _all_gather_moment_opt_params(
         self, state_dict, group_info, moment_opt_param_names
     ):
@@ -1145,3 +1157,28 @@ def _bucket_tensors_with_group_size(self, group_info, group_size):
                     current_size = 0
 
         return group_mapping, size_mapping
+
+    def convert_state_dict_with_rank_unique_name(self, state_dict):
+        cur_rank = dist.get_rank()
+        tensor_names = list(state_dict.keys())
+
+        for name in tensor_names:
+            tensor = state_dict[name]
+            if not tensor.is_dist():
+                continue
+            if "slice@" not in name:
+                continue
+
+            if "_moment" in name or "_pow_acc" in name or "_master" in name:
+                rank_name = f"{name}_rank{cur_rank}"
+                state_dict[rank_name] = state_dict[name]
+
+            del state_dict[name]
+
+    def convert_state_dict_with_origin_name(self, state_dict):
+        tensor_names = list(state_dict.keys())
+        for name in list(state_dict.keys()):
+            if "_rank" in name:
+                no_rank_name = name.split("_rank")[0]
+                state_dict[no_rank_name] = state_dict[name]
+                del state_dict[name]
diff --git a/python/paddle/distributed/checkpoint/save_state_dict.py b/python/paddle/distributed/checkpoint/save_state_dict.py
@@ -273,7 +273,7 @@ def save_state_dict(
         if coordinator_rank == paddle.distributed.get_rank():
             logger.debug(f"metadata:{metadata}")
             paddle.save(metadata, os.path.join(path, f"{unique_id}.metadata"))
-        logger.debug(f"local_state_dict:{local_state_dict}")
+
         dedup_tensor(
             local_state_dict, local_storage_metadata, metadata.storage_metadata
         )
diff --git a/test/auto_parallel/pir/sharding_tensor_fusion_save_load.py b/test/auto_parallel/pir/sharding_tensor_fusion_save_load.py
@@ -68,21 +68,24 @@ def __init__(self):
         self._amp_dtype = os.getenv("amp_dtype", 'float16')
         self._amp_level = os.getenv("amp_level", 'O0')
         self._init_loss_scaling = 1024.0
-        self.mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        self.mesh = dist.ProcessMesh([0, 1], dim_names=["dp"])
         self._in_pir_mode = paddle.base.framework.get_flags(
             "FLAGS_enable_pir_api"
         )["FLAGS_enable_pir_api"]
         self.num_batch = 2
+        self.save_unbalanced_param = int(
+            os.getenv("save_unbalanced_param", '1')
+        )
 
     def set_random_seed(self, seed):
         random.seed(seed)
         np.random.seed(seed)
         paddle.seed(seed)
 
     def create_data_loader(self, return_dict=False):
-        images = np.random.rand(BATCH_SIZE, IMAGE_SIZE).astype('float32')
-        labels = np.random.rand(BATCH_SIZE, CLASS_NUM).astype('float32')
-        dataset = RandomDataset(images, labels, BATCH_SIZE, return_dict)
+        images = np.random.rand(100, IMAGE_SIZE).astype('float32')
+        labels = np.random.rand(100, CLASS_NUM).astype('float32')
+        dataset = RandomDataset(images, labels, 100, return_dict)
         loader = DataLoader(dataset, batch_size=BATCH_SIZE)
         return loader
 
@@ -138,22 +141,25 @@ def run_dy2static(self):
         lr_scheduler = paddle.optimizer.lr.LinearWarmup(
             learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
         )
-        opt = paddle.optimizer.Adam(
+        opt = paddle.optimizer.AdamW(
             learning_rate=lr_scheduler,
             parameters=layer.parameters(),
         )
+        opt = dist.shard_optimizer(opt, dist.ShardingStage1("dp", self.mesh))
         dist_loader = dist.shard_dataloader(
             dataloader=data_loader,
             meshes=[self.mesh],
             input_keys=["image", "label"],
-            shard_dims=['x'],
+            shard_dims=["dp"],
         )
 
         loss_fn = nn.MSELoss()
         strategy = dist.Strategy()
         strategy.sharding.enable = True
         strategy.sharding.degree = 2
         strategy.sharding.stage = 1
+        strategy.sharding.enable_stage1_tensor_fusion = True
+        strategy.sharding.save_unbalanced_param = self.save_unbalanced_param
 
         if self._amp:
             layer, opt = paddle.amp.decorate(
@@ -195,30 +201,10 @@ def run_dy2static(self):
                 loss_md5 = hashlib.md5(array_bytes).hexdigest()
                 loss_before_save.append(loss_md5)
 
-                if int(dist.get_rank()) in [2, 3, 6, 7]:
-                    assert loss is not None
-                else:
-                    assert loss is None
-
             if step >= 9:
                 break
 
-        # check pir dist_model save&load
-        paddle.enable_static()
-        model_file_path = os.path.join(
-            self._ckpt_path,
-            "rank_" + str(paddle.distributed.get_rank()) + ".pd_dist_model",
-        )
-        paddle.save(
-            dist_model._engine._pir_dist_main_progs["train"], model_file_path
-        )
-        loaded_model = paddle.load(model_file_path)
-        self.check_program_equal(
-            dist_model._engine._pir_dist_main_progs["train"], loaded_model
-        )
-        paddle.disable_static()
         paddle.distributed.barrier()
-
         time.sleep(10)
         loss_after_load = []
         for step, inputs in enumerate(dist_loader()):
@@ -230,24 +216,21 @@ def run_dy2static(self):
             if step == 2:
                 state_dict = dist_model.state_dict()
                 dist.load_state_dict(state_dict, self._ckpt_path)
+                dist_model.set_state_dict(state_dict)
             if step > 2:
                 numpy_array = np.array(loss)
                 array_bytes = numpy_array.tobytes()
                 loss_md5 = hashlib.md5(array_bytes).hexdigest()
                 loss_after_load.append(loss_md5)
 
-                if int(dist.get_rank()) == 1:
-                    assert loss is not None
-                else:
-                    assert loss is None
             if step >= 9:
                 break
 
         return (loss_before_save, loss_after_load)
 
     def run_test_case(self):
         loss = self.run_dy2static()
-        if int(dist.get_rank()) == 1:
+        if int(dist.get_rank()) == 0:
             assert len(loss[0]) == len(loss[1])
             for i in range(len(loss[0])):
                 assert loss[0][i] == loss[1][i]
diff --git a/test/auto_parallel/pir/test_sharding_tensor_fusion_save_load.py b/test/auto_parallel/pir/test_sharding_tensor_fusion_save_load.py

Original file line number	Diff line number	Diff line change
`@@ -273,7 +273,7 @@ def save_state_dict(`
`273`	`273`	`if coordinator_rank == paddle.distributed.get_rank():`
`274`	`274`	`logger.debug(f"metadata:{metadata}")`
`275`	`275`	`paddle.save(metadata, os.path.join(path, f"{unique_id}.metadata"))`
`276`		`- logger.debug(f"local_state_dict:{local_state_dict}")`
	`276`	`+`
`277`	`277`	`dedup_tensor(`
`278`	`278`	`local_state_dict, local_storage_metadata, metadata.storage_metadata`
`279`	`279`	`)`