cerebras.modelzoo.common.input_utils.get_streaming_batch_size#
- cerebras.modelzoo.common.input_utils.get_streaming_batch_size(effective_batch_size, global_rank=None)[source]#
Returns the streaming batch size of the given 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 (int) – The effective batch size of the model.
global_rank (Optional[int]) – The global rank of the task to return the streaming batch size for. If None, it returns the streaming batch size of the current task.
- Returns
The local batch size to be streamed by the given task. If queried on the user node (used when compiling the model), this returns the original effective batch size as passed in the argument.
- Return type
int