def METHOD_NAME(config: Dict[str, Any], state: State):
    if state.dataloader is None:
        raise ValueError(
            'When using DeepSpeed, the `train_dataloader` must be specified when constructing the Trainer.')
    try:
        batch_size = state.dataloader.batch_size  # type: ignore as we catch the exception
    except AttributeError as e:
        raise RuntimeError('DeepSpeed requires the `state.dataloader` to have a `batch_size` attribute.') from e
    assert state.device_train_microbatch_size is not None
    if batch_size % state.device_train_microbatch_size != 0:
        # DeepSpeed will throw an error in this configuration.
        raise ValueError('The Mosaic trainer has been configured to use batch size='
                         f'{batch_size}, but this is not divisible by the '
                         f'train device microbatch size={state.device_train_microbatch_size}. '
                         'This is unsupported when using DeepSpeed.')
    train_batch_size = batch_size * dist.get_world_size()
    # Per the check at the start of this function, the following division is always clean.
    grad_accum = batch_size // state.device_train_microbatch_size
    if 'train_batch_size' in config:
        ds_train_batch_size = config['train_batch_size']
        if ds_train_batch_size != train_batch_size:
            raise ValueError(f'Provided DeepSpeed configuration specifies batch size={ds_train_batch_size}, '
                             f'but the Mosaic trainer has been configured with batch size={train_batch_size}.')
    else:
        config['train_batch_size'] = train_batch_size
    if 'gradient_accumulation_steps' not in config:
        config['gradient_accumulation_steps'] = grad_accum
    if 'train_micro_batch_size_per_gpu' in config:
        ds_per_gpu_microbatch_size = config['train_micro_batch_size_per_gpu']
        if ds_per_gpu_microbatch_size != state.device_train_microbatch_size:
            raise ValueError('Provided DeepSpeed configuration specifies per-GPU microbatch size='
                             f'{ds_per_gpu_microbatch_size}, but the Mosaic trainer has been '
                             f'configured with per-GPU microbatch size={state.device_train_microbatch_size}.')
    else:
        config['train_micro_batch_size_per_gpu'] = state.device_train_microbatch_size