cerebras.modelzoo.common.input_utils.get_streaming_batch_size#
- cerebras.modelzoo.common.input_utils.get_streaming_batch_size(effective_batch_size: int) int[source]#
Returns the streaming batch size of the current task.
In a Wafer-Scaler Cluster setup with more than 1 CS-X node, the batch size used in compile and specified by user is the effective batch size at which gradient updates are done. However, each worker node streams a local batch of data to a given CS-X node to enable data parallel training.
This helper method returns the local batch size that the current task should use given the effective batch size.
- Parameters
effective_batch_size – The effective batch size of the model.
- Returns
The local batch size to be streamed by this task. If queried on the user node (used when compiling the model), this returns the original effective batch size as passed in the argument.