[Auto-Parallel] adapt optimizer_sharded_state_dict in FlexCheckpoint (PaddlePaddle#76305)

Author: Xing-lil

python/paddle/distributed/flex_checkpoint/dcp/sharded_weight.py

@@ -61,6 +61,8 @@ def __init__(
         self.key = key
         if local_tensor.is_dist():
             self.local_tensor = local_tensor._local_value()
+            # Note: The local_tensor must keep the same name with the original tensor. Otherwise, the static_to_struct_mapping will be wrong.
+            self.local_tensor.name = local_tensor.name
             self.local_shape = local_tensor._local_shape
         else:
             self.local_tensor = local_tensor

python/paddle/optimizer/adamw.py

@@ -797,11 +797,20 @@ def _generate_base_static_name(vname):
 
             # Determine tensor partitioning scheme
             if _MOMENT_NAME in optim_state_type:
-                optimizer_sharded_state_dict[unified_name] = (
-                    create_sharded_weight_with_new_local(
-                        unified_name, tensor, sharded_weight
+                if tensor.is_dist():
+                    optimizer_sharded_state_dict[unified_name] = ShardedWeight(
+                        key=unified_name,
+                        local_tensor=tensor,
+                        local_shape=tensor.shape,
+                        global_shape=tensor.shape,
+                        global_offset=sharded_weight.global_offset,
+                    )
+                else:
+                    optimizer_sharded_state_dict[unified_name] = (
+                        create_sharded_weight_with_new_local(
+                            unified_name, tensor, sharded_weight
+                        )
                     )
-                )
             else:  # Non-momentum parameters
                 optimizer_sharded_state_dict[unified_name] = ShardedWeight(
                     key=unified_name,
@@ -817,10 +826,19 @@ def _generate_base_static_name(vname):
                 struct_name = static_to_struct_mapping[key]
                 sharded_weight = model_sharded_state_dict[struct_name]
                 unified_name = f"{struct_name}.w_0"
-                optimizer_sharded_state_dict[unified_name] = (
-                    create_sharded_weight_with_new_local(
-                        unified_name, tensor, sharded_weight
+                if tensor.is_dist():
+                    optimizer_sharded_state_dict[unified_name] = ShardedWeight(
+                        key=unified_name,
+                        local_tensor=tensor,
+                        local_shape=tensor.shape,
+                        global_shape=tensor.shape,
+                        global_offset=sharded_weight.global_offset,
+                    )
+                else:
+                    optimizer_sharded_state_dict[unified_name] = (
+                        create_sharded_weight_with_new_local(
+                            unified_name, tensor, sharded_weight
+                        )
                     )
-                )
 
         return optimizer_sharded_state_dict

test/auto_parallel/semi_auto_parallel_for_flex_checkpoint.py

@@ -11,7 +11,7 @@
 # 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 os
 import random
 import tempfile
 
@@ -21,6 +21,10 @@
 import paddle.distributed as dist
 from paddle import nn
 
+MODEL_STATE_DIC = "model_state"
+OPTIMIZER_STATE_DIC = "optimizer_state"
+MASTER_WEIGHT_DIC = "master_weight"
+
 
 class SimpleModel(nn.Layer):
     def __init__(self, hidden_size=3072, layer_num=2):
@@ -50,44 +54,94 @@ def create_model_and_optimizer(self):
             learning_rate=0.001, parameters=model.parameters()
         )
         opt = dist.shard_optimizer(opt, dist.ShardingStage1("dp", self.mesh))
+        model, opt = paddle.amp.decorate(
+            model, optimizers=opt, level='O2', master_grad=True
+        )
         return model, opt
 
     def run_training_and_save(self):
         model, opt = self.create_model_and_optimizer()
+
         for step in range(3):
-            inputs = paddle.randn([self.batch_size, self.hidden_size])
-            labels = paddle.randn([self.batch_size, self.hidden_size])
+            inputs = paddle.ones(
+                [self.batch_size, self.hidden_size], dtype='float16'
+            )
+            labels = paddle.ones(
+                [self.batch_size, self.hidden_size], dtype='float16'
+            )
             inputs = dist.shard_tensor(inputs, self.mesh, [dist.Shard(0)])
             logits = model(inputs)
             loss = paddle.nn.functional.mse_loss(logits, labels)
             loss.backward()
-            opt.step()
+            if step == 2:
+                loss_md5 = loss._md5sum()
+            else:
+                opt.step()
             print(f"Train step {step}, loss: {loss.item()}")
+
         save_md5 = [p._md5sum() for p in model.parameters()]
+
+        # save model and optimizer
+        model_state_dict_path = os.path.join(self.ckpt_path, MODEL_STATE_DIC)
+        opt_state_dict_path = os.path.join(self.ckpt_path, OPTIMIZER_STATE_DIC)
+        master_weights_path = os.path.join(self.ckpt_path, MASTER_WEIGHT_DIC)
         sharded_state_dict = model.sharded_state_dict()
-        dist.save_state_dict(sharded_state_dict, self.ckpt_path)
-        return save_md5
+        dist.save_state_dict(sharded_state_dict, model_state_dict_path)
+        optimizer_states = {}
+        master_weights = {}
+        opt_sharded_state_dict = opt.sharded_state_dict(sharded_state_dict)
+        for k, v in opt_sharded_state_dict.items():
+            if k.endswith(".w_0"):
+                master_weights[k] = v
+            else:
+                optimizer_states[k] = v
+        dist.save_state_dict(optimizer_states, opt_state_dict_path)
+        dist.save_state_dict(master_weights, master_weights_path)
+        return save_md5, loss_md5
 
     def run_loading_and_validation(self):
         model, opt = self.create_model_and_optimizer()
+
+        # load model and optimizer
+        model_state_dict_path = os.path.join(self.ckpt_path, MODEL_STATE_DIC)
+        master_weights_path = os.path.join(self.ckpt_path, MASTER_WEIGHT_DIC)
+        opt_states_path = os.path.join(self.ckpt_path, OPTIMIZER_STATE_DIC)
         sharded_state_dict = model.sharded_state_dict()
-        dist.load_state_dict(sharded_state_dict, self.ckpt_path)
+        dist.load_state_dict(sharded_state_dict, model_state_dict_path)
+        opt_sharded_state_dict = opt.sharded_state_dict(sharded_state_dict)
+        opt_states = {}
+        master_weights = {}
+        for k, v in opt_sharded_state_dict.items():
+            if k.endswith(".w_0"):
+                master_weights[k] = v
+            else:
+                opt_states[k] = v
+        dist.load_state_dict(opt_states, opt_states_path)
+        dist.load_state_dict(master_weights, master_weights_path)
+
         load_md5 = [p._md5sum() for p in model.parameters()]
-        for step in range(3):
-            inputs = paddle.randn([self.batch_size, self.hidden_size])
-            labels = paddle.randn([self.batch_size, self.hidden_size])
+
+        for step in range(1):
+            inputs = paddle.ones(
+                [self.batch_size, self.hidden_size], dtype='float16'
+            )
+            labels = paddle.ones(
+                [self.batch_size, self.hidden_size], dtype='float16'
+            )
             inputs = dist.shard_tensor(inputs, self.mesh, [dist.Shard(0)])
             logits = model(inputs)
             loss = paddle.nn.functional.mse_loss(logits, labels)
             loss.backward()
             opt.step()
+            loss_md5 = loss._md5sum()
             print(f"Train step {step}, loss: {loss.item()}")
-        return load_md5
+        return load_md5, loss_md5
 
     def run_test(self):
-        save_param_md5sum = self.run_training_and_save()
-        load_param_md5sum = self.run_loading_and_validation()
+        save_param_md5sum, loss_md5 = self.run_training_and_save()
+        load_param_md5sum, loss_md5_reload = self.run_loading_and_validation()
         np.testing.assert_equal(save_param_md5sum, load_param_md5sum)
+        np.testing.assert_equal(loss_md5, loss_md5_reload)
 
 
 if __name__ == '__main__':