togethercomputer/StripedHyena-Nous-7B

1import torch2 3 4def column_split(x, num_heads, head_size):5    """Split a tensor with `num_heads` alongside the head dimension, instead of6    across heads. Fixed to three projections7    """8 9    x_reshaped = x.reshape(10        x.shape[0],11        num_heads,12        3 * head_size,13    )14 15    x2, x1, v = (16        x_reshaped[:, :, :head_size],17        x_reshaped[18            :,19            :,20            head_size : 2 * head_size,21        ],22        x_reshaped[:, :, 2 * head_size :],23    )24    x2, x1, v = (25        x2.reshape(x2.shape[0], -1),26        x1.reshape(x1.shape[0], -1),27        v.reshape(v.shape[0], -1),28    )29    return x2, x1, v30 31 32def get_init_from_string(init_str):33    if type(init_str) == str:34        if init_str == "torch.nn.init.zeros_":35            return torch.nn.init.zeros_36        elif init_str == "torch.nn.init.xavier_uniform_":37            return torch.nn.init.xavier_uniform_38        elif init_str == "torch.nn.init.xavier_normal_":39            return torch.nn.init.xavier_normal_40        else:41            raise ValueError(f"Unrecognized init {init_str}")42 43 44def print_rank_0(message, debug=False, end="\n"):45    """Print from rank 0 only."""46    if torch.distributed.is_initialized():47        if torch.distributed.get_rank() == 0:48            print(message, flush=True, end=end)49    else:50        print(message, flush=True, end=end)51 52 53class dotdict(dict):54    """dot.notation access to dictionary attributes"""55 56    __getattr__ = dict.get57    __setattr__ = dict.__setitem__58    __delattr__ = dict.__delitem__59 60 61def ensure_divisibility(numerator, denominator):62    """Ensure that numerator is divisible by the denominator."""63    assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator)64 65 66def divide(numerator, denominator):67    """Ensure that numerator is divisible by the denominator and return68    the division value."""69    ensure_divisibility(numerator, denominator)70    return numerator // denominator71 72 73class VocabUtility:74    """Split the vocabulary into `world_size` chunks amd return the75    first and last index of the vocabulary belonging to the `rank`76    partition: Note that indices in [first, last]"""77 78    @staticmethod79    def vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank, world_size):80        index_f = rank * per_partition_vocab_size81        index_l = index_f + per_partition_vocab_size82        return index_f, index_l83 84    @staticmethod85    def vocab_range_from_global_vocab_size(global_vocab_size, rank, world_size):86        per_partition_vocab_size = divide(global_vocab_size, world_size)87        return VocabUtility.vocab_range_from_per_partition_vocab_size(88            per_partition_vocab_size, rank, world_size89        )90