[shardformer] added embedding gradient check (#4124)

tests/kit/model_zoo/registry.py

@@ -70,6 +70,8 @@ class ModelZooRegistry(dict):
         for k, v in self.items():
             if keyword in k:
                 new_dict[k] = v
+
+        assert len(new_dict) > 0, f'No model found with keyword {keyword}'
         return new_dict
 
 

tests/test_shardformer/test_model/test_shard_bert.py

@@ -18,20 +18,35 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
     org_loss.backward()
     shard_loss.backward()
 
-    # check grad equality
+    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}"
+
+    # check grad
+
     if org_model.__class__.__name__ == 'BertModel':
-        org_grad = org_model.encoder.layer[0].attention.self.query.weight.grad
-        shard_grad = sharded_model.encoder.layer[0].attention.self.query.weight.grad
+        bert = org_model
+        sharded_bert = sharded_model
     else:
-        org_grad = org_model.bert.encoder.layer[0].attention.self.query.weight.grad
-        shard_grad = sharded_model.bert.encoder.layer[0].attention.self.query.weight.grad
+        bert = org_model.bert
+        sharded_bert = sharded_model.bert
+
+    # compare self attention grad
+    org_grad = bert.encoder.layer[0].attention.self.query.weight.grad
+    shard_grad = sharded_bert.encoder.layer[0].attention.self.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_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}"
 
-    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}"
+    # compare embedding grad
+    org_grad = bert.embeddings.word_embeddings.weight.grad
+    shard_grad = sharded_bert.embeddings.word_embeddings.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_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}"
 

tests/test_shardformer/test_model/test_shard_bloom.py

@@ -18,20 +18,36 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
     org_loss.backward()
     shard_loss.backward()
 
-    # check grad equality
+    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}"
+
+    # unwrap model
     if org_model.__class__.__name__ == 'BloomModel':
-        org_grad = org_model.h[0].self_attention.query_key_value.weight.grad
-        shard_grad = sharded_model.h[0].self_attention.query_key_value.weight.grad
+        bloom = org_model
+        sharded_bloom = sharded_model
     else:
-        org_grad = org_model.transformer.h[0].self_attention.query_key_value.weight.grad
-        shard_grad = sharded_model.transformer.h[0].self_attention.query_key_value.weight.grad
+        bloom = org_model.transformer
+        sharded_bloom = sharded_model.transformer
+
+    # check attention grad
+    org_grad = bloom.h[0].self_attention.query_key_value.weight.grad
+    shard_grad = sharded_bloom.h[0].self_attention.query_key_value.weight.grad
+
+    shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)]
+    torch.distributed.all_gather(shard_grad_list, shard_grad)
+    all_shard_grad = torch.cat(shard_grad_list, dim=0)
+
+    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}"
+
+    # check embedding weights
+    org_grad = bloom.word_embeddings.weight.grad
+    shard_grad = sharded_bloom.word_embeddings.weight.grad
 
     shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)]
     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}"
 

tests/test_shardformer/test_model/test_shard_gpt2.py

@@ -18,20 +18,36 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
     org_loss.backward()
     shard_loss.backward()
 
-    # check grad equality
+    assert torch.allclose(org_loss, shard_loss,
+                          atol=1e-5), f"shard model loss is not equal to origin model loss\n{org_loss}\n{shard_loss}"
+
+    # unwrap model
     if org_model.__class__.__name__ == 'GPT2Model':
-        org_grad = org_model.h[0].mlp.c_fc.weight.grad
-        shard_grad = sharded_model.h[0].mlp.c_fc.weight.grad
+        org_model = org_model
+        sharded_model = sharded_model
     else:
-        org_grad = org_model.transformer.h[0].mlp.c_fc.weight.grad
-        shard_grad = sharded_model.transformer.h[0].mlp.c_fc.weight.grad
+        org_model = org_model.transformer
+        sharded_model = sharded_model.transformer
+
+    # check mlp grad
+    org_grad = org_model.h[0].mlp.c_fc.weight.grad
+    shard_grad = sharded_model.h[0].mlp.c_fc.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=1)
 
-    assert torch.allclose(org_loss, shard_loss,
-                          atol=1e-5), f"shard model loss is not equal to origin 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 origin model grad\n{org_grad}\n{all_shard_grad}"
+
+    # check embedding weights
+    org_grad = org_model.wte.weight.grad
+    shard_grad = sharded_model.wte.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_grad, all_shard_grad,
         atol=1e-5), f"shard model grad is not equal to origin model grad\n{org_grad}\n{all_shard_grad}"

tests/test_shardformer/test_model/test_shard_llama.py

