Cleanup

optimum/neuron/distributed/leaf_functions_for_pipeline_parallelism.py

@@ -0,0 +1,18 @@
+import torch
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+
+
+leaf_prepare_4d_causal_attention_mask = torch.fx._symbolic_trace._create_wrapped_func(
+    _prepare_4d_causal_attention_mask
+)
+
+leaf_prepare_4d_causal_attention_mask = torch.fx._symbolic_trace._create_wrapped_func(
+    _prepare_4d_causal_attention_mask
+)
+
+leaf_prepare_4d_causal_attention_maskfor_sdpa = torch.fx._symbolic_trace._create_wrapped_func(
+    _prepare_4d_causal_attention_mask_for_sdpa
+)

optimum/neuron/utils/neuron_cc_wrapper

@@ -0,0 +1,9 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from optimum.neuron.utils.optimum_neuron_cc_wrapper import main
+
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

optimum/neuron/utils/optimum_neuron_cc_wrapper.py

@@ -0,0 +1,28 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from libneuronxla.neuron_cc_wrapper import main as neuron_cc_wrapper_main
+
+from .cache_utils import get_hf_hub_cache_repos
+from .hub_neuronx_cache import hub_neuronx_cache
+
+
+def main():
+    with hub_neuronx_cache(cache_repo_id=get_hf_hub_cache_repos()[0]):
+        return neuron_cc_wrapper_main()
+
+
+if __name__ == "__main__":
+    main()

tests/distributed/distributed.py

@@ -97,8 +97,7 @@ class DistributedExec(ABC):
     exec_timeout: int = TEST_TIMEOUT
 
     @abstractmethod
-    def run(self):
-        ...
+    def run(self): ...
 
     def __call__(self, request=None):
         self._fixture_kwargs = self._get_fixture_kwargs(request, self.run)

tests/test_cache_utils.py

@@ -83,9 +83,9 @@ def test_get_neuron_cache_path(self):
         assert get_neuron_cache_path() is None
 
         custom_cache_dir_name = Path("_this/is_/my1/2custom/cache/dir")
-        os.environ[
-            "NEURON_CC_FLAGS"
-        ] = f"--some --parameters --here --cache_dir={custom_cache_dir_name} --other --paremeters --here"
+        os.environ["NEURON_CC_FLAGS"] = (
+            f"--some --parameters --here --cache_dir={custom_cache_dir_name} --other --paremeters --here"
+        )
 
         self.assertEqual(get_neuron_cache_path(), custom_cache_dir_name)
 
@@ -99,9 +99,9 @@ def _test_set_neuron_cache_path(self, new_cache_path):
         set_neuron_cache_path(new_cache_path, ignore_no_cache=True)
         self.assertEqual(get_neuron_cache_path(), Path(new_cache_path))
 
-        os.environ[
-            "NEURON_CC_FLAGS"
-        ] = "--some --parameters --here --cache_dir=original_cache_dir --other --paremeters"
+        os.environ["NEURON_CC_FLAGS"] = (
+            "--some --parameters --here --cache_dir=original_cache_dir --other --paremeters"
+        )
         set_neuron_cache_path(new_cache_path)
         self.assertEqual(get_neuron_cache_path(), Path(new_cache_path))