[AutoParallel] fix bugs in DistModel (#61313)

deepllz · web-flow · commit 7f0949b5423e · 2024-01-31T17:19:38.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -1222,8 +1222,13 @@ def _to_lodtensor(tensor: paddle.Tensor):
                         tensor._local_value().get_tensor()
                     )
                 else:
+                    # infer dtype from tensor
+                    if tensor.is_integer():
+                        dtype = paddle.iinfo(tensor.dtype).dtype
+                    else:
+                        dtype = paddle.finfo(tensor.dtype).dtype
                     tensor_np_value = np.zeros(
-                        tensor._local_value().shape, dtype=np.float32
+                        tensor._local_value().shape, dtype=dtype
                     )
                     lodtensor.set(
                         tensor_np_value,
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -304,6 +304,8 @@ def _prepare_data_spec_from_dataloader(self, dataloader):
                     DistrubutedInputSpec.from_dtensor(item, name)
                 )
 
+        inputs_spec = self._validate_spec(inputs_spec)
+        labels_spec = self._validate_spec(labels_spec)
         return inputs_spec, labels_spec
 
     def _prepare_data_spec(self, data, split, batch_size):
@@ -1835,9 +1837,11 @@ def _validate_spec(self, specs):
         specs = auto_utils.to_list(specs)
         if specs is not None:
             for i, spec in enumerate(specs):
-                if not isinstance(spec, InputSpec):
+                if not isinstance(spec, InputSpec) and not isinstance(
+                    spec, DistrubutedInputSpec
+                ):
                     raise TypeError(
-                        "'spec' must be object of class `paddle.static.InputSpec`."
+                        "'spec' must be object of class `paddle.static.InputSpec` or `DistrubutedInputSpec`."
                     )
                 if spec.name is None:
                     raise ValueError(
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_llama_pp_gradmerge.py b/test/auto_parallel/hybrid_strategy/semi_auto_llama_pp_gradmerge.py
@@ -19,6 +19,7 @@
 from semi_auto_parallel_llama_model import (
     LlamaForCausalLMAuto,
     LlamaPretrainingCriterionAuto,
+    get_mesh,
     set_global_mesh,
 )
 
@@ -233,8 +234,13 @@ def run_llama(self, to_static=0):
                 num_workers=0,
             )
 
-            dist_model, dist_loader = dist.to_static(
-                model, train_dataloader, criterion, optimizer, strategy=strategy
+            dist_loader = dist.shard_dataloader(
+                train_dataloader,
+                meshes=[get_mesh(0), get_mesh(-1)],
+                shard_dims="dp",
+            )
+            dist_model = dist.to_static(
+                model, dist_loader, criterion, optimizer, strategy=strategy
             )
 
             def validate_batch(batch):
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_parallel_llama_model.py b/test/auto_parallel/hybrid_strategy/semi_auto_parallel_llama_model.py
@@ -536,7 +536,7 @@ def _prepare_decoder_attention_mask(
                     combined_attention_mask = dist.shard_tensor(
                         combined_attention_mask,
                         mesh,
-                        [dist.Replicate(), dist.Replicate()],
+                        [dist.Replicate() for _ in range(len(mesh._shape))],
                     )
                     expanded_attn_mask = (
                         expanded_attn_mask & combined_attention_mask
@@ -648,7 +648,9 @@ def forward(
                 (batch_size, seq_length)
             )
             position_ids = dist.shard_tensor(
-                position_ids, mesh, [dist.Replicate(), dist.Replicate()]
+                position_ids,
+                mesh,
+                [dist.Replicate() for _ in range(len(mesh._shape))],
             )
 
         attention_mask = self._prepare_decoder_attention_mask(

Original file line number	Diff line number	Diff line change
`@@ -536,7 +536,7 @@ def _prepare_decoder_attention_mask(`
`536`	`536`	`combined_attention_mask = dist.shard_tensor(`
`537`	`537`	`combined_attention_mask,`
`538`	`538`	`mesh,`
`539`		`- [dist.Replicate(), dist.Replicate()],`
	`539`	`+ [dist.Replicate() for _ in range(len(mesh._shape))],`
`540`	`540`	`)`
`541`	`541`	`expanded_attn_mask = (`
`542`	`542`	`expanded_attn_mask & combined_attention_mask`
`@@ -648,7 +648,9 @@ def forward(`
`648`	`648`	`(batch_size, seq_length)`
`649`	`649`	`)`
`650`	`650`	`position_ids = dist.shard_tensor(`
`651`		`- position_ids, mesh, [dist.Replicate(), dist.Replicate()]`
	`651`	`+ position_ids,`
	`652`	`+ mesh,`
	`653`	`+ [dist.Replicate() for _ in range(len(mesh._shape))],`
`652`	`654`	`)`
`653`	`655`
`654`	`656`	`attention_mask = self._prepare_decoder_attention_mask(`