US 12,561,120 B1Granted

Collective operation and compute operation pipelining

AssigneeAmazon Technologies, Inc.

InventorsHongbin Zheng, Qingrui Liu, Shruthi Hebbur Prasanna Kumar, Donald John Kretsch, Yuwen Jia

Life of this patent

Estimating the term…

Loading patent…

Abstract

Techniques for collective operation and compute operation pipelining on integrated circuit devices configured to perform operations associated with machine learning models are described herein. A first node or integrated circuit device can generate a first input subtensor and a second input subtensor based on the first input tensor. Several collective operations and compute operations can be performed on the first input subtensor and second input subtensor in combination with input subtensors from other integrated circuit devices to generate output subtensors. The first node can generate an output tensor based on the output subtensors.

Claims20

Dependent claims

A computer-program product tangibly embodied in a non-transitory machine-readable storage medium, the computer-program product including instructions configured to cause a processor to perform operations for executing a compiler, the operations comprising:

receiving a sequence of operations for operating a machine learning model in computer code form to be converted into machine instructions for execution on a first integrated circuit device, the sequence of operations including an all-gather operation, one or more compute operations, and a reduce-scatter operation, each operation of the sequence of operations 8 including a rank dimension and a highest non-rank dimension;
pipelining the sequence of operations by:
modifying the computer code, for each operation of the sequence of operations, by dividing the highest non-rank dimension by a tiling factor size and adding a new highest dimension of the tiling factor size;
determining a thread axis for an output of the all-gather operation;
determining respective thread axes for each operation of the sequence of operations based on the thread axis for the output of the all-gather operation;
modifying the computer code such that each operation of the sequence of operations is a singular thread axis dimension loop associated with the respective thread axis of the operation, each singular thread axis dimension loop having a size of the tiling factor;
and
modifying the computer code to fuse the singular thread axis dimension loops for each operation of the sequence of operations into a pipeline loop having a size of the tiling factor;
and
generating the machine instructions for the first integrated circuit device based on the modified computer code.

7 more elements

A computer-implemented method comprising:

receiving, by a first node, a first input tensor;
generating, by the first node, a first input subtensor and a second input subtensor based on the first input tensor;
performing, by the first node, a first all-gather operation to store the first input subtensor and a third input subtensor from a second node;
after completion of the first all-gather operation:
performing, by the first node, a first compute operation to generate a first compute subtensor based on the first input subtensor and the third input subtensor;
performing, by the first node, a second all-gather operation to store the second input subtensor and a fourth input subtensor from the second node;
after completion of the first compute operation:
performing, by the first node, a second compute operation to generate a second compute subtensor based on the second input subtensor and the fourth input subtensor;
performing, by the first node, a first reduce-scatter operation to generate a first output subtensor based on the first compute subtensor and a third compute subtensor 17 from the second node;
after completion of the second compute operation, performing, by the first node, a second reduce-scatter operation to generate a second output subtensor based on the second compute subtensor and a fourth compute subtensor from the second node;
and
generating, by the first node, an output tensor based on the first output subtensor and the second output subtensor.

9 more elements

Dependencies are inferred from the claim text.