Dynamic parallelism mode switching for operating artificial intelligence model

A method comprising:

• maintaining, by at least one processor, a cache memory that has a memory layout compatible across a plurality of parallelism modes of operations for an artificial intelligence model, the artificial intelligence model comprising a neural network model, the artificial intelligence model being deployed across a plurality of processing units;
• receiving, by the at least one processor, a set of inference requests for processing by the artificial intelligence model;
• analyzing, by the at least one processor, a set of computational workload conditions associated with the artificial intelligence model;
• selecting, by the at least one processor and for the artificial intelligence model, a current parallelism mode from the plurality of parallelism modes based on the set of computational workload conditions, the plurality of parallelism modes comprising a tensor parallelism mode and a sequence parallelism mode, the tensor parallelism mode causing a set of parameters of the artificial intelligence model to be distributed across the plurality of processing units such that each processing unit receives, stores, and uses a different subset of the set of parameters of the artificial intelligence model, the tensor parallelism mode causing each inference request of the set of inference requests to be distributed to each processing unit of the plurality of processing units for inference processing, the sequence parallelism mode causing the set of parameters to be distributed to each processing unit of the plurality of processing such that each processing unit receives, stores, and uses an entire set of parameters of the artificial intelligence model, the sequence parallelism mode causing different portions of each inference request of the set of inference requests to be distributed across the plurality of processing units for inference processing;
• using, by the at least one processor, the current parallelism mode to process the set of inference requests using the artificial intelligence model and the cache memory;
• receiving, by the at least one processor and from the artificial intelligence model, a set of inference responses for the set of inference requests;
• and
• during processing of the set of inference requests and after processing a subset of tokens through early exit logic that filters tokens based on computational significance, switching, by the at least one processor, between different parallelism modes of the plurality of parallelism modes within a single forward pass through the neural network model based on layer-specific computational requirements of the neural network model.

A machine-storage medium storing instructions that when executed by a machine, cause the machine to perform operations comprising:

• maintaining a cache memory that has a memory layout compatible across a plurality of parallelism modes of operations for an artificial intelligence model, the artificial intelligence model comprising a neural network model, the artificial intelligence model being deployed across a plurality of processing units;
• receiving a set of inference requests for processing by the artificial intelligence model;
• analyzing a set of computational workload conditions associated with the artificial intelligence model;
• selecting, for the artificial intelligence model, a current parallelism mode from the plurality of parallelism modes based on the set of computational workload conditions, the plurality of parallelism modes comprising a tensor parallelism mode and a sequence parallelism mode, the tensor parallelism mode causing a set of parameters of the artificial intelligence model to be distributed across the plurality of processing units such that each processing unit receives, stores, and uses a different subset of the set of parameters of the artificial intelligence model, the tensor parallelism mode causing each inference request of the set of inference requests to be distributed to each processing unit of the plurality of processing units for inference processing, the sequence parallelism mode causing the set of parameters to be distributed to each processing unit of the plurality of processing such that each processing unit receives, stores, and uses an entire set of parameters of the artificial intelligence model, the sequence parallelism mode causing different portions of each inference request of the set of inference requests to be distributed across the plurality of processing units for inference processing;
• using the current parallelism mode to process the set of inference requests using the artificial intelligence model and the cache memory;
• receiving, from the artificial intelligence model, a set of inference responses for the set of inference requests;
• and
• during processing of the set of inference requests and after processing a subset of tokens through early exit logic that filters tokens based on computational significance, switching between different parallelism modes of the plurality of parallelism modes within a single forward pass through the neural network model based on layer-specific computational requirements of the neural network model.