additional fmt fixes

Signed-off-by: 1000850000 user <aaron.chew1@ibm.com>
foundation-model-stack · Aug 29, 2024 · 29362a4 · 29362a4
1 parent 9cb999e
commit 29362a4
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/tests/acceleration/test_acceleration_dataclasses.py b/tests/acceleration/test_acceleration_dataclasses.py
@@ -25,8 +25,8 @@
 )
 from tuning.config.acceleration_configs.attention_and_distributed_packing import (
     AttentionAndDistributedPackingConfig,
-    PaddingFree,
     MultiPack,
+    PaddingFree,
 )
 from tuning.config.acceleration_configs.fused_ops_and_kernels import (
     FastKernelsConfig,

diff --git a/tests/acceleration/test_acceleration_framework.py b/tests/acceleration/test_acceleration_framework.py
@@ -40,8 +40,8 @@
 )
 from tuning.config.acceleration_configs.attention_and_distributed_packing import (
     AttentionAndDistributedPackingConfig,
-    PaddingFree,
     MultiPack,
+    PaddingFree,
 )
 from tuning.config.acceleration_configs.fused_ops_and_kernels import (
     FastKernelsConfig,
@@ -514,7 +514,8 @@ def test_framework_initialize_and_trains_with_aadp():
                     model_args,
                     data_args,
                     train_args,
-                    attention_and_distributed_packing_config=attention_and_distributed_packing_config,
+                    attention_and_distributed_packing_config=\
+                        attention_and_distributed_packing_config,
                 )
 
         # spy inside the train to ensure that the ilab plugin is called
@@ -573,7 +574,8 @@ def test_padding_free_plugin_raises_error_with_untokenized_dataset():
                         model_args,
                         data_args,
                         train_args,
-                        attention_and_distributed_packing_config=attention_and_distributed_packing_config,
+                        attention_and_distributed_packing_config=\
+                            attention_and_distributed_packing_config,
                     )
 
 
@@ -596,6 +598,7 @@ def test_error_raised_with_paddingfree_and_flash_attn_disabled():
             attention_and_distributed_packing_config=attention_and_distributed_packing_config,
         )
 
+
 def test_error_raised_with_multipack_and_paddingfree_disabled():
     """Ensure error raised when padding-free is not used with flash attention"""
     with pytest.raises(

diff --git a/tuning/config/acceleration_configs/acceleration_framework_config.py b/tuning/config/acceleration_configs/acceleration_framework_config.py
@@ -123,7 +123,9 @@ def _verify_configured_dataclasses(self):
             # this also ensures that the attention implementation for multipack
             # will be flash attention as sfttrainer will enforce flash attn to be
             # set for padding free
-            assert self.padding_free is not None, "`--multipack` is currently only supported with `--padding_free`"
+            assert (
+                self.padding_free is not None
+            ), "`--multipack` is currently only supported with `--padding_free`"
 
     @staticmethod
     def from_dataclasses(*dataclasses: Type):