[QEff. Finetuning]: Tests for Pipeline Parallelism and updated documentation

QEfficient/finetune/experimental/core/config_manager.py

@@ -21,6 +21,7 @@
 from transformers.hf_argparser import HfArgumentParser
 
 from QEfficient.finetune.experimental.core.logger import Logger
+from QEfficient.finetune.experimental.core.utils import constants
 from QEfficient.finetune.experimental.core.utils.dist_utils import is_main_process
 from QEfficient.utils.device_utils import is_nsp_free
 
@@ -855,7 +856,7 @@ def get_model_config(self) -> Dict[str, Any]:
             training_dtype = training_config.get("torch_dtype")
             if training_dtype:
                 # Convert from training format (fp16/bf16) to model format (float16/bfloat16)
-                dtype_mapping = {"fp16": "float16", "bf16": "bfloat16"}
+                dtype_mapping = dtype_mapping = constants.DTYPE_MAPPING
                 model_config["torch_dtype"] = dtype_mapping.get(training_dtype, "auto")
 
         return model_config

QEfficient/finetune/experimental/core/model.py

@@ -105,16 +105,6 @@ def _resolve_auto_class(auto_class_name: str) -> Type:
             )
         return getattr(transformers, auto_class_name)
 
-    # def _build_quant_config(self) -> Optional[BitsAndBytesConfig]:
-    #     if not self.model_kwargs.get("load_in_4bit"):
-    #         return None
-    #     return BitsAndBytesConfig(
-    #         load_in_4bit=True,
-    #         bnb_4bit_quant_type=self.model_kwargs.get("bnb_4bit_quant_type", "nf4"),
-    #         bnb_4bit_compute_dtype=self.model_kwargs.get("bnb_4bit_compute_dtype", torch.float16),
-    #         bnb_4bit_use_double_quant=self.model_kwargs.get("bnb_4bit_use_double_quant", True),
-    #     )
-
     def configure_model_kwargs(self) -> Dict[str, Any]:
         """Hook for subclasses to tweak HF `.from_pretrained` kwargs."""
 

QEfficient/finetune/experimental/core/utils/constants.py

@@ -0,0 +1 @@
+DTYPE_MAPPING = {"fp16": "float16", "bf16": "bfloat16"}