Updated spec with Tanner feedback

Author: gilbertmike

accelforge/frontend/workload.py

@@ -3,13 +3,11 @@
 """
 
 import copy
-from itertools import product
-import itertools
-import logging
 import re
-from typing import Annotated, Any, TypeAlias
+from typing import Annotated, Any, TypeAlias, Union
 
 import pydot
+from pydantic import Discriminator, Tag
 
 from accelforge.util.parallel import _SVGJupyterRender
 
@@ -19,6 +17,7 @@
     EvalableList,
     EvalableModel,
     EvalsTo,
+    _get_tag,
 )
 from accelforge.util._visualization import _pydot_graph
 from accelforge.frontend.renames import (
@@ -833,22 +832,29 @@ class Adapter(EvalableModel):
     """
     name: EvalsTo[str]
 
-    adapter_class: EvalsTo[str]
-    """
-    The class of the adapter. Supported classes are:
-    - copy
-    - collective
-    """
 
+class CopyAdapter(Adapter):
     tensor: EvalsTo[str]
 
 
+EinsumOrAdapter = Annotated[
+    Union[
+        Annotated["Einsum", Tag("Einsum")],
+        Annotated["CopyAdapter", Tag("Copy")],
+    ],
+    Discriminator(lambda v: _get_tag(v, default="Einsum")),
+]
+
+
 class Workload(EvalableModel):
     """
     The workload specification as a cascade of Einsums, with each Einsum being a
     computation step in the workload.
     """
 
+    einsums_and_adapters: EvalableList[EinsumOrAdapter] = EvalableList()
+    """Einsums and adapters in the workload."""
+
     einsums: EvalableList[Einsum] = EvalableList()
     """ The Einsums in the workload. """
 
