switch instructlab dolomite

Signed-off-by: Yu Chin Fabian Lim <[email protected]>

src/instructlab/training/main_ds.py

@@ -11,6 +11,9 @@
 # Third Party
 from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
 from deepspeed.runtime.zero.utils import ZeRORuntimeException
+from instructlab.dolomite.enums import GradientCheckpointingMethod
+from instructlab.dolomite.gradient_checkpointing import apply_gradient_checkpointing
+from instructlab.dolomite.hf_models import GPTDolomiteForCausalLM
 from torch.distributed import ReduceOp, all_reduce
 from tqdm import tqdm
 from transformers import AutoModelForCausalLM, get_scheduler
@@ -89,7 +92,7 @@ def setup_model(args, tokenizer, train_loader, grad_accum):
 
     if args.is_granite:
         # Third Party
-        from dolomite_engine.hf_models.models import GPTDolomiteForCausalLM
+        # from dolomite_engine.hf_models.models import GPTDolomiteForCausalLM
 
         model = GPTDolomiteForCausalLM.from_pretrained(
             args.model_name_or_path,
@@ -202,8 +205,8 @@ def setup_model(args, tokenizer, train_loader, grad_accum):
     # for both lora and full here
     if args.is_granite:
         # Third Party
-        from dolomite_engine.enums import GradientCheckpointingMethod
-        from dolomite_engine.gradient_checkpointing import apply_gradient_checkpointing
+        # from dolomite_engine.enums import GradientCheckpointingMethod
+        # from dolomite_engine.gradient_checkpointing import apply_gradient_checkpointing
 
         block_name = model._no_split_modules[0]
         apply_gradient_checkpointing(

src/instructlab/training/utils.py

@@ -13,6 +13,7 @@
 import warnings
 
 # Third Party
+from instructlab.dolomite.hf_models import export_to_huggingface
 from rich.logging import RichHandler
 from torch import distributed as dist
 from torch.distributed import get_rank, is_initialized
@@ -539,7 +540,6 @@ def save_hf_format_ds(args, model, tokenizer, samples_seen, convert_granite=True
             from tempfile import TemporaryDirectory
 
             # Third Party
-            from dolomite_engine.hf_models import export_to_huggingface
             from safetensors.torch import save_file
 
             with TemporaryDirectory("w") as tmpdir: