Update hybrid_parallel_plugin.py

fix

fix

Author: flybird11111

colossalai/booster/plugin/hybrid_parallel_plugin.py

@@ -931,7 +931,7 @@ class HybridParallelPlugin(PipelinePluginBase):
         pp_style (str, optional): The style for pipeline parallelism. Defaults to '1f1b'.
         num_model_chunks (int, optional): The number of model chunks for interleaved pipeline parallelism. Defaults to 1.
         enable_metadata_cache (bool, optional): Whether to enable metadata cache for pipeline parallelism. Defaults to True.
-        make_vocab_size_divisible_by (int, optional): it's used when padding the vocabulary size, to make it choose an faster kenel. Default to 128.
+        make_vocab_size_divisible_by (int, optional): it's used when padding the vocabulary size, to make it choose an faster kenel. Default to 64.
     """
 
     def __init__(

colossalai/shardformer/layer/embedding.py

@@ -192,7 +192,8 @@ def __init__(
 
         super().__init__(self.num_embeddings, num_embeddings, weight)
 
-        self.resize_embedding_weight()
+        if weight.shape[0] < self.num_embeddings:
+            self.resize_embedding_weight()
 
         if weight is None:
             self.reset_parameters()
@@ -306,7 +307,8 @@ def __init__(
 
         # resize vocabulary size
         super().__init__(self.num_embeddings, num_embeddings, weight)
-        self.resize_embedding_weight()
+        if not is_distributed_tensor(self.weight):
+            self.resize_embedding_weight()
 
         # deal with tensor parallelism
         self.num_embeddings_per_partition = divide(self.num_embeddings, tensor_parallel_size)

colossalai/shardformer/policies/llama.py

@@ -297,7 +297,6 @@ def module_policy(self):
                     ],
                 )
             }
-        print("new_item", new_item)
         policy.update(new_item)
 
         if self.pipeline_stage_manager:

colossalai/shardformer/policies/opt.py

@@ -65,7 +65,6 @@ def module_policy(self):
         if self.shard_config.enable_tensor_parallelism:
             embedding_cls = VocabParallelEmbedding1D
         else:
-            # TODO when not tie weight and not pad the vocab size
             if self.tie_weight:
                 embedding_cls = PaddingEmbedding
 

colossalai/shardformer/policies/t5.py

@@ -11,6 +11,8 @@
     FusedRMSNorm,
     Linear1D_Col,
     Linear1D_Row,
+    PaddingEmbedding,
+    PaddingLMHead,
     RMSNorm,
     VocabParallelEmbedding1D,
     VocabParallelLMHead1D,
@@ -35,22 +37,18 @@ def config_sanity_check(self):
         pass
 
     def preprocess(self):
-        # reshape the embedding layer
-        r"""
-        Reshape the Embedding layer to make the embedding dimension divisible by world_size
-        """
-        # TODO padding the vocab size in VocabParallelEmbedding1D
-        # vocab_size = self.model.config.vocab_size
-        # if self.shard_config.enable_tensor_parallelism:
-        #     world_size = self.shard_config.tensor_parallel_size
-        #     multiple = world_size * self.shard_config.make_vocab_size_divisible_by
-        # else:
-        #     multiple = self.shard_config.make_vocab_size_divisible_by
-        # if vocab_size % multiple != 0:
-        #     new_vocab_size = vocab_size + multiple - vocab_size % multiple
-        #     self.model.resize_token_embeddings(new_vocab_size)
+        self.tie_weight = self.tie_weight_check()
         return self.model
 
+    def tie_weight_check(self):
+        input_embedding = self.model.get_input_embeddings()
+        output_embedding = self.model.get_output_embeddings()
+        return (
+            input_embedding is not None
+            and output_embedding is not None
+            and id(input_embedding.weight) == id(output_embedding.weight)
+        )
+
     def module_policy(self):
         from transformers.models.t5.modeling_t5 import (
             T5Attention,
@@ -64,6 +62,13 @@ def module_policy(self):
 
         policy = {}
 
+        embedding_cls = None
+        if self.shard_config.enable_tensor_parallelism:
+            embedding_cls = VocabParallelEmbedding1D
+        else:
+            if self.tie_weight:
+                embedding_cls = PaddingEmbedding
+
         if self.shard_config.enable_fused_normalization:
             norm_cls = FusedRMSNorm
         else:
@@ -80,11 +85,6 @@ def module_policy(self):
                         suffix="dropout",
                         target_module=DropoutForParallelInput,
                     ),
-                    SubModuleReplacementDescription(
-                        suffix="embed_tokens",
-                        target_module=VocabParallelEmbedding1D,
-                        kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
-                    ),
                 ]
             )
             policy[T5LayerSelfAttention] = ModulePolicyDescription(
@@ -180,6 +180,17 @@ def module_policy(self):
                 ]
             )
 
+        if embedding_cls is not None:
+            self.append_or_create_submodule_replacement(
+                description=SubModuleReplacementDescription(
+                    suffix="embed_tokens",
+                    target_module=embedding_cls,
+                    kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
+                ),
+                policy=policy,
+                target_key=T5Stack,
+            )
+
         # optimization configuration
         self.append_or_create_submodule_replacement(
             description=SubModuleReplacementDescription(
@@ -371,11 +382,18 @@ def module_policy(self):
 
         policy = super().module_policy()
 
+        embedding_cls = None
         if self.shard_config.enable_tensor_parallelism:
+            embedding_cls = VocabParallelEmbedding1D
+        else:
+            if self.tie_weight:
+                embedding_cls = PaddingEmbedding
+
+        if embedding_cls is not None:
             self.append_or_create_submodule_replacement(
                 description=SubModuleReplacementDescription(
                     suffix="shared",
-                    target_module=VocabParallelEmbedding1D,
+                    target_module=embedding_cls,
                     kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
                 ),
                 policy=policy,
@@ -408,23 +426,44 @@ def module_policy(self):
 
         policy = super().module_policy()
 
+        embedding_cls = None
         if self.shard_config.enable_tensor_parallelism:
+            embedding_cls = VocabParallelEmbedding1D
+        else:
+            if self.tie_weight:
+                embedding_cls = PaddingEmbedding
+
+        if embedding_cls is not None:
             self.append_or_create_submodule_replacement(
-                description=[
-                    SubModuleReplacementDescription(
-                        suffix="shared",
-                        target_module=VocabParallelEmbedding1D,
-                        kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
-                    ),
-                    SubModuleReplacementDescription(
-                        suffix="lm_head",
-                        target_module=VocabParallelLMHead1D,
-                        kwargs=dict(
-                            gather_output=True,
-                            make_vocab_size_divisible_by=self.shard_config.make_vocab_size_divisible_by,
-                        ),
-                    ),
-                ],
+                description=SubModuleReplacementDescription(
+                    suffix="shared",
+                    target_module=embedding_cls,
+                    kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
+                ),
+                policy=policy,
+                target_key=T5ForConditionalGeneration,
+            )
+
+        if self.shard_config.enable_tensor_parallelism:
+            self.append_or_create_submodule_replacement(
+                description=SubModuleReplacementDescription(
+                    suffix="lm_head",
+                    target_module=VocabParallelLMHead1D,
+                    kwargs={
+                        "gather_output": True,
+                        "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by,
+                    },
+                ),
+                policy=policy,
+                target_key=T5ForConditionalGeneration,
+            )
+        else:
+            self.append_or_create_submodule_replacement(
+                description=SubModuleReplacementDescription(
+                    suffix="lm_head",
+                    target_module=PaddingLMHead,
+                    kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
+                ),
                 policy=policy,
                 target_key=T5ForConditionalGeneration,
             )
@@ -475,11 +514,18 @@ def module_policy(self):
 
         policy = super().module_policy()
 
+        embedding_cls = None
         if self.shard_config.enable_tensor_parallelism:
+            embedding_cls = VocabParallelEmbedding1D
+        else:
+            if self.tie_weight:
+                embedding_cls = PaddingEmbedding
+
+        if embedding_cls is not None:
             self.append_or_create_submodule_replacement(
                 description=SubModuleReplacementDescription(
                     suffix="shared",
-                    target_module=VocabParallelEmbedding1D,
+                    target_module=embedding_cls,
                     kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
                 ),
                 policy=policy,

colossalai/shardformer/shard/sharder.py

@@ -39,7 +39,6 @@ def shard(self) -> List[Dict[int, Tensor]]:
         self._preprocess()
         # get shared params before release unheld layers, this avoid misjudgment of shared params (None is None)
         shared_params = self.policy.get_shared_params()
-        print("shared_params", shared_params)
         held_layers = self._release_unheld_layers()
         self._replace_module(include=held_layers)
         self._materialize()

tests/test_shardformer/test_model/test_shard_opt.py

@@ -60,7 +60,7 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
 
     # optimizer executes step
     org_optimizer.step()
-    # sharded_optimizer.step()
+    sharded_optimizer.step()
 
     # check last hidden state & loss
     if stage_manager is None or stage_manager.is_last_stage():

tests/test_shardformer/test_model/test_shard_t5.py

@@ -203,16 +203,14 @@ def check_t5_3d(rank, world_size, port):
     colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
     run_t5_3d_test()
 
-# TODO padding the vocab size in VocabParallelEmbedding1D
-@pytest.mark.skip("padding the vocab size in VocabParallelEmbedding1D")
+
 @pytest.mark.dist
 @rerun_if_address_is_in_use()
 @clear_cache_before_run()
 def test_t5():
     spawn(check_t5, 4)
 
-# TODO padding the vocab size in VocabParallelEmbedding1D
-@pytest.mark.skip("padding the vocab size in VocabParallelEmbedding1D")
+
 @pytest.mark.largedist
 @rerun_if_address_is_in_use()
 @clear_cache_before_run()