def update_config_with_batch_size_info(
    cfg: dict[str, Any],
    device_train_batch_size: Union[int, float],
    device_train_microbatch_size: Union[int, float, Literal['auto']],
    device_train_grad_accum: Union[int, Literal['auto']],
) -> dict[str, Any]:
    """Update the config with batch size information.

    Args:
        cfg (Dict[str, Any]): The config to update.
        device_train_batch_size (Union[int, float]): The batch size of the training dataset for each device.
        device_train_microbatch_size (Union[int, float, Literal['auto']]): The microbatch size of the training dataset for each device.
        device_train_grad_accum (Union[int, Literal['auto']]): The gradient accumulation settings for each device.

    Returns:
        Dict[str, Any]: The updated config.
    """
    cfg['n_gpus'] = dist.get_world_size()
    cfg['device_train_batch_size'] = device_train_batch_size
    cfg['device_train_microbatch_size'] = device_train_microbatch_size
    cfg['device_train_grad_accum'] = device_train_grad_accum
    # Safely set `device_eval_batch_size` if not provided by user
    if 'device_eval_batch_size' not in cfg:
        if cfg['device_train_microbatch_size'] == 'auto':
            cfg['device_eval_batch_size'
               ] = 1  # TODO debug auto eval microbatching
        else:
            cfg['device_eval_batch_size'] = cfg['device_train_microbatch_size']
    return cfg
Update the config with batch size information.

    Args:
        cfg (Dict[str, Any]): The config to update.
        device_train_batch_size (Union[int, float]): The batch size of the training dataset for each device.
        device_train_microbatch_size (Union[int, float, Literal['auto']]): The microbatch size of the training dataset for each device.
        device_train_grad_accum (Union[int, Literal['auto']]): The gradient accumulation settings for each device.

    Returns:
        Dict[str, Any]: The updated config.
    
llmfoundry/utils/config_utils.py