[shardformer] Pytree fix (#4533)

* pytree test * test bert * test bert * test bert * revise * add register * add register
2025-09-16 06:30:41 +00:00 · 2023-09-04 17:52:23 +08:00
parent 508ca36fe3
commit 24c0768795
5 changed files with 81 additions and 17 deletions
--- a/tests/test_shardformer/test_model/test_shard_bert.py
+++ b/tests/test_shardformer/test_model/test_shard_bert.py
@@ -179,6 +179,7 @@ def run_bert_3d_test(test_config):
    sub_model_zoo = model_zoo.get_sub_registry('transformers_bert')

    for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
+
        check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config)

    clear_layout_converter()