Merge pull request THUDM#31 from THUDM/develop

Develop
gg-big-org · Dec 9, 2022 · afdee38 · afdee38
2 parents 564ee2d + 41b48c2
commit afdee38
Show file tree

Hide file tree

Showing 9 changed files with 1,782 additions and 1 deletion.
diff --git a/codegeex/megatron/tools/pretrain_codegeex.py b/codegeex/megatron/tools/pretrain_codegeex.py
@@ -70,7 +70,7 @@ def model_provider(pre_process=True, post_process=True):
                 mp_rank = mpu.get_tensor_model_parallel_rank()
                 if os.path.isdir(args.load_state):
                     model_path = os.path.join(
-                        args.load_state, f"model_mp_rank_{mp_rank}.pt"
+                        args.load_state, "mp_rank_{:02d}_model_states.pt".format(mp_rank)
                     )
                 else:
                     model_path = args.load_state

diff --git a/codegeex/paddle/__init__.py b/codegeex/paddle/__init__.py
@@ -0,0 +1 @@
+from .codegeex_model import CodeGeeXModel