@@ -891,11 +897,6 @@ class Workload(EvalableModel):
     matching tensors as persistent. Example: "weight" or "~(Outputs | Intermediates)".
     """
 
-    adapters: EvalableList[Adapter] = EvalableList()
-    """
-    Adapters for shared tensors.
-    """
-
     def _for_einsum(self, einsum_name: EinsumName) -> "Workload":
         """Return a copy of the workload with only the Einsum with the given name."""
         new = self.model_copy(deep=False)
@@ -913,6 +914,8 @@ def __init__(self, **data):
                 processed_einsums.append(_parse_einsum_entry(einsum_entry))
             data["einsums"] = processed_einsums
 
+        self._canonicalize_einsums_and_adapters(data)
+
         super().__init__(**data)
 
     def model_post_init(self, __context__=None) -> None:
@@ -953,6 +956,26 @@ def _validate(self):
                         )}"
                     )
 
+    def _canonicalize_einsums_and_adapters(self, data):
+        """
+        Checks that only one of einsums or einsums_and_adapters is specified
+        and fills in the other.
+        """
+        has_einsums = "einsums" in data and len(data["einsums"]) > 0
+        has_einsums_and_adapters = "einsums_and_adapters" in data and len(data["einsums_and_adapters"]) > 0
+
+        if has_einsums and has_einsums_and_adapters:
+            raise EvaluationError("Please use only one of keywords einsums or einsums_and_adapters")
+
+        def is_einsum(d):
+            return not hasattr(d, "tag") or d.tag == "!Einsum"
+
+        if has_einsums:
+            data["einsums_and_adapters"] = data["einsums"]
+            data["einsums"] = [e for e in data["einsums_and_adapters"] if is_einsum(e)]
+        elif has_einsums_and_adapters:
+            data["einsums"] = [e for e in data["einsums_and_adapters"] if is_einsum(e)]
+
     @property
     def einsum_names(self) -> list[EinsumName]:
         """Returns the names of the Einsums in the workload."""
@@ -1334,9 +1357,3 @@ def get_compute_intensity(self, einsum_name: str) -> float:
             self.get_tensor_size(tensor)
             for tensor in self.einsums[einsum_name].tensor_names
         )
-
-    def get_adapter_for(self, tensor_name: TensorName) -> Adapter:
-        for adapter in self.adapters:
-            if adapter.tensor == tensor_name:
-                return adapter
-        return ValueError(f"No adapter found for tensor {tensor_name}")

accelforge/util/_basetypes.py

@@ -4,6 +4,7 @@
 import os
 from pathlib import Path
 import re
+import ruamel.yaml
 from pydantic import BaseModel, ConfigDict, Tag, ValidationError
 from pydantic.main import IncEx
 from pydantic_core.core_schema import (
@@ -59,7 +60,7 @@
 Ts = TypeVarTuple("Ts")
 
 
-def _get_tag(value: Any) -> str:
+def _get_tag(value: Any, default: str=None) -> str:
     if not isinstance(value, dict):
         return value.__class__.__name__
     tag = None
@@ -81,14 +82,22 @@ def try_index(attr: str) -> str:
             break
         if tag := try_index(attr):
             break
+
+    def is_yaml_tag_none(tag):
+        return isinstance(tag, ruamel.yaml.tag.Tag) and tag.value is None
+
+    if (tag is None or is_yaml_tag_none(tag)) and not default is None:
+        tag = default
+
     if tag is None:
         raise ValueError(
             f"No tag found for {value}. Either set the type field " "or use a YAML tag."
         )
-    tag = str(tag)
-    if tag.startswith("!"):
-        tag = tag[1:]
-    return tag
+    else:
+        tag = str(tag)
+        if tag.startswith("!"):
+            tag = tag[1:]
+        return tag
 
 
 def _uninstantiable(cls):

tests/input_files/adapters/gpt3_6.7B.yaml

@@ -0,0 +1,64 @@
+# Each tensor is shaped by a set of ranks, denoted by capital letters
+# For example: Q is shaped by (B, M, H, E)
+# We'll use lower-case letters to index into the ranks
+# For example: Q[b, m, h, e] is the tensor Q at index (b, m, h, e)
+
+# When making a projection list, it's equivalent to the Einsum subscript notation, so:
+# Q projection [b, m, h, e] means that b indexes into B, m indexes into M...
+# When making a projection dict, it's equivalent to the Einsum subscript/superscript notation, so:
+# K projection { B: b, M: p, H: h, E: e } means that b indexes into B, p indexes into M...
+
+# Renames take a tensor name and turn them into a canonical name that we can use in
+# architecture constraints. For example, we want to use the words "input", "weight", and
+# "output" to refer to the tensors of an Einsum, but the Einsum QK has no clear "weight"
+# or "input" because both Q and K are inputs. So we rename K to be weight.
+
+workload:
+  rank_sizes:
+    {% set BATCH_SIZE = BATCH_SIZE | default(1) %}
+    {% set N_TOKENS = N_TOKENS | default(8192) %}
+    B: {{BATCH_SIZE}}
+    P: {{N_TOKENS}}
+    M: {{N_TOKENS}}
+    H: 32
+    E: 128
+    F: 128
+    D: 4096 # = e * h
+    C: 16384
+    J: 4096
+    G: 4096
+
+  bits_per_value: {All: 8}
+  persistent_tensors: weight - Intermediates
+
+  einsums:
+  - !Copy
+    name: copy_I
+    tensor: I
+  - "V[b, m, h, e] = I[b, m, d] * WV[h, e, d]"
+  - "K[b, m, h, e] = I[b, m, d] * WK[h, e, d]"
+  - "Q[b, m, h, e] = I[b, m, d] * WQ[h, e, d]"
+
+  - einsum: "QK[b, m, p, h] = Q[b, m, h, e] * K[b, M: p, h, e]"
+    renames: {input: Q}
+  - "QK_softmax[b, m, p, h] = QK[b, m, p, h]"
+
+  - einsum: "AV[b, m, h, f] = QK_softmax[b, m, p, h] * V[b, M: p, h, E: f]"
+    renames: {input: QK_softmax}
+  - "Z[b, m, g] = AV[b, m, h, f] * WZ[h, f, g]"
+  - "FFA[b, m, c] = Z[b, m, g] * WFFA[g, c]"
+  - "FFB[b, m, j] = FFA[b, m, c] * WFFB[c, j]"
+
+renames:
+  einsums:
+  - name: default
+    tensor_accesses:
+    - name: input
+      source: Inputs & Intermediates if len(All) == 3 else Inputs
+      expected_count: 1
+    - name: output
+      source: Outputs
+      expected_count: 1
+    - name: weight
+      source: ~(input | output)
+      expected_count: 1 if len(All) == 3 else 0

tests/test_adapter.py

@@ -0,0 +1,15 @@
+import unittest
+from pathlib import Path
+
+import accelforge as af
+
+INPUT_FILES_DIR = Path(__file__).parent / "input_files" / "adapters"
+
+
+class TestParsing(unittest.TestCase):
+    def test_gpt3(self):
+        spec = af.Spec.from_yaml(
+            INPUT_FILES_DIR / "gpt3_6.7B.yaml"
+        )
+        self.assertEqual(spec.workload.einsum_names, ["V", "K", "Q", "QK_softmax", "Z", "FFA", "FFB"])
+