def vocab_size_with_padding(vocab_size, make_vocab_size_divisible_by, world_size):
    """Pad vocab size so it is divisible by model parallel size and
    still having GPU friendly size."""

    after = vocab_size
    multiple = make_vocab_size_divisible_by * world_size

    while (after % multiple) != 0:
        after += 1

    return after