mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
19 lines
524 B
19 lines
524 B
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
from .._common_utils import divide
|
|
|
|
|
|
def vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank):
|
|
index_f = rank * per_partition_vocab_size
|
|
index_l = index_f + per_partition_vocab_size
|
|
return index_f, index_l
|
|
|
|
|
|
def vocab_range_from_global_vocab_size(global_vocab_size, rank, world_size):
|
|
per_partition_vocab_size = divide(global_vocab_size, world_size)
|
|
return vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank)
|
|
|
|
|
|
|