Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/shardformer/policies/autopolicy.py b/colossalai/shardformer/policies/autopolicy.py
--- a/colossalai/shardformer/policies/autopolicy.py
+++ b/colossalai/shardformer/policies/autopolicy.py
@@ -71,6 +71,14 @@ class PolicyLocation:
     "transformers.models.gpt2.modeling_gpt2.GPT2ForSequenceClassification":
         PolicyLocation(file_name="gpt2", class_name="GPT2ForSequenceClassificationPolicy"),
 
+    # ViT
+    "transformers.models.vit.modeling_vit.ViTModel":
+        PolicyLocation(file_name="vit", class_name="ViTPolicy"),
+    "transformers.models.vit.modeling_vit.ViTForImageClassification":
+        PolicyLocation(file_name="vit", class_name="ViTForImageClassificationPolicy"),
+    "transformers.models.vit.modeling_vit.ViTForMaskedImageModeling":
+        PolicyLocation(file_name="vit", class_name="ViTForMaskedImageModelingPolicy"),
+
     # OPT
     "transformers.models.opt.modeling_opt.OPTModel":
         PolicyLocation(file_name="opt", class_name="OPTModelPolicy"),
diff --git a/colossalai/shardformer/policies/vit.py b/colossalai/shardformer/policies/vit.py
--- a/colossalai/shardformer/policies/vit.py
+++ b/colossalai/shardformer/policies/vit.py
@@ -2,11 +2,11 @@
 
 import torch.nn as nn
 
-from colossalai.shardformer.layer import DropoutForReplicatedInput, FusedLayerNorm, Linear1D_Col, Linear1D_Row
+from colossalai.shardformer.layer import DropoutForReplicatedInput, DropoutForParallelInput, FusedLayerNorm, Linear1D_Col, Linear1D_Row
 
 from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
 
-__all__ = ['ViTPolicy']
+__all__ = ['ViTPolicy', 'ViTForImageClassificationPolicy', 'ViTForMaskedImageModelingPolicy']
 
 
 class ViTPolicy(Policy):
@@ -15,96 +15,104 @@ def config_sanity_check(self):
         pass
 
     def preprocess(self):
-        # Resize embedding
-        vocab_size = self.model.config.vocab_size
-        world_size = self.shard_config.tensor_parallel_size
-
-        if vocab_size % world_size != 0:
-            new_vocab_size = vocab_size + world_size - vocab_size % world_size
-            self.model.resize_token_embeddings(new_vocab_size)
-
         return self.model
 
     def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
         from transformers.models.vit.modeling_vit import ViTEmbeddings, ViTLayer
 
-        base_policy = {
-            ViTEmbeddings:
-                ModulePolicyDescription(sub_module_replacement=[
-                    SubModuleReplacementDescription(
-                        suffix="dropout",
-                        target_module=DropoutForReplicatedInput,
-                    )
-                ]),
-            ViTLayer:
-                ModulePolicyDescription(attribute_replacement={
-                    "attention.attention.num_attention_heads":
-                        self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
-                    "attention.attention.all_head_size":
-                        self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
-                },
+        policy = {}
+
+        if self.shard_config.enable_tensor_parallelism:
+            policy[ViTEmbeddings] = ModulePolicyDescription(attribute_replacement={},
+                                        param_replacement=[],
                                         sub_module_replacement=[
                                             SubModuleReplacementDescription(
-                                                suffix="attention.attention.query",
-                                                target_module=Linear1D_Col,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="attention.attention.key",
-                                                target_module=Linear1D_Col,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="attention.attention.value",
-                                                target_module=Linear1D_Col,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="attention.attention.dropout",
-                                                target_module=DropoutForParallelInput,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="attention.output.dense",
-                                                target_module=Linear1D_Row,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="attention.output.dropout",
-                                                target_module=DropoutForParallelInput,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="intermediate.dense",
-                                                target_module=Linear1D_Col,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="output.dense",
-                                                target_module=Linear1D_Row,
-                                            ),
-                                            SubModuleReplacementDescription(
-                                                suffix="output.dropout",
-                                                target_module=DropoutForParallelInput,
-                                            ),
-                                        ]),
-        }
-
-        # optimization configuration
-        if self.shard_config.enable_fused_normalization:
-            base_policy[ViTAttention].sub_module_replacement.extend([
-                SubModuleReplacementDescription(
-                    suffix="layernorm_before",
-                    target_module=FusedLayerNorm,
-                ),
-                SubModuleReplacementDescription(
-                    suffix="layernorm_after",
-                    target_module=FusedLayerNorm,
+                                                suffix="dropout",
+                                                target_module=DropoutForReplicatedInput,
+                                            )
+                                        ])
+            
+            policy[ViTLayer] = ModulePolicyDescription(
+                    attribute_replacement={
+                        "attention.attention.num_attention_heads":
+                            self.model.config.num_attention_heads//self.shard_config.tensor_parallel_size,
+                        "attention.attention.all_head_size":
+                            self.model.config.hidden_size//self.shard_config.tensor_parallel_size,
+                    },
+                    param_replacement=[],
+                    sub_module_replacement=[
+                        SubModuleReplacementDescription(
+                            suffix="attention.attention.query",
+                            target_module=Linear1D_Col,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="attention.attention.key",
+                            target_module=Linear1D_Col,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="attention.attention.value",
+                            target_module=Linear1D_Col,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="attention.attention.dropout",
+                            target_module=DropoutForParallelInput,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="attention.output.dense",
+                            target_module=Linear1D_Row,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="attention.output.dropout",
+                            target_module=DropoutForReplicatedInput,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="intermediate.dense",
+                            target_module=Linear1D_Col,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="output.dense",
+                            target_module=Linear1D_Row,
+                        ),
+                        SubModuleReplacementDescription(
+                            suffix="output.dropout",
+                            target_module=DropoutForReplicatedInput,
+                        ),
+                    ]
                 )
-            ])
-            base_policy[ViTModel].sub_module_replacement.append(
-                SubModuleReplacementDescription(
-                    suffix="layernorm",
-                    target_module=FusedLayerNorm,
-                ))
-
-        return base_policy
 
