Tree-based network architecture for accelerating machine learning collective operations

A method for processing computations in a tree-based network architecture, the method comprising:

• computing, by each of a plurality of compute-memory stacks in a multi-chip package, a respective computation;
• aggregating, by a chip level input/output (I/O) die connected to the plurality of compute-memory stacks in the multi-chip package, the respective computations to generate a chip level aggregated computation;
• aggregating, by a server level I/O die of a server level aggregator in a server, the chip level aggregated computation with additional chip level aggregated computations to generate a server level aggregated computation;
• and
• outputting, by the server level aggregator, the server level aggregated computation.

A large model processing unit comprising a plurality of tree-based network architectures, each of the tree-based network architectures comprising:

• a server comprising a plurality of multi-chip packages connected to a server level aggregator;
• the multi-chip packages each comprising a plurality of compute-memory stacks connected to a chip level input/output (I/O) die, each chip level I/O die configured to aggregate computations performed by the compute-memory stacks to generate a chip level aggregated computation and output the chip level aggregated computation to the server level aggregator;
• and
• the server level aggregator comprising a server level I/O die configured to aggregate the chip level aggregated computation from each of the plurality of multi-chip packages to generate a server level aggregated computation and output the server level aggregated computation.