@@ -23,7 +23,10 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
     org_loss.backward()
     shard_loss.backward()
 
-    # check grad
+    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}"
+
+    # unwrap model
     if hasattr(org_model, 'model'):
         llama_model = org_model.model
         shard_llama_model = sharded_model.model
@@ -31,14 +34,21 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
         llama_model = org_model
         shard_llama_model = sharded_model
 
+    # check attention grad
     org_grad = llama_model.layers[0].self_attn.q_proj.weight.grad
     shard_grad = shard_llama_model.layers[0].self_attn.q_proj.weight.grad
     shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)]
     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_grad, all_shard_grad,
+                          atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{shard_grad}"
 
-    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}"
+    # check embedding grad
+    org_grad = llama_model.embed_tokens.weight.grad
+    shard_grad = shard_llama_model.embed_tokens.weight.grad
+    shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)]
+    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_grad, all_shard_grad,
                           atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{shard_grad}"
 

tests/test_shardformer/test_model/test_shard_opt.py

@@ -28,7 +28,10 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
     org_loss.backward()
     shard_loss.backward()
 
-    # check grad
+    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}"
+
+    # unwrap model
     if hasattr(org_model, 'model'):
         opt_model = org_model.model
         shard_opt_model = sharded_model.model
@@ -36,16 +39,23 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
         opt_model = org_model
         shard_opt_model = sharded_model
 
+    # check attention grad
     org_grad = opt_model.decoder.layers[0].self_attn.q_proj.weight.grad
     shard_grad = shard_opt_model.decoder.layers[0].self_attn.q_proj.weight.grad
-
     shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)]
-    shard_grad = torch.distributed.all_gather(shard_grad_list, 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{shard_grad}"
+                          atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}"
+
+    # check embedding grad
+    org_grad = opt_model.decoder.embed_tokens.weight.grad
+    shard_grad = shard_opt_model.decoder.embed_tokens.weight.grad
+    shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)]
+    torch.distributed.all_gather(shard_grad_list, shard_grad)
+    all_shard_grad = torch.cat(shard_grad_list, dim=0)
+    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}"
 
 
 def check_OPTModel(rank, world_size, port):
@@ -65,3 +75,7 @@ def check_OPTModel(rank, world_size, port):
 @clear_cache_before_run()
 def test_OPTModel():
     spawn(check_OPTModel, 4)
+
+
+if __name__ == '__main__':
+    test_OPTModel()

tests/test_shardformer/test_model/test_shard_t5.py

@@ -21,19 +21,43 @@ def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transfo
     org_loss.backward()
     shard_loss.backward()
 
-    # check grad equality
+    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}"
+
+    # check attention grad
     org_grad = org_model.encoder.block[0].layer[0].SelfAttention.q.weight.grad
     shard_grad = sharded_model.encoder.block[0].layer[0].SelfAttention.q.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{shard_grad}"
 
+    # check self attention embed
+    org_grad = org_model.encoder.block[0].layer[0].SelfAttention.relative_attention_bias.weight.grad
+    shard_grad = sharded_model.encoder.block[0].layer[0].SelfAttention.relative_attention_bias.weight.grad
+    shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)]
+    torch.distributed.all_gather(shard_grad_list, shard_grad)
+    all_shard_grad = torch.cat(shard_grad_list, dim=1)
+    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}"
+
+    # check token embedding grad
+    org_grad = org_model.shared.weight.grad
+
+    # check weights are tied
+    if hasattr(org_model, 'lm_head'):
+        assert org_model.shared.weight.data.data_ptr() == org_model.lm_head.weight.data.data_ptr()
+        assert sharded_model.shared.weight.data.data_ptr() == sharded_model.lm_head.weight.data.data_ptr()
+
+    shard_grad = sharded_model.shared.weight.grad
+    shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)]
+    torch.distributed.all_gather(shard_grad_list, shard_grad)
+    all_shard_grad = torch.cat(shard_grad_list, dim=0)
+    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}"
+
 
 def check_t5(rank, world_size, port):
     disable_existing_loggers()
@@ -44,7 +68,6 @@ def check_t5(rank, world_size, port):
     for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
         org_model, sharded_model = build_model(model_fn)
         check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn)
-
     torch.cuda.empty_cache()
 
 
@@ -56,4 +79,4 @@ def test_t5():
 
 
 if __name__ == "__main__":
-    test_t5()
+    test_t5()

tests/test_shardformer/test_model/test_shard_vit.py

@@ -45,6 +45,7 @@ def check_vit(rank, world_size, port):
 
 
 @pytest.mark.dist
+@pytest.mark.skip
 @rerun_if_address_is_in_use()
 @clear_cache_before_run()
 def test_vit():