+        return policy
+  
+    
     def new_model_class(self):
         return None
 
     def postprocess(self):
         return self.model
+
+class ViTForImageClassificationPolicy(ViTPolicy):
+
+     def module_policy(self):
+        from transformers.models.vit.modeling_vit import ViTForImageClassification
+
+        policy = super().module_policy()
+        if self.shard_config.enable_tensor_parallelism:
+            new_item = {
+                ViTForImageClassification:
+                ModulePolicyDescription(sub_module_replacement=[
+                                        SubModuleReplacementDescription(suffix="classifier",
+                                                                            target_module=Linear1D_Col,
+                                                                            kwargs=dict(gather_output=True))
+                                        ])
+            }
+            policy.update(new_item)
+        return policy
+
+class ViTForMaskedImageModelingPolicy(ViTPolicy):
+    
+    def module_policy(self):
+        policy = super().module_policy()
+        return policy
+    
+
+        
+


Test patch:
diff --git a/tests/kit/model_zoo/transformers/__init__.py b/tests/kit/model_zoo/transformers/__init__.py
--- a/tests/kit/model_zoo/transformers/__init__.py
+++ b/tests/kit/model_zoo/transformers/__init__.py
@@ -5,3 +5,4 @@
 from .llama import *
 from .opt import *
 from .t5 import *
+from .vit import *
diff --git a/tests/kit/model_zoo/transformers/vit.py b/tests/kit/model_zoo/transformers/vit.py
new file mode 100644
--- /dev/null
+++ b/tests/kit/model_zoo/transformers/vit.py
@@ -0,0 +1,60 @@
+import torch
+import transformers
+
+from ..registry import ModelAttribute, model_zoo
+
+# ===============================
+# Register single-sentence VIT
+# ===============================
+
+config = transformers.ViTConfig(num_hidden_layers=4,
+                                hidden_size=128,
+                                intermediate_size=256,
+                                num_attention_heads=4)
+
+# define data gen function
+def data_gen():
+    pixel_values = torch.randn(1, 3, 224, 224)
+    return dict(pixel_values = pixel_values)
+
+def data_gen_for_masked_image_modeling():
+    data = data_gen()
+    num_patches = (config.image_size // config.patch_size) ** 2
+    bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
+    data['bool_masked_pos'] = bool_masked_pos
+    return data
+
+# define output transform function
+output_transform_fn = lambda x: x
+
+# function to get the loss
+loss_fn_for_vit_model = lambda x : x.pooler_output.mean()
+loss_fn_for_image_classification = lambda x : x.logits.mean()
+loss_fn_for_masked_image_modeling = lambda x : x.loss
+
+# register the following models
+# transformers.ViTModel,
+# transformers.ViTForMaskedImageModeling,
+# transformers.ViTForImageClassification,
+model_zoo.register(name = 'transformers_vit',
+                    model_fn = lambda : transformers.ViTModel(config),
+                    data_gen_fn = data_gen,
+                    output_transform_fn = output_transform_fn,
+                    loss_fn = loss_fn_for_vit_model,
+                    model_attribute = ModelAttribute(has_control_flow=True))
+
+model_zoo.register(name = 'transformers_vit_for_masked_image_modeling',
+                    model_fn = lambda : transformers.ViTForMaskedImageModeling(config),
+                    data_gen_fn = data_gen_for_masked_image_modeling,
+                    output_transform_fn = output_transform_fn,
+                    loss_fn = loss_fn_for_masked_image_modeling,
+                    model_attribute = ModelAttribute(has_control_flow=True))
+
+model_zoo.register(name = 'transformers_vit_for_image_classification',
+                    model_fn = lambda : transformers.ViTForImageClassification(config),
+                    data_gen_fn = data_gen,
+                    output_transform_fn = output_transform_fn,
+                    loss_fn = loss_fn_for_image_classification,
+                    model_attribute = ModelAttribute(has_control_flow=True))
+
+
diff --git a/tests/test_shardformer/test_model/test_shard_vit.py b/tests/test_shardformer/test_model/test_shard_vit.py
--- a/tests/test_shardformer/test_model/test_shard_vit.py
+++ b/tests/test_shardformer/test_model/test_shard_vit.py
@@ -1,9 +1,18 @@
+import os
+
 import pytest
 import torch
 
 import colossalai
 from colossalai.logging import disable_existing_loggers
-from colossalai.testing import assert_hf_output_close, clear_cache_before_run, rerun_if_address_is_in_use, spawn
+from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor
+from colossalai.testing import (
+    assert_hf_output_close,
+    clear_cache_before_run,
+    parameterize,
+    rerun_if_address_is_in_use,
+    spawn,
+)
 from tests.kit.model_zoo import model_zoo
 from tests.test_shardformer.test_model._utils import build_model, run_forward
 
@@ -12,44 +21,59 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
     # check forward
     org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn,
                                                                  output_transform_fn, loss_fn)
-    assert_hf_output_close(org_output, shard_output)
-
+    assert_hf_output_close(org_output, shard_output, atol=1e-4, rtol=1e-4)
     # do backward
     org_loss.backward()
     shard_loss.backward()
 
-    # check grad
-    org_grad = org_model.encoder.layer[0].attention.attention.query.weight.grad
-    shard_grad = sharded_model.encoder.layer[0].attention.attention.query.weight.grad
-
-    shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)]
-    shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad)
-    all_shard_grad = torch.cat(shard_grad_list, dim=0)
-
     assert torch.allclose(org_loss, shard_loss,
                           atol=1e-5), f"shard model loss is not equal to orgin model loss\n{org_loss}\n{shard_loss}"
-    assert torch.allclose(org_grad, all_shard_grad,
-                          atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}"
 
+    # unwrap model
+    if org_model.__class__.__name__ == 'ViTModel':
+        vit_model = org_model
+        shard_vit_model = sharded_model
+    else:
+        vit_model = org_model.vit
+        shard_vit_model = sharded_model.vit
 
-def check_vit(rank, world_size, port):
-    disable_existing_loggers()
-    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    # check attention grad
+    org_grad = vit_model.encoder.layer[0].attention.attention.query.weight.grad
+    shard_grad = shard_vit_model.encoder.layer[0].attention.attention.query.weight.grad
+    shard_weight = shard_vit_model.encoder.layer[0].attention.attention.query.weight
 
+    if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight):
+        shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)]
+        shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad)
+        all_shard_grad = torch.cat(shard_grad_list, dim=0)
+    else:
+        all_shard_grad = shard_grad
+    assert torch.allclose(org_grad, all_shard_grad,
+                          atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{shard_grad}"
+    
+
+@parameterize('enable_fused_normalization', [True, False])
+@parameterize('enable_tensor_parallelism', [True, False])
+def run_vit_test(enable_fused_normalization, enable_tensor_parallelism):
     sub_model_zoo = model_zoo.get_sub_registry('transformers_vit')
     for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
-        org_model, sharded_model = build_model(world_size, model_fn)
+        org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism)
         check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn)
-
     torch.cuda.empty_cache()
 
 
+def check_vit(rank, world_size, port):
+    disable_existing_loggers()
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    run_vit_test()
+
+
 @pytest.mark.dist
 @pytest.mark.skip
 @rerun_if_address_is_in_use()
 @clear_cache_before_run()
 def test_vit():
-    spawn(check_vit, 4)
+    spawn(check_vit, 2)
 
 
 if __name__ == "__main__":