Fix checkpoints for Gaudi

remove _orig_mod prefix from checkpoints from torch.compile trained
model

Signed-off-by: Jianhong-Zhang <[email protected]>

Author: Jianhong-Zhang

src/instructlab/training/hpu_utils.py

@@ -1,3 +1,4 @@
+import os
 import torch
 from functools import lru_cache
 
@@ -47,3 +48,15 @@ def simple_bucket(length):
 
 def bucket(length):
     return simple_bucket(length)
+
+
+def save_hpu_model(model, output_dir):
+    from safetensors.torch import save_file
+
+    state_dict = model.state_dict()
+    remove_prefix = "_orig_mod."
+    clean_state_dict = {
+        k[len(remove_prefix) :] if k.startswith(remove_prefix) else k: v
+        for k, v in state_dict.items()
+    }
+    save_file(clean_state_dict, os.path.join(output_dir, "model.safetensors"))

src/instructlab/training/utils.py

@@ -50,7 +50,11 @@
     QuantizeDataType,
     TrainingArgs,
 )
-from instructlab.training.hpu_utils import is_torch_hpu_available, bucket
+from instructlab.training.hpu_utils import (
+    is_torch_hpu_available,
+    bucket,
+    save_hpu_model,
+)
 
 logger = logging.getLogger("instructlab.training")
 
@@ -1033,12 +1037,15 @@ def _get_state_dict_patched(model, unwrap=False):
             model.module.unmerge_adapter()
 
     if not is_lora:
-        accelerator.save_model(
-            model,
-            save_directory=output_dir,
-            max_shard_size="5GB",
-            safe_serialization=True,
-        )
+        if is_torch_hpu_available() and os.getenv("HPU_ENABLE_TORCH_COMPILE", False):
+            save_hpu_model(model, output_dir)
+        else:
+            accelerator.save_model(
+                model,
+                save_directory=output_dir,
+                max_shard_size="5GB",
+                safe_serialization=True,
+            )
 
     if args.use_dolomite and convert_dolomite and accelerator.is_main_process:
         # export doesnt like the directory to exist