In-switch embedding bag pooling
An apparatus and method for reducing the memory bandwidth of executing machine learning models. A computing system includes two or more processing nodes, each including at least one or more processors and a corresponding local memory. Switch circuitry communicates with at least the local memories and a system memory of the computing system. The switch includes multiple direct memory access (DMA) interfaces. Each of one or more processing nodes stores multiple embedding rows of embedding tables. A processor of the processing node identifies two or more embedding rows as source operands of a reduction operation. The switch executes memory access requests to retrieve data of the two or more embedding rows from the corresponding local memory, and generates a result by performing the reduction operation. The switch sends the result to the local memory.
1 . An integrated circuit comprising:
a switch configured to be coupled to a local memory of a processing node and to a system memory, the switch comprising circuitry configured to:
receive, from the processing node, an indication that identifies a plurality of embedding rows and a reduction operation to be performed on the plurality of embedding rows;
retrieve the plurality of embedding rows from at least one of the system memory, the local memory of the processing node, a local memory of another processing node, or a remote memory;
perform the reduction operation on the plurality of embedding rows to generate a reduced result; and
send the reduced result to the local memory of the processing node.
2 . The integrated circuit as recited in claim 1 , wherein the plurality of embedding rows are entries of an embedding table of a machine learning model.
3 . The integrated circuit as recited in claim 1 , wherein in response to receiving a combined operation, the switch is further configured to:
transfer data of an entry of a first data array stored in a first processing node to a second data array stored in a second processing node; and
write, to the entry of the first data array, result data generated by performing a second type of reduction operation using data of the first data array.
4 . The integrated circuit as recited in claim 3 , wherein in response to receiving a second combined operation, the switch is further configured to:
transfer data of an entry of the second data array to the first data array; and
write, to that entry of the second data array, result data generated by performing a third type of reduction operation using data of the second data array.
5 . The integrated circuit as recited in claim 3 , wherein the switch is configured to communicate with a local memory of the first processing node using a first communication protocol and to communicate with a local memory of the second processing node using a second, different communication protocol.
6 . The integrated circuit as recited in claim 1 , wherein the reduction operation comprises one of a summing operation, maximum operation, or averaging operation on the plurality of embedding rows.
7 . The integrated circuit as recited in claim 5 , wherein a first processor is configured to process tasks and access data stored in the local memory of the first processing node and the local memory of the second processing node while the switch generates the result data.
8 . A method comprising:
receiving, by a switch of an integrated circuit and from a processing node coupled to the switch, an indication that identifies a plurality of embedding rows and a reduction operation to be performed on the plurality of embedding rows;
retrieving, by the switch, the plurality of embedding rows from at least one of a system memory, a local memory of the processing node, a local memory of another processing node, or a remote memory;
performing, by the switch, the reduction operation on the plurality of embedding rows to generate a reduced result; and
sending, by the switch, the reduced result to the local memory of the processing node.
9 . The method as recited in claim 8 , wherein the plurality of embedding rows are entries of an embedding table of a machine learning model.
10 . The method as recited in claim 8 , wherein in response to receiving a combined operation, the method further comprises:
transferring data of an entry of a first data array stored in a first processing node to a second data array stored in a second processing node; and
writing, to the entry of the first data array, result data generated by performing a second type of reduction operation using data of the first data array.
11 . The method as recited in claim 10 , wherein in response to receiving a second combined operation, the method further comprises:
transferring data of an entry of the second data array to the first data array; and
writing, to that entry of the second data array, result data generated by performing a third type of reduction operation using data of the second data array.
12 . The method as recited in claim 10 , wherein the reduction operation comprises one of a summing operation, averaging operation, or a maximum operation.
13 . The method as recited in claim 10 , further comprising communicating, by the switch, with a local memory of the first processing node using a first communication protocol and communicating with a local memory of the second processing node using a second, different communication protocol.
14 . The method as recited in claim 13 , further comprising processing tasks, by a first processor, and accessing, by the first processor, data stored in the local memory of the first processing node and the local memory of the second processing node while the switch generates the result data.
15 . A computing system comprising:
a first processing node comprising a processor and a local memory storing a plurality of embedding rows of an embedding table; and
a switch coupled to the local memory of the processing node and to a system memory accessible by the computing system, the switch comprising circuitry configured to:
receive, from the processor of the processing node, an indication that identifies a subset of the embedding rows and a reduction operation to be performed on the subset of the embedding rows;
retrieve the subset of the embedding rows from at least one of the system memory, the local memory of the processing node, a local memory of another processing node, or a remote memory;
perform the reduction operation on the subset of the embedding rows to generate a reduced result; and
send the reduced result to the local memory of the processing node.
16 . The computing system as recited in claim 15 , wherein the indication further identifies a set of offsets and indices, and wherein the switch is configured to determine the subset of the embedding rows based on the set of offsets and indices.
17 . The computing system as recited in claim 15 , wherein retrieving the subset of the embedding rows comprises issuing, by the switch, a plurality of direct memory access requests that target non-contiguous memory locations storing respective portions of the embedding rows.
18 . The computing system as recited in claim 17 , wherein performing the reduction operation comprises performing a plurality of partial reductions on subsets of elements of the embedding rows and combining the partial reductions to generate the reduced result.
19 . The computing system as recited in claim 17 , wherein the switch is further configured to schedule retrieval of the subset of embedding rows based on programmable thresholds stored in configuration registers of the switch.
20 . The computing system as recited in claim 17 , wherein the processor is configured to access the local memory using a first communication path while the switch retrieves the subset of the embedding rows using a second communication path different from the first communication path.