diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..aa8086d --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,37 @@ +name: CI + +on: + pull_request: + branches: [main] + push: + branches: [main] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + lint: + if: github.repository == 'apple/coreai-torch' + runs-on: [self-hosted, macos, tahoe, ARM64] + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6 + - name: Ensure uv + run: | + command -v uv >/dev/null 2>&1 || curl -LsSf https://astral.sh/uv/install.sh | sh + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + - run: uv run --extra dev ruff check . + - run: uv run --extra dev ruff format --check . + + python-test: + if: github.repository == 'apple/coreai-torch' + runs-on: [self-hosted, macos, tahoe, ARM64] + timeout-minutes: 60 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6 + - name: Ensure uv + run: | + command -v uv >/dev/null 2>&1 || curl -LsSf https://astral.sh/uv/install.sh | sh + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + - run: uv run --extra test pytest tests/ -n auto -m "not slow" diff --git a/coreai_torch/_compression/utils.py b/coreai_torch/_compression/utils.py index 13bcd3d..be371fe 100644 --- a/coreai_torch/_compression/utils.py +++ b/coreai_torch/_compression/utils.py @@ -99,6 +99,7 @@ def repeat_tensor_as(tensor: torch.Tensor, target_shape: torch.Size) -> torch.Te ) return repeated_tensor + def wrap_for_parametrization( compression_module_class: type[torch.nn.Module], ) -> type[torch.nn.Module]: diff --git a/docs/coreai-core/tutorials/construct-a-graph.ipynb b/docs/coreai-core/tutorials/construct-a-graph.ipynb index 107dd6b..a10b05d 100644 --- a/docs/coreai-core/tutorials/construct-a-graph.ipynb +++ b/docs/coreai-core/tutorials/construct-a-graph.ipynb @@ -76,11 +76,10 @@ "\n", "import numpy as np\n", "\n", - "from coreai.authoring import AIModelAsset, AIProgram, Module, TensorSpec\n", - "\n", "# Graph-building primitives — pending re-export from coreai.authoring.\n", "from coreai._compiler.dialects import coreai as ops\n", - "from coreai._compiler.ir import Value" + "from coreai._compiler.ir import Value\n", + "from coreai.authoring import AIModelAsset, AIProgram, Module, TensorSpec" ] }, { @@ -175,6 +174,7 @@ " ) -> Annotated[Value, output_spec]:\n", " return ops.add(x, x)\n", "\n", + "\n", "module.verify()\n", "print(\"Module verified.\")" ] diff --git a/docs/coreai-core/tutorials/run-an-aimodel.ipynb b/docs/coreai-core/tutorials/run-an-aimodel.ipynb index 02f2c91..4d878f9 100644 --- a/docs/coreai-core/tutorials/run-an-aimodel.ipynb +++ b/docs/coreai-core/tutorials/run-an-aimodel.ipynb @@ -64,7 +64,6 @@ "from pathlib import Path\n", "\n", "import numpy as np\n", - "\n", "from coreai.authoring import AIModelAsset\n", "from coreai.runtime import InferenceFunction, NDArray\n", "\n", @@ -108,10 +107,10 @@ "from shutil import rmtree\n", "from typing import Annotated\n", "\n", - "from coreai.authoring import AIProgram, Module, TensorSpec\n", "# Pending re-export from coreai.authoring; see the previous tutorial.\n", "from coreai._compiler.dialects import coreai as ops\n", "from coreai._compiler.ir import Value\n", + "from coreai.authoring import AIProgram, Module, TensorSpec\n", "\n", "# Reconstruct asset.\n", "if asset_path.exists():\n", @@ -126,6 +125,7 @@ " ) -> Annotated[Value, TensorSpec(shape=[2, 3], dtype=np.float32, name=\"y\")]:\n", " return ops.add(x, x)\n", "\n", + "\n", "AIProgram(module).save_asset(asset_path)\n", "print(f\"created {asset_path}\")" ] diff --git a/docs/getting-started/quickstart.ipynb b/docs/getting-started/quickstart.ipynb index 42687cc..b282217 100644 --- a/docs/getting-started/quickstart.ipynb +++ b/docs/getting-started/quickstart.ipynb @@ -68,8 +68,6 @@ "metadata": {}, "outputs": [], "source": [ - "import torch\n", - "\n", "example_input = (torch.randn(1, 10),)\n", "exported = torch.export.export(model, args=example_input)" ] @@ -145,10 +143,10 @@ "outputs": [], "source": [ "import tempfile\n", - "import torch\n", "from pathlib import Path\n", "\n", "import numpy as np\n", + "import torch\n", "from coreai.runtime import NDArray\n", "\n", "\n", @@ -240,9 +238,6 @@ "metadata": {}, "outputs": [], "source": [ - "import tempfile\n", - "\n", - "\n", "async def run():\n", " with tempfile.TemporaryDirectory() as tmpdir:\n", " asset = coreai_program.save_asset(Path(tmpdir) / \"mobilenet_v2_example.aimodel\")\n", @@ -274,6 +269,7 @@ "outputs": [], "source": [ "import torch\n", + "\n", "import coreai_torch\n", "\n", "model = SimpleModel().eval()\n", diff --git a/docs/guides/composite-ops.ipynb b/docs/guides/composite-ops.ipynb index 0face7c..df93289 100644 --- a/docs/guides/composite-ops.ipynb +++ b/docs/guides/composite-ops.ipynb @@ -97,11 +97,6 @@ "metadata": {}, "outputs": [], "source": [ - "import torch\n", - "import torch.nn as nn\n", - "from coreai_torch.composite_ops import RMSNormImpl\n", - "\n", - "\n", "class RMSNorm(nn.Module):\n", " \"\"\"Convenience wrapper that owns the learnable scale parameter.\"\"\"\n", "\n", @@ -149,7 +144,6 @@ "outputs": [], "source": [ "import torch\n", - "import coreai_torch\n", "\n", "coreai_program = (\n", " TorchConverter()\n", @@ -185,10 +179,10 @@ "outputs": [], "source": [ "import tempfile\n", - "import torch\n", "from pathlib import Path\n", "\n", "import numpy as np\n", + "import torch\n", "from coreai.runtime import NDArray\n", "\n", "\n", diff --git a/docs/guides/conversion-workflows.ipynb b/docs/guides/conversion-workflows.ipynb index 88faff8..72a3144 100644 --- a/docs/guides/conversion-workflows.ipynb +++ b/docs/guides/conversion-workflows.ipynb @@ -86,6 +86,7 @@ "outputs": [], "source": [ "import torch\n", + "\n", "import coreai_torch\n", "\n", "model = MyModel().eval()\n", @@ -133,7 +134,7 @@ "source": [ "import torch\n", "import torch.nn as nn\n", - "import coreai_torch\n", + "\n", "from coreai_torch import ExternalizeSpec, TorchConverter\n", "from coreai_torch.composite_ops import RMSNormImpl\n", "\n", diff --git a/docs/guides/custom-metal-kernels.ipynb b/docs/guides/custom-metal-kernels.ipynb index 6bca4e3..baaaecf 100644 --- a/docs/guides/custom-metal-kernels.ipynb +++ b/docs/guides/custom-metal-kernels.ipynb @@ -122,6 +122,7 @@ "outputs": [], "source": [ "import torch\n", + "\n", "from coreai_torch import get_decomp_table\n", "\n", "model = AddModel().eval()\n", @@ -182,6 +183,7 @@ "source": [ "import torch\n", "\n", + "\n", "def torch_matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n", " return torch.matmul(x, y)\n", "\n", @@ -222,6 +224,7 @@ "import torch\n", "import torch.nn as nn\n", "\n", + "\n", "def torch_sincos(x: torch.Tensor) -> list[torch.Tensor]:\n", " return [torch.sin(x), torch.cos(x)]\n", "\n", diff --git a/docs/guides/custom-op-lowering.ipynb b/docs/guides/custom-op-lowering.ipynb index 396e215..be468b4 100644 --- a/docs/guides/custom-op-lowering.ipynb +++ b/docs/guides/custom-op-lowering.ipynb @@ -92,6 +92,7 @@ "outputs": [], "source": [ "import torch\n", + "\n", "from coreai_torch import get_decomp_table\n", "\n", "model = ScaledAddModel().eval()\n", @@ -198,9 +199,9 @@ "metadata": {}, "outputs": [], "source": [ + "import numpy as np\n", "import torch\n", "import torch.nn as nn\n", - "import numpy as np\n", "\n", "from coreai_torch._utils import get_operand\n", "\n", diff --git a/docs/guides/externalization.ipynb b/docs/guides/externalization.ipynb index 6a4fea5..931e749 100644 --- a/docs/guides/externalization.ipynb +++ b/docs/guides/externalization.ipynb @@ -86,8 +86,6 @@ "import torch\n", "import torch.nn as nn\n", "\n", - "from coreai_torch import ExternalizeSpec\n", - "\n", "\n", "class RMSNormComposite(nn.Module):\n", " def __init__(self, axes=-1, eps=1e-5, version=1):\n", @@ -101,6 +99,7 @@ " inv_rms = torch.rsqrt((x_f32 * x_f32).mean(self.axes, keepdim=True) + self.eps)\n", " return (input * inv_rms).to(input.dtype) * scale\n", "\n", + "\n", "model = RMSNormComposite().eval()\n", "sample = (torch.randn(10), torch.randn(10))" ] diff --git a/tests/compression/test_compression.py b/tests/compression/test_compression.py index 4064733..a626a64 100644 --- a/tests/compression/test_compression.py +++ b/tests/compression/test_compression.py @@ -14,7 +14,6 @@ import torch from coreai.authoring import AIProgram from coreai.runtime import AIModel, NDArray - from torch import nn from torch.export.exported_program import ExportedProgram @@ -30,6 +29,7 @@ # We add "./tests/coreai" path, in order to use some existing utils sys.path.append(str(Path(__file__).parents[2])) + def _scale_shape( input_shape: tuple[int, ...], axis: int, @@ -57,6 +57,7 @@ async def lower_to_coreai( converter = TorchConverter().add_exported_program(coreaten_program) return converter.to_coreai() + async def _validate_execution( coreai_program: AIProgram, torch_out: torch.Tensor, @@ -82,6 +83,7 @@ async def _validate_execution( atol=atol, ) + @pytest.mark.parametrize( "nbits", [4, 8], @@ -1056,9 +1058,7 @@ def test_ir( "" if granularity == "per_tensor" else f"{activation_shape[axis]}x" ) - msg = ( - "TODO: reshape on consts such as offset and scale is not const eliminated" - ) + msg = "TODO: reshape on consts such as offset and scale is not const eliminated" pytest.xfail(reason=msg) truth = f""" // CHECK-LABEL: coreai.graph @main diff --git a/tests/test_externalize.py b/tests/test_externalize.py index 3a1f4ff..fe8a74e 100644 --- a/tests/test_externalize.py +++ b/tests/test_externalize.py @@ -3032,10 +3032,12 @@ def forward( @pytest.mark.ir @pytest.mark.flaky(reruns=3) -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) def test_externalize_rms_norm_with_quantized_linears_ir() -> None: """IR check: Quantized (int4) weights retain si4 dtype when externalization re-exports the model.""" @@ -3065,7 +3067,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 sample = (torch.randn(2, 24),) @@ -3109,10 +3111,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: @pytest.mark.flaky(reruns=3) -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) async def test_externalize_rms_norm_with_quantized_linears() -> None: """Quantized (int4) weights retain si4 dtype when externalization re-exports the model. @@ -3153,7 +3157,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 sample = (torch.randn(2, 24),) @@ -3176,10 +3180,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: @pytest.mark.ir @pytest.mark.flaky(reruns=3) -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) def test_externalize_gather_mm_with_quantized_rhs_ir() -> None: """IR check: Quantized expert weight flows as rhs into an externalized GatherMM composite.""" num_experts = 4 @@ -3217,7 +3223,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 x = torch.randn(2, 1, 1, in_dim) indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16) @@ -3263,10 +3269,12 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: @pytest.mark.flaky(reruns=3) -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) async def test_externalize_gather_mm_with_quantized_rhs() -> None: """Quantized expert weight flows as rhs into an externalized GatherMM composite. @@ -3318,7 +3326,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 x = torch.randn(2, 1, 1, in_dim) indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16) @@ -3345,10 +3353,12 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: @pytest.mark.ir @pytest.mark.flaky(reruns=3) -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) def test_externalize_multiple_composites_with_quantized_weights_ir() -> None: """IR check: Multiple composite ops (RMSNorm + SDPA) externalized with quantized linears.""" head_dim = 16 @@ -3390,7 +3400,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 sample = (torch.randn(1, 4, embed_dim),) @@ -3433,10 +3443,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: @pytest.mark.flaky(reruns=3) -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) async def test_externalize_multiple_composites_with_quantized_weights() -> None: """Multiple composite ops (RMSNorm + SDPA) externalized with quantized linears. @@ -3484,7 +3496,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 sample = (torch.randn(1, 4, embed_dim),) @@ -3512,10 +3524,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: @pytest.mark.ir -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) def test_externalize_gather_mm_combined_with_rms_norm_ir() -> None: """IR check: GatherMM + RMSNorm both externalized alongside quantized weights.""" num_experts = 4 @@ -3553,7 +3567,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 x = torch.randn(2, in_dim) indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16) @@ -3590,10 +3604,13 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: """ filecheck_pattern(ir, check_file=pattern) -@pytest.mark.skip(reason=( - "transform_with_custom_compression_ops has been deprecated. Consider removing " - "these tests or use an alternative way to generate quantized weights" -)) + +@pytest.mark.skip( + reason=( + "transform_with_custom_compression_ops has been deprecated. Consider removing " + "these tests or use an alternative way to generate quantized weights" + ) +) async def test_externalize_gather_mm_combined_with_rms_norm() -> None: """GatherMM + RMSNorm both externalized alongside quantized weights. @@ -3637,7 +3654,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: ) quantizer = PostTrainingQuantizer(model, quantization_config) model = cast("nn.Module", quantizer.compress()) - transform_with_custom_compression_ops(model) + transform_with_custom_compression_ops(model) # noqa: F821 x = torch.randn(2, in_dim) indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16)