fix: smaller batch for llama in FHE

zama-ai · Jan 5, 2025 · 6f6cfb9 · 6f6cfb9
1 parent a4a248e
commit 6f6cfb9
Show file tree

Hide file tree

Showing 3 changed files with 153 additions and 1,955 deletions.
diff --git a/src/concrete/ml/torch/hybrid_model.py b/src/concrete/ml/torch/hybrid_model.py
@@ -18,6 +18,7 @@
 from brevitas.quant_tensor import QuantTensor
 from concrete.fhe import Configuration
 from torch import nn
+from tqdm.autonotebook import tqdm
 
 from ..common.utils import MAX_BITWIDTH_BACKWARD_COMPATIBLE, HybridFHEMode
 from ..deployment.fhe_client_server import FHEModelClient, FHEModelDev, FHEModelServer
@@ -577,8 +578,6 @@ def compile_model(
 
         self.configuration = configuration
 
-        from tqdm import tqdm
-
         for name in tqdm(self.module_names, desc="Compiling FHE layers"):
             remote_module = self._get_module_by_name(self.model, name)
             assert isinstance(remote_module, RemoteModule)

diff --git a/src/concrete/ml/torch/lora.py b/src/concrete/ml/torch/lora.py
@@ -6,7 +6,7 @@
 import torch
 from torch import Tensor, nn
 from torch.utils.data import DataLoader
-from tqdm import tqdm
+from tqdm.autonotebook import tqdm
 
 from ..common.utils import assert_true
 from .hybrid_backprop_linear import CustomLinear