From a62ee55a81d794a7ee1abe8d82498d7cc76f529e Mon Sep 17 00:00:00 2001
From: Saaketh Narayan <narayan.saaketh@gmail.com>
Date: Wed, 14 Feb 2024 12:39:15 -0800
Subject: [PATCH] Make turbo an optional dependency (#964)

---
 llmfoundry/optim/lion8b.py | 14 ++++++++++++--
 setup.py                   |  7 +++++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/llmfoundry/optim/lion8b.py b/llmfoundry/optim/lion8b.py
index f76d29b1c7..0026fcef06 100644
--- a/llmfoundry/optim/lion8b.py
+++ b/llmfoundry/optim/lion8b.py
@@ -235,7 +235,12 @@ def __init__(self, data: Optional[torch.Tensor], try_quantize: bool = True):
         self._f_encode = None
         self._f_decode = None
         if self._try_quantize:
-            from turbo import dequantize_signed, quantize_signed
+            try:
+                from turbo import dequantize_signed, quantize_signed
+            except ModuleNotFoundError:
+                raise NotImplementedError(
+                    'The Lion 8b optimizer requires installing mosaicml-turbo. ',
+                    'Please `pip install llm-foundry[turbo]` to install it.')
             self._f_encode = quantize_signed
             self._f_decode = dequantize_signed
 
@@ -396,7 +401,12 @@ def lion8b_step_fused(grads: torch.Tensor,
                 f'Weights must be f32 or match grad dtype {grads.dtype}')
 
     # ------------------------------------------------ actual function call
-    from turbo import lion8b_step_cuda
+    try:
+        from turbo import lion8b_step_cuda
+    except ModuleNotFoundError:
+        raise NotImplementedError(
+            'The Lion 8b optimizer requires installing mosaicml-turbo. ',
+            'Please `pip install llm-foundry[turbo]` to install it.')
     return lion8b_step_cuda(grads=grads,
                             weights=weights,
                             momentums=momentums,
diff --git a/setup.py b/setup.py
index 8fdfc2197b..3dc838a7ab 100644
--- a/setup.py
+++ b/setup.py
@@ -100,13 +100,16 @@
 
 extra_deps['gpu'] = [
     'flash-attn==1.0.9',
-    'mosaicml-turbo==0.0.8',
     # PyPI does not support direct dependencies, so we remove this line before uploading from PyPI
     'xentropy-cuda-lib@git+https://github.com/HazyResearch/flash-attention.git@v1.0.9#subdirectory=csrc/xentropy',
 ]
+
+extra_deps['turbo'] = [
+    'mosaicml-turbo==0.0.8',
+]
+
 extra_deps['gpu-flash2'] = [
     'flash-attn==2.5.0',
-    'mosaicml-turbo==0.0.8',
 ]
 
 extra_deps['peft'] = [