diff --git a/tests/acceleration/test_acceleration_framework.py b/tests/acceleration/test_acceleration_framework.py index 0582782cb..b925f216a 100644 --- a/tests/acceleration/test_acceleration_framework.py +++ b/tests/acceleration/test_acceleration_framework.py @@ -434,7 +434,7 @@ def test_framework_intialized_properly_foak(): fusedops_kernels_config = FusedOpsAndKernelsConfig( fused_lora=FusedLoraConfig(base_layer="auto_gptq", fused_lora=True), fast_kernels=FastKernelsConfig( - fast_loss=True, fast_rsm_layernorm=True, fast_rope_embeddings=True + fast_loss=True, fast_rms_layernorm=True, fast_rope_embeddings=True ), ) diff --git a/tuning/config/acceleration_configs/acceleration_framework_config.py b/tuning/config/acceleration_configs/acceleration_framework_config.py index 75f521097..5eb7d3df2 100644 --- a/tuning/config/acceleration_configs/acceleration_framework_config.py +++ b/tuning/config/acceleration_configs/acceleration_framework_config.py @@ -129,12 +129,10 @@ def _verify_configured_dataclasses(self): # Check that fused lora must be activated with either auto_gptq or bitsandbytes if self.fused_lora is not None: - assert ( - self.bitsandbytes is not None or self.auto_gptq is not None - ), "`--fused_lora` must be accompanied by a quantized base layer"\ + assert self.bitsandbytes is not None or self.auto_gptq is not None, ( + "`--fused_lora` must be accompanied by a quantized base layer" " `--auto_gptq` or `--bitsandbytes`." - - + ) @staticmethod def from_dataclasses(*dataclasses: Type): diff --git a/tuning/config/acceleration_configs/fused_ops_and_kernels.py b/tuning/config/acceleration_configs/fused_ops_and_kernels.py index 4777394c8..a8a66a9a4 100644 --- a/tuning/config/acceleration_configs/fused_ops_and_kernels.py +++ b/tuning/config/acceleration_configs/fused_ops_and_kernels.py @@ -70,6 +70,7 @@ class FusedOpsAndKernelsConfig: # fast kernels fast_kernels: FastKernelsConfig = None + def __post_init__(self): # ensure nested dataclasses initialized ensure_nested_dataclasses_initialized(self)