add model.yaml save to context to io_dump

Signed-off-by: HuiyingLi <[email protected]>
HuiyingLi · Oct 21, 2024 · 4ea3505 · 4ea3505
1 parent 4c3de48
commit 4ea3505
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/checkpoint_converters/convert_nemo1_to_nemo2.py b/scripts/checkpoint_converters/convert_nemo1_to_nemo2.py
@@ -161,7 +161,7 @@ def skip_fp8_load(x):
         #Corresponding to Connector: on_import_ckpt
         if hasattr(trainer.model, "__io__") and hasattr(trainer.model.tokenizer, '__io__'):
             trainer.model.__io__.tokenizer = trainer.model.tokenizer.__io__
-        TrainerContext.from_trainer(trainer).io_dump(ckpt_to_context_subdir(args.output_path))
+        TrainerContext.from_trainer(trainer).io_dump(ckpt_to_context_subdir(args.output_path), yaml_attrs=["model"])
 
     #remove tmp dir
     if os.path.isdir(tokenizer_tmp_